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

__all__ = ['ProjectPath', 'RunManager']

# Cell
import pandas as pd
from datetime import date
import scipy.io as sio
import requests
import IPython
import matplotlib.pyplot as plt
from nbdev.imports import test_eq
from geoget.download import run_all
from .core import filter_files, ls, Path, InOutPath
from .geo import Region
from .data import *
from .predict import predict_nrt
Path.ls = ls

# Cell
class ProjectPath():
    def __init__(self, path:Path, config_dir='config', ladsweb_dir='ladsweb',
                 hotspots_dir='hotspots', dataset_dir='dataset', outputs_dir='outputs'):
        if isinstance(path, str): path = Path(path)
        self.path     = path
        self.config   = path/config_dir
        self.ladsweb  = path/ladsweb_dir
        self.hotspots = path/hotspots_dir
        self.dataset  = path/dataset_dir
        self.outputs  = path/outputs_dir
        self.create_dirs([self.path, self.config, self.ladsweb, self.hotspots,
                          self.dataset, self.outputs])

    @staticmethod
    def create_dirs(paths_list):
        "Create directories in list."
        for o in paths_list: o.mkdir(exist_ok=True)

# Cell
class RunManager():
    def __init__(self, project_path:ProjectPath, region, time='today',
                 product:str='VIIRS750', days=64):
        self.path    = project_path
        self.time    = self.init_time(time)
        self.product = product
        self.region  = region
        self.days    = days

    def init_time(self, time):
        if time == 'today':
            time = pd.Timestamp(date.today())
        elif time == 'yesterday':
            time = pd.Timestamp(date.today())-pd.Timedelta(days=1)
        return time

    def last_n_days(self, time:pd.Timestamp, days):
        return pd.date_range(start=time-pd.Timedelta(days=days-1), periods=days,
                              freq='D')

    def check_data(self):
        "Check existing and missing files in dataset folder."
        times = self.last_n_days(self.time, self.days)
        files, missing_files = [], []
        for t in times:
            tstr = t.strftime('%Y%m%d')
            file = self.path.dataset/f'{self.product}{self.region}_{tstr}.mat'
            if file.is_file():
                files.append(file)
            else:
                missing_files.append(file)
        return {'files': files, 'missing_files': missing_files}

    def get_download_dates(self):
        "Find for which new dates the files need to be downloaded."
        files = self.check_data()['files']
        if len(files) == 0:
            start = self.last_n_days(self.time, self.days)[0]
        else:
            start = pd.Timestamp(files[-1].stem.split('_')[-1])+pd.Timedelta(days=1)
        start = start.strftime('%Y-%m-%d 00:00:00')
        end = self.time.strftime('%Y-%m-%d 23:59:59')
        return start, end

    def update_hotspots(self, location, mode='7d', save=True):
        """Update hotspots file with new data.
          location is according to the data url naming format
          mode can be on of: 24h, 48h, 7d"""
        url = f'https://firms.modaps.eosdis.nasa.gov/' \
                   f'active_fire/viirs/text/VNP14IMGTDL_NRT_{location}_{mode}.csv'
        files = self.path.hotspots.ls(include=['.csv', f'hotspots{self.region}'])
        frp = [pd.read_csv(f) for f in files]
        frp = pd.concat([*frp, pd.read_csv(url)], axis=0, sort=False
                        ).drop_duplicates().reset_index(drop=True)
        if save:
            frp.to_csv(self.path.hotspots/f'hotspots{self.region}.csv', index=False)
            print(f'hotspots{self.region}.csv updated')
        else: return frp

    def download_viirs(self):
        "Download viirs data needed for the dataset."
        tstart, tend = self.get_download_dates()
        region = Region.load(f'{self.path.config}/R_{self.region}.json')
        viirs_downloader = VIIRS750_download(region, tstart, tend)
        viirs_downloader_list = viirs_downloader.split_times()
        run_all(viirs_downloader_list, self.path.ladsweb)

    def preprocess_dataset(self):
        "Apply pre-processing to the rawdata and saves results in dataset directory."
        paths = InOutPath(f'{self.path.ladsweb}', f'{self.path.dataset}')
        R = Region.load(f'{self.path.config}/R_{self.region}.json')
        bands = ['Reflectance_M5', 'Reflectance_M7', 'Reflectance_M10', 'Radiance_M12',
                 'Radiance_M15', 'SolarZenithAngle', 'SatelliteZenithAngle']
        print('\nPre-processing data...')
        viirs = Viirs750Dataset(paths, R, bands=bands)
        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 = ActiveFires(f'{self.path.hotspots}/hotspots{self.region}.csv')
        viirs.process_all(proc_funcs=[merge_tiles, mir_calc, rename, bfilter, act_fires])

    def init_model_weights(self, weight_files:list):
        "Downloads model weights if they don't exist yet on config directory."
        local_files = []
        for w in weight_files:
            file_save = self.path.config/w
            if not file_save.is_file():
                print(f'Downloading model weights {w}')
                url = f'https://github.com/mnpinto/banet_weights/raw/master/model/{w}'
                file = requests.get(url)
                open(str(file_save), 'wb').write(file.content)
            local_files.append(file_save)
        return local_files

    def get_preds(self, weight_files:list, threshold=0.5, save=True):
        "Computes BA-Net predictions ensembling the models in the weight_files list."
        local_files = self.init_model_weights(weight_files)
        iop = InOutPath(self.path.dataset, self.path.outputs, mkdir=False)
        predict_nrt(iop, self.time, local_files, self.region, threshold, save)