from util.geos_nav import get_navigation
from util.setup_cloud_products import model_path_cld_frac, model_path_cld_opd
from aeolus.datasource import CLAVRx
import os
from deeplearning.cloud_fraction_fcn_abi import SRCNN as SRCNN_CLD_FRAC
from deeplearning.cloud_opd_fcn_abi import SRCNN as SRCNN_CLD_OPD
from util.util import get_cartopy_crs, write_cld_prods_file_nc4
import numpy as np
import time


def infer_cloud_products(clvrx_path, output_dir, full_disk=True, satellite='GOES16', domain='FD', pattern=None):
    # -- location of the trained model
    ckpt_dir_s = os.listdir(model_path_cld_frac)
    ckpt_dir_cld_frac = model_path_cld_frac + ckpt_dir_s[0]

    ckpt_dir_s = os.listdir(model_path_cld_opd)
    ckpt_dir_cld_opd = model_path_cld_opd + ckpt_dir_s[0]

    # -- Navigation parameters
    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
    nav = get_navigation(satellite, domain)
    cc = np.arange(xlen)
    ll = np.arange(ylen)
    x_rad = cc * nav.CFAC + nav.COFF
    y_rad = ll * nav.LFAC + nav.LOFF

    # -- Create a model instance and initialize with trained model above
    nn_cld_frac = SRCNN_CLD_FRAC()
    nn_cld_frac.setup_inference(ckpt_dir_cld_frac)

    nn_cld_opd = SRCNN_CLD_OPD()
    nn_cld_opd.setup_inference(ckpt_dir_cld_opd)

    if pattern is not None:
        clvrx_ds = CLAVRx(clvrx_path, pattern=pattern)
    else:
        clvrx_ds = CLAVRx(clvrx_path)

    for pname, t_start, t_stop, dto in clvrx_ds:
        clvrx_str_time = dto.strftime('%Y-%m-%d_%H:%M')
        out_file = output_dir + 'cloud_fraction_opd_' + clvrx_str_time + '.nc'

        t0 = time.time()
        if full_disk:
            cld_frac = nn_cld_frac.run_inference_full_disk(pname, None)
            cld_opd, opd = nn_cld_opd.run_inference_full_disk(pname, None)
        else:
            cld_frac = nn_cld_frac.run_inference(pname, None)
            cld_opd, opd = nn_cld_opd.run_inference(pname, None)

        clr_cat = cld_frac == 0
        cld_cat = cld_frac == 4
        cld_opd[clr_cat] = opd[clr_cat]
        cld_opd[cld_cat] = opd[cld_cat]

        cld_opd[np.logical_and(cld_opd < 0.0, cld_opd > 160.0)] = -1.0
        cld_opd[np.isnan(opd)] = -1.0

        write_cld_prods_file_nc4(clvrx_str_time, out_file, cld_frac, cld_opd, x_rad, y_rad, None, None,
                                satellite=satellite, domain=domain, has_time=True)
        t1 = time.time()
        print('total time: ', (t1 - t0))
        print('Done: ', pname)