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, run_evaluate_static, run_evaluate_static_full_disk
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):
    ckpt_dir_s = os.listdir(model_path)
    ckpt_dir = model_path + ckpt_dir_s[0]

    clvrx_ds = CLAVRx(clvrx_path)
    for fname, t_start, t_stop in clvrx_ds:
        dto = clvrx_ds.get_datetime(fname)
        clvrx_str_time = dto.strftime('%Y-%m-%d_%H:%M')
        out_file = output_dir + 'cloud_fraction_' + clvrx_str_time
        if full_disk:
            run_evaluate_static_full_disk(fname, out_file, ckpt_dir)
        else:
            run_evaluate_static(fname, out_file, ckpt_dir)


def infer_cloud_fraction_new(clvrx_path, output_dir, full_disk=True, satellite='GOES16', domain='FD'):
    # 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)

    clvrx_ds = CLAVRx(clvrx_path)
    for fname, t_start, t_stop in clvrx_ds:
        dto = clvrx_ds.get_datetime(fname)
        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)