from util.geos_nav import get_navigation
from util.setup_cloud_fraction import model_path
from aeolus.datasource import CLAVRx
import os
from deeplearning.cloud_fraction_fcn_abi import SRCNN
from util.util import get_cartopy_crs, write_cld_frac_file_nc4
import numpy as np


def infer_cloud_fraction(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)
    ckpt_dir = model_path + 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 = SRCNN()
    nn.setup_inference(ckpt_dir)

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

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

        if full_disk:
            cld_frac = nn.run_inference_full_disk(fname, None)
        else:
            cld_frac = nn.run_inference(fname, None)

        write_cld_frac_file_nc4(clvrx_str_time, out_file, cld_frac, x_rad, y_rad, None, None, satellite=satellite,
                                domain=domain, has_time=True)
        print('Done: ', clvrx_path + fname)