# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/06_cli.ipynb (unless otherwise specified).

__all__ = ['banet_viirs375_download', 'banet_viirs750_download', 'banet_create_dataset', 'banet_dataset2tiles',
           'banet_predict_monthly', 'banet_predict_times', 'banet_nrt_run']

# Cell
import calendar
import pandas as pd
from fastscript import call_parse, Param
import pdb
import os
import IPython
import pdb

from geoget.download import run_all

from .core import InOutPath, Path, ls, dict2json
from .data import *
from .geo import Region
from .predict import predict_month, predict_time
from .models import BA_Net
#from banet.train import train_model
from .nrt import ProjectPath, RunManager

# Cell
Path.ls = ls
_bands =  ['Reflectance_M5', 'Reflectance_M7', 'Reflectance_M10', 'Radiance_M12',
           'Radiance_M15', 'SolarZenithAngle', 'SatelliteZenithAngle']
_regions_path = '../data/regions'
_weight_files = ['models/banetv0.20-val2017-fold0.pth',
                 'models/banetv0.20-val2017-fold1.pth',
                 'models/banetv0.20-val2017-fold2.pth']
_weight_files = [os.path.expanduser('~/.banet/') + o for o in _weight_files]

# Cell
@call_parse
def banet_viirs375_download(region:Param("Region name", str),
    tstart:Param("Start of serach window yyyy-mm-dd HH:MM:SS", str),
    tend:Param("End of search windo yyyy-mm-dd HH:MM:SS", str),
    path_save:Param("Path to save the outputs of the request", str),
    regions_path:Param("Path for region json files", str)):
    Path(path_save).mkdir(exist_ok=True)
    region = Region.load(f'{regions_path}/R_{region}.json')
    viirs_downloader1 = VIIRS375_download(region, tstart, tend)
    viirs_downloader2 = VIIRS750_download(region, tstart, tend, bands=['SolarZenithAngle',
                      'SatelliteZenithAngle'])
    viirs_downloader_list1 = viirs_downloader1.split_times()
    viirs_downloader_list2 = viirs_downloader2.split_times()
    viirs_downloader_list = [*viirs_downloader_list1, *viirs_downloader_list2]
    print(f'Splitting request into {len(viirs_downloader_list)} orders.')
    run_all(viirs_downloader_list, path_save)

# Cell
@call_parse
def banet_viirs750_download(region:Param("Region name", str),
    tstart:Param("Start of serach window yyyy-mm-dd HH:MM:SS", str),
    tend:Param("End of search windo yyyy-mm-dd HH:MM:SS", str),
    path_save:Param("Path to save the outputs of the request", str),
    regions_path:Param("Path for region json files", str)):
    Path(path_save).mkdir(exist_ok=True)
    region = Region.load(f'{regions_path}/R_{region}.json')
    viirs_downloader = VIIRS750_download(region, tstart, tend)
    viirs_downloader_list = viirs_downloader.split_times()
    print(f'Splitting request into {len(viirs_downloader_list)} orders.')
    run_all(viirs_downloader_list, path_save)

# Cell
@call_parse
def banet_create_dataset(region:Param("Region name", str),
                   viirs_path:Param("Input path for VIIRS raw data", str),
                   fires_path:Param("Input path for Active Fires csv", str),
                   save_path:Param("Path to save outputs", str),
                   regions_path:Param("Path where region defenition files are stored", str),
                   mcd64_path:Param("Input path for MCD64 raw data", str)=None,
                   cci51_path:Param("Input path for FireCCI51 raw data", str)=None,
                   bands:Param("List of bands to use as inputs for VIIRS raw data", str)=_bands,
                   year:Param("Set to process a single year instead of all available", int)=None):

    paths = InOutPath(f'{viirs_path}', f'{save_path}')
    R = Region.load(f'{regions_path}/R_{region}.json')

    # VIIRS750
    print('\nCreating dataset for VIIRS750')
    viirs = Viirs750Dataset(paths, R, bands=bands)
    viirs.filter_times(year)
    merge_tiles = MergeTiles('SatelliteZenithAngle')
    mir_calc = MirCalc('SolarZenithAngle', 'Radiance_M12', 'Radiance_M15')
    rename = BandsRename(['Reflectance_M5', 'Reflectance_M7'], ['Red', 'NIR'])
    bfilter = BandsFilter(['Red', 'NIR', 'MIR'])
    act_fires = ActiveFiresLog(f'{fires_path}/hotspots{R.name}.csv')
    viirs.process_all(proc_funcs=[BandsAssertShape(), merge_tiles, mir_calc,
                                  rename, bfilter, act_fires])

    # MCD64A1C6
    if mcd64_path is not None:
        print('\nCreating dataset for MCD64A1C6')
        paths.input_path = Path(mcd64_path)
        mcd = MCD64Dataset(paths, R)
        mcd.match_times(viirs)
        mcd.process_all()

    # FireCCI51
    if cci51_path is not None:
        print('\nCreating dataset for FireCCI51')
        paths.input_path = Path(cci51_path)
        cci51 = FireCCI51Dataset(paths, R)
        cci51.match_times(viirs)
        cci51.process_all()

# Cell
@call_parse
def banet_dataset2tiles(region:Param("Region name", str),
                  input_path:Param("Input path for dataset", str),
                  output_path:Param("Output path for tiles dataset", str),
                  size:Param("Tiles size", int)=128,
                  step:Param("Step size of moving window to create tiles", int)=100,
                  year:Param("Set to process a single year instead of all available", int)=None):

    iop = InOutPath(input_path, output_path)
    r2t = Region2Tiles(iop, 'VIIRS750', 'MCD64A1C6', regions=[region],
                       bands=[['Red', 'NIR', 'MIR', 'FRP'], ['bafrac']],
                       size=size, step=step)
    if year is None: r2t.process_all()
    else: r2t.process_all(include=[f'_{year}'])

# Cell
@call_parse
def banet_predict_monthly(region:Param("Region name", str),
                    input_path:Param("Input path for dataset", str),
                    output_path:Param("Output path for tiles dataset", str),
                    year:Param("Set to process a single year instead of all available", int),
                    weight_files:Param("List of pth weight files", list)=_weight_files):

    iop = InOutPath(input_path, f'{output_path}')
    times = pd.DatetimeIndex([pd.Timestamp(o.stem.split('_')[-1])
                              for o in iop.src.ls(include=['.mat'])])
    times = times[times.year == year]
    tstart, tend = times.min(), times.max()
    month_start = (tstart + pd.Timedelta(days=31)).month
    for m in range(month_start, tend.month):
        print(f'Generating maps for {calendar.month_name[m]} {year}:')
        t = pd.Timestamp(f'{year}-{m}-01')
        predict_month(iop, t, weight_files, region)

# Cell
@call_parse
def banet_predict_times(region:Param("Region name", str),
                    tstart:Param("Start of search window yyyy-mm-01", str),
                    tend:Param("End of search window yyyy-mm-01", str),
                    input_path:Param("Input path for dataset", str),
                    output_path:Param("Output path for tiles dataset", str),
                    regions_path:Param("Path for region json files", str),
                    product:Param("Name of product (default VIIRS750)", str)="VIIRS750",
                    output:Param("Name of file to save results", str)="data",
                    weight_files:Param("List of pth weight files", list)=_weight_files):

    iop = InOutPath(input_path, f'{output_path}')

    tstart = pd.Timestamp(tstart) - pd.Timedelta(days=15)
    tstart = pd.Timestamp(f'{tstart.year}-{tstart.month}-01')
    tend = pd.Timestamp(tend) + pd.Timedelta(days=75)
    tend = pd.Timestamp(f'{tend.year}-{tend.month}-01') - pd.Timedelta(days=1)
    times = pd.date_range(tstart, tend, freq='D')

    R = Region.load(f'{regions_path}/R_{region}.json')
    predict_time(iop, times, weight_files, R, product=product, output=output)

# Cell
@call_parse
def banet_nrt_run(region:Param("Region name", str),
                  left:Param("Left limit of the bounding box.", float),
                  bottom:Param("Bottom limit of the bounding box.", float),
                  right:Param("Right limit of the bounding box.", float),
                  top:Param("Top limit of the bounding box.", float),
                  project_path:Param("Root directory of the project", str),
                  hotspots_region:Param("Hotspots region name", str),
                  product:Param("Product name", str)='VIIRS750',
                  area_epsg:Param("EPSG to compute burned area", int)=3763,
                  pixel_size:Param("Pixel size", float)=0.01,
                  time:Param("Day for the run", str, choices=["today", "yesterday"])="today",
                  threshold:Param("Threshold to apply to output of the model", float)=0.5,
                  skip_hotspots:Param("Skip download of ladsweb data", bool)=False,
                  skip_ladsweb:Param("Skip download of ladsweb data", bool)=False,
                  skip_preprocess:Param("Skip download of ladsweb data", bool)=False,
                  skip_getpreds:Param("Skip download of ladsweb data", bool)=False,
                  skip_postprocess:Param("Skip post processing", bool)=False,
                  single_fold:Param("Use only one model", bool)=False,
                  max_size:Param("Max size to split region on inference and avoid memory error", int)=2000):
    paths = ProjectPath(project_path)
    weight_files = ['banetv0.20-val2017-fold0.pth',
                    'banetv0.20-val2017-fold1.pth',
                    'banetv0.20-val2017-fold2.pth']
    if single_fold: weight_files = ['banetv0.20-val2017-fold0.pth']
    manager = RunManager(paths, region, time=time, product=product, max_size=max_size)
    R = {'name': region, 'bbox': [left, bottom, right, top], 'pixel_size': pixel_size}
    dict2json(R, paths.config/f'R_{region}.json')
    if not skip_hotspots: manager.update_hotspots(hotspots_region)
    if not skip_ladsweb: manager.download_viirs()
    if not skip_preprocess: manager.preprocess_dataset()
    if not skip_getpreds: manager.get_preds(weight_files, threshold=threshold, max_size=max_size)
    if not skip_postprocess: manager.postprocess(filename=f'ba_{manager.time.strftime("%Y%m%d")}',
                    threshold=threshold, area_epsg=area_epsg)