from __future__ import print_function
import cmor
import numpy

nlat = 90
nlon = 180


def mywrite(data=None, time_vals=None, append_to=None, cell_bounds=None):
    breq = "100000. 80000. 80000. 68000. 68000. 56000. 56000. 44000. 44000. 31000. 31000. 18000. 18000.  0.".split()

    bnds_req = []
    for b in breq:
        bnds_req.append(float(b))

    bnds_req = numpy.array(bnds_req)
    bnds_req.shape = (7, 2)

    print(bnds_req[-2], bnds_req.shape)

    levs = []

    for b in bnds_req:
        levs.append((b[0] + b[1]) / 2.)

    levs = numpy.array(levs)

    print (levs)

    ipth = "Test"
    if append_to is None:
        mode = cmor.CMOR_REPLACE
    else:
        mode = cmor.CMOR_APPEND
    print ('Mode in python:', mode)
    cmor.setup(inpath=ipth,
               set_verbosity=cmor.CMOR_NORMAL,
               netcdf_file_action=mode,
               logfile=None)

    cmor.dataset_json("Test/common_user_input.json")

    cmor.load_table("Tables/CMIP6_Amon.json")

    dlat = 180 / nlat
    dlon = 360. / nlon
    lats = numpy.arange(-90 + dlat / 2., 90, dlat)
    bnds_lat = numpy.arange(-90, 90 + dlat, dlat)
    lons = numpy.arange(0 + dlon / 2., 360., dlon) - 180.
    bnds_lon = numpy.arange(0, 360. + dlon, dlon) - 180.

    plevs = numpy.array([100000., 92500., 85000., 70000.,
                         60000., 50000., 40000., 30000., 25000., 20000.,
                         15000., 10000., 7000., 5000., 3000., 2000., 1000., 500., 100.])

    itim = cmor.axis(table_entry='time', units='month since 2008')
    #itim = cmor.axis(table_entry='time',units='month since 2008',coord_vals=numpy.arange(0,12,1))
    ilat = cmor.axis(
        table_entry='latitude',
        coord_vals=lats,
        units='degrees_north',
        cell_bounds=bnds_lat)
    ilon = cmor.axis(
        table_entry='longitude',
        coord_vals=lons,
        units='degrees_east',
        cell_bounds=bnds_lon)
    print ('so far', itim, ilat, ilon)
    ilev = cmor.axis(table_entry="plev19", coord_vals=plevs, units="Pa")

    iv = cmor.variable(
        table_entry='ta', axis_ids=numpy.array(
            (itim, ilev, ilat, ilon)), units='K')

    # cmor.write(iv,data)
    if append_to is None:
        print ('time:', time_vals)
        print ('bnds:', cell_bounds)
        # ,file_suffix="with-appending")
        cmor.write(iv, data, time_vals=time_vals, time_bnds=cell_bounds)
    else:
        print ('Ok writing with a suffix', append_to)
        cmor.write(
            iv,
            data,
            time_vals=time_vals,
            file_suffix=append_to,
            time_bnds=cell_bounds)
        print ('and back')
    file = cmor.close(iv, file_name=True)
    print ('Ok dumped to:', file)
    cmor.close()
    return file


ntime = 12
data = numpy.random.random((ntime, 19, nlat, nlon)) + 280.


f1 = mywrite(data=data[:6], time_vals=numpy.arange(
    0, 6, 1), cell_bounds=numpy.arange(0, 7, 1))
print ('First part: ', f1)
f2 = mywrite(data=data[6:], time_vals=numpy.arange(
    6, 12, 1), cell_bounds=numpy.arange(6, 13, 1), append_to=f1)
print (f2)