#****** Conflict with 'orog' in formula terms *****
from __future__ import print_function
import cmor
import numpy

error_flag = cmor.setup(inpath='Test', netcdf_file_action=cmor.CMOR_REPLACE)

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

# creates 1 degree grid
nlat = 180
nlon = 360
alats = numpy.arange(180) - 89.5
bnds_lat = numpy.arange(181) - 90
alons = numpy.arange(360) + .5
bnds_lon = numpy.arange(361)
cmor.load_table("TestTables/CMIP6_Emon.json")
ilat = cmor.axis(
    table_entry='latitude',
    units='degrees_north',
    length=nlat,
    coord_vals=alats,
    cell_bounds=bnds_lat)

ilon = cmor.axis(
    table_entry='longitude',
    length=nlon,
    units='degrees_east',
    coord_vals=alons,
    cell_bounds=bnds_lon)

lev = 1
ntimes = 12
plevs = (numpy.arange(lev) + 1) * 1.E4


itim = cmor.axis(
    table_entry='time',
    units='months since 2030-1-1',
    length=ntimes,
    interval='1 month')

zlevs = numpy.array((0.1, ))
zlev_bnds = numpy.array((0., .2, ))
table_entry = 'hybrid_height'
if table_entry == 'hybrid_height':
    ilev = cmor.axis(
        table_entry='hybrid_height',
        # table_entry='standard_sigma',
        # table_entry='standard_hybrid_sigma',
        units='m',
        length=lev,
        coord_vals=zlevs,
        cell_bounds=zlev_bnds)

    p0 = 0.5e4
##   p0 = 1.e5
# a_coeff = (/ 0.1, 0.2, 0.3, 0.2, 0.1 /)
    a_coeff = numpy.array((0.2, ))
    b_coeff = numpy.array((0.0, ))

# a_coeff_bnds=(/0.,.15, .25, .25, .15, 0./)
    a_coeff_bnds = numpy.array((0., .3,))
    b_coeff_bnds = numpy.array((0., .05,))

# error_flag = cmor.zfactor(
# zaxis_id=ilev,
# zfactor_name='ptop',
# units='Pa',
# zfactor_values = p0)

    error_flag = cmor.zfactor(
        zaxis_id=ilev,
        zfactor_name='b',
        axis_ids=numpy.array((ilev, )),
        zfactor_values=b_coeff,
        zfactor_bounds=b_coeff_bnds)

# error_flag = cmor.zfactor(
# zaxis_id=ilev,
# zfactor_name='lev',
##         axis_ids= numpy.array(( ilev, )),
# units='m',
##         zfactor_values = a_coeff,
# zfactor_bounds = a_coeff_bnds )

    data2d = numpy.random.random((180, 360)).astype('f') * 8000

    zfactor_id = cmor.zfactor(
        zaxis_id=ilev,
        zfactor_name='orog',
        axis_ids=numpy.array((ilon, ilat)),
        units='m',
        zfactor_values=data2d)
else:
    print ('yep working case')
    ilev = cmor.axis(
        table_entry='standard_sigma',
        units='1',
        length=lev,
        coord_vals=zlevs,
        cell_bounds=zlev_bnds)

    p0 = 0.5E4
    a_coeff = numpy.array((0.2, 0.4, 0.6, 0.8, 0.95))
    b_coeff = numpy.array((0.0, 0.1, 0.2, 0.5, 0.8))

    a_coeff_bnds = numpy.array((0., .3, .5, .7, .9, 1.))
    b_coeff_bnds = numpy.array((0., .05, .15, .35, .65, 1.))

    error_flag = cmor.zfactor(
        zaxis_id=ilev,
        zfactor_name='ptop',
        units='Pa',
        zfactor_values=p0)

    error_flag = cmor.zfactor(
        zaxis_id=ilev,
        zfactor_name='sigma',
        axis_ids=numpy.array((ilev, )),
        zfactor_values=a_coeff,
        zfactor_bounds=a_coeff_bnds)

    data2d = numpy.random.random((180, 360)).astype('f') - 97000.

    zfactor_id = cmor.zfactor(
        zaxis_id=ilev,
        zfactor_name='ps',
        axis_ids=numpy.array((ilon, ilat, itim)),
        units='Pa')

print ("ILEV is:", ilev)
var3d_ids = cmor.variable(
    table_entry='concdust',
    units='kg m-3',
    axis_ids=numpy.array((ilev, ilon, ilat, itim)),
    missing_value=1.0e28,
    original_name='cocoa is good, but concoa is better')


for it in range(ntimes):

    time = numpy.array((it))
    bnds_time = numpy.array((it, it + 1))
    data3d = numpy.random.random((lev, 360, 180)).astype('f') * 40.

    error_flag = cmor.write(
        var_id=var3d_ids,
        data=data3d,
        ntimes_passed=1,
        time_vals=time,
        time_bnds=bnds_time)


error_flag = cmor.close()