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

__all__ = ['VIIRS375_download', 'VIIRS750_download', 'BandsFilter', 'BandsRename', 'MergeTiles', 'BandsAssertShape',
           'ActiveFiresLog', 'InterpolateAngles', 'MirCalc', 'BaseDataset', 'mask2nan', 'nan2zero', 'group_files',
           'ViirsDataset', 'Viirs375Dataset', 'Viirs750Dataset', 'MCD64Dataset', 'FireCCI51Dataset', 'AusCoverDataset',
           'MTBSDataset', 'ICNFDataset', 'Region2Tiles']

# Cell
import numpy as np
import pandas as pd
import re
import gc
import sys
from pathlib import Path
from rasterio.coords import disjoint_bounds
from tqdm import tqdm
import scipy.io as sio
from functools import partial
import netCDF4
from netCDF4 import Dataset
from pyhdf.SD import SD, SDC
from geopandas import GeoDataFrame
from shapely.geometry import Point
import matplotlib.pyplot as plt
from concurrent.futures import ThreadPoolExecutor
from warnings import warn
from nbdev.imports import test_eq
from pyresample.geometry import SwathDefinition, AreaDefinition, CRS, create_area_def
from pyresample import kd_tree

from geoget.download import *
from .core import *
from .geo import *
from IPython.core.debugger import set_trace

# Cell
class VIIRS375_download(Ladsweb):
    "Utility for downloading VIIRS 375m data to create the dataset."
    def __init__(self, region, tstart, tend, bands=None):
        product = 'NPP_VIAES_L1'
        collection = '5000'
        if bands is None:
            bands = ['Reflectance_I1', 'Reflectance_I2', 'Reflectance_I3',
                      'Radiance_I4', 'Radiance_I5']
        super().__init__(product, collection, tstart, tend, list(region.bbox),
                         bands, daynight='D', repPixSize=region.pixel_size)

# Cell
class VIIRS750_download(Ladsweb):
    "Utility for downloading VIIRS 750m data to create the dataset."
    def __init__(self, region, tstart, tend, bands=None):
        product = 'NPP_VMAES_L1'
        collection = '5000'
        if bands is None:
            bands = ['Reflectance_M5', 'Reflectance_M7', 'Reflectance_M10',
                      'Radiance_M12', 'Radiance_M15', 'SolarZenithAngle',
                      'SatelliteZenithAngle']
        super().__init__(product, collection, tstart, tend, list(region.bbox),
                         bands, daynight='D', repPixSize=region.pixel_size)

# Cell
class BandsFilter():
    """Remove bands not in to_keep list from the dictionary."""
    def __init__(self, to_keep: list):
        self.to_keep = to_keep if isinstance(to_keep, list) else [to_keep]

    def __call__(self, data:dict, *args, **kwargs) -> dict:
        keys = [k for k in data]
        for k in keys:
            if k not in self.to_keep:
                del data[k]
        return data

# Cell
class BandsRename():
    def __init__(self, input_names:list, output_names:list):
        self.input_names = input_names if isinstance(input_names, list) else [input_names]
        self.output_names = output_names if isinstance(output_names, list) else [output_names]

    def __call__(self, data:dict, *args, **kwargs) -> dict:
        for i, o in zip(self.input_names, self.output_names):
            data[o] = data.pop(i)
        return data

# Cell
class MergeTiles():
    def __init__(self, band:str, ignore:list=[]):
        self.band = band
        self.ignore = ignore

    def __call__(self, data:dict, *args, **kwargs) -> dict:
        d = np.nanmean(np.array(data[self.band]), axis=(1,2))
        d = np.array(np.array(d).argsort())
        masks = np.array(data[self.band])[d]
        for k in data:
            if k not in self.ignore:
                data_aux = np.zeros_like(data[k][0])*np.nan
                for dband, mask in zip(np.array(data[k])[d], masks):
                    I = (np.isnan(data_aux)) & (~np.isnan(mask))
                    data_aux[I] = dband[I]
                data[k] = data_aux
        return data

# Cell
class BandsAssertShape():
    def __call__(self, data:dict, max_size=None, *args, **kwargs) -> dict:
        for k in kwargs['cls'].bands:
            rshape = kwargs['cls'].region.shape
            if isinstance(data[k], list):
                for d in data[k]:
                    shape = d.shape
                    if len(shape) == 3: # first is time
                        shape = shape[1:]
                    if shape != rshape:
                        error = f'{k} shape {shape} does not match region shape {rshape}'
                        raise Exception(error)
            else:
                shape = data[k].shape
                if len(shape) == 3: # first is time
                    shape = shape[1:]
                if shape != rshape:
                    error = f'{k} shape {shape} does not match region shape {rshape}'
                    raise Exception(error)
        return data

# Cell
class ActiveFiresLog():
    """Get active fires, interpolate to grid and apply log1p."""
    def __init__(self, file):
        self.file = file
        self.lon = None
        self.lat = None
        self.df = self.load_csv()

    def load_csv(self):
        return pd.read_csv(self.file, parse_dates=['acq_date']).set_index('acq_date')

    def __call__(self, data, time, *args, **kwargs):
        R = kwargs['cls'].region.new()
        if "R" in data and isinstance(data["R"], Region):
            R = data["R"]
        if R.pixel_size < 0.01:
            Rbase = R.new(pixel_size=0.01)
            self.lon, self.lat = Rbase.coords()
        else:
            self.lon, self.lat = R.coords()
        frp = self.df[self.df.index == time]

        if len(frp) > 0:
            geometry = [Point(xy) for xy in zip(frp['longitude'], frp['latitude'])]
            frp = GeoDataFrame(frp, geometry=geometry)
            out = rasterize(frp, 'frp', Rbase, merge_alg='add')
            if R.pixel_size < 0.01:
                out = downsample(out, src_tfm=Rbase.transform, dst_tfm=R.transform,
                                 dst_shape=R.shape, resampling='bilinear')
            out[out==0] = np.nan
        else: out = np.zeros(R.shape)*np.nan
        data['FRP'] = np.log1p(out)
        return data

# Cell
class InterpolateAngles():
    """Interpolated Angles to working resolution."""
    def __init__(self, R_base:Region, R_interp:Region, bands:list):
        self.R_base   = R_base
        self.R_interp = R_interp
        self.bands    = bands

    def __call__(self, data:dict, *args, **kwargs):
        R_interp = self.R_interp
        if "R" in data and isinstance(data["R"], Region):
            R_interp = data["R"]
        for b in self.bands:
            data[b] = [downsample(data[b][i], src_tfm=self.R_base.transform,
                                 dst_tfm=R_interp.transform,
                                 dst_shape=R_interp.shape,
                                 resampling='bilinear')
                       for i in range(len(data[b]))]
        return data

# Cell
class MirCalc():
    def __init__(self, solar_zenith_angle:str, mir_radiance:str, tir_radiance:str,
                 output_name:str='MIR'):
        self.sza = solar_zenith_angle
        self.r_mir = mir_radiance
        self.r_tir = tir_radiance
        self.output_name = output_name

    def __call__(self, data:dict, *args, **kwargs):
        sza = data[self.sza]
        mir = data[self.r_mir]
        tir = data[self.r_tir]
        data[self.output_name] = self.refl_mir_calc(mir, tir, sza, sensor=kwargs['cls'].name)
        return data

    def refl_mir_calc(self, mir, tir, sza, sensor):
        """
        Computes the MIR reflectance from MIR radiance and Longwave IR radiance.
        sensor can be "VIIRS375" or "VIIRS750"
        sza is the solar zenith angle
        for VIIRS375, mir is band I4 and tir band I5
        for VIIRS750, mir is band M12 and tir band M15
        returns a matrix of MIR reflectances with the same shape as mir and tir inputs.
        Missing values are represented by 0.
        """
        lambda_M12= 3.6966
        lambda_M15=10.7343
        lambda_I4 = 3.7486
        lambda_I5 = 11.4979

        c1 = 1.1911e8 # [ W m-2 sr-1 (micrometer -1)-4 ]
        c2 = 1.439e4 # [ K micrometer ]
        E_0_mir_M12 = 11.7881 # M12 newkur_semcab
        E_0_mir_I4= 11.2640 # I4 newkur_semcab

        if sensor=='VIIRS375':
            lambda_mir = lambda_I4
            lambda_tir = lambda_I5
            E_0_mir = E_0_mir_I4
        elif sensor=='VIIRS750':
            lambda_mir = lambda_M12
            lambda_tir = lambda_M15
            E_0_mir = E_0_mir_M12
        else: raise NotImplementedError(
            f'refl_mir_calc not implemented for {sensor}. Available options are VIIRS750 and VIIRS375.')

        miu_0=np.cos((sza*np.pi)/180)

        mir[mir <= 0] = np.nan
        tir[tir <= 0] = np.nan

        # Brighness temperature
        a1 = (lambda_tir**5)
        a = c1/(a1*tir)
        logaritmo = np.log(a+1)
        divisor = lambda_tir*logaritmo
        T = (c2/divisor)
        del a, logaritmo, divisor

        # Plank function
        divisor2 = (lambda_mir*T)
        exponencial = np.exp(c2/divisor2)
        b = c1*(lambda_mir**-5)
        BT_mir = b/(exponencial-1)
        del divisor2, exponencial, b, T

        # MIR reflectance
        c = (E_0_mir*miu_0)/np.pi
        termo1 = (mir-BT_mir)
        termo2 = (c-BT_mir)
        Refl_mir = termo1/termo2
        Refl_mir[Refl_mir <= 0] = 0
        return Refl_mir

# Cell
class BaseDataset():
    def __init__(self, name:str, paths:InOutPath, region:Region,
                 times:pd.DatetimeIndex=None, bands:list=None):
        self.paths = paths
        self.region = region
        self.name = name
        self.times = times
        self.bands = bands

        if self.times is None:
            self.times = self.find_dates()

    def list_files(self, time:pd.Timestamp) -> list:
        "This method should return a list of filenames corresponding to the given Timestamp."
        pass

    def find_dates(self):
        """This method should return a pd.DatetimeIndex
        with list of dates present in the data available in the input path."""
        pass

    def match_times(self, other, on='month'):
        "Set the times attribute to match the times of other dataset."
        if on != 'month':
            raise NotImplementedError('match_times is only implemented on month.')

        ym_other = sorted(set([(t.year, t.month) for t in other.times]))
        out = []
        for t in self.times:
            if (t.year, t.month) in ym_other:
                out.append(t)
        self.times = pd.DatetimeIndex(out)

    def filter_times(self, year):
        """To select only a specific year. This can be usefull for testing and
           for adding more new years and avoid reprocessing all the dataset."""
        if year is not None:
            self.times = self.times[self.times.year == year]

    def open(self, files:list, crop=None) -> dict:
        """This method is used to open a file or list of files for a given
        time period and returns a dictionary with the data ready to be passed
        to the processing functions."""
        pass

    def save(self, time:pd.Timestamp, data:dict, crop=None):
        "Saves data in a single file for a specific timestamp in netcdf4 format."
        scale_factor = 100
        mask_value = 65535
        tstr = time.strftime('%Y%m%d')
        filename = f'{self.paths.dst}/{self.name}{self.region.name}_{tstr}.nc'
        mode = "r+" if Path(filename).is_file() and crop is not None else "w"
        with netCDF4.Dataset(filename, mode, format="NETCDF4") as rootgrp:
            if mode == "w":
                latD = rootgrp.createDimension("lat", self.region.height)
                lonD = rootgrp.createDimension("lon", self.region.width)
            for k in data:
                if mode == "w":
                    a = rootgrp.createVariable(k,"u2",("lat","lon",), zlib=True)
                    a.setncattr('scale', scale_factor)
                    a.setncattr('mask', mask_value)
                else: a = rootgrp[k]
                v = data[k]*scale_factor
                v[np.isnan(v)] = mask_value
                assert np.sum(data[k] < 0) == 0
                if np.sum(data[k]*scale_factor > mask_value) > 0:
                    warn(f'Clipping data to {mask_value-1} for {k}.')
                    v[v>mask_value] = mask_value-1
                if crop is None:
                    a[:] = v.astype(np.uint16)
                else:
                    a[crop[0]:crop[1],crop[2]:crop[3]] = v.astype(np.uint16)

    def process_one(self, time:pd.Timestamp, proc_funcs:list=[], save=True, max_size=None,
                    **proc_funcs_kwargs):
        """This method defines a processing pipeline consisting of opening the file
        using the `open` method, applying each of the `proc_funcs` to the output of the previous
        and `save` the processed data using save method."""
        tstr = time.strftime('%Y%m%d')
        files = self.list_files(time)
        #try:
        tstr = time.strftime('%Y%m%d')
        filename = f'{self.paths.dst}/{self.name}{self.region.name}_{tstr}.nc'
        if not Path(filename).is_file():
            if len(files) > 0:
                if max_size is None:
                    si = [[0, self.region.shape[0], 0, self.region.shape[1]]]
                else:
                    si = [[max(0,j*max_size),(j+1)*max_size,max(0,i*max_size),(i+1)*max_size]
                           for i in range((self.region.shape[1]-1)//max_size+1)
                           for j in range((self.region.shape[0]-1)//max_size+1)]
                for c in si:
                    #print(f'Processing crop {c}')
                    data = self.open(files, crop=c)
                    kwargs = {'cls': self, **proc_funcs_kwargs}
                    for f in proc_funcs:
                        data = f(data, time, **kwargs)
                    if save:
                        self.save(time, data, crop=c)
                        del data
                        gc.collect()
                    else: return data
            else:
                warn(f'No files for {time}. Skipping to the next time.')
        else: warn(f'Skipping files for {time}. File already exists.')
        #except:
        #    msg = f'Unable to process files for {time}. Check if files are corrupted. Skipping to the next time. { sys.exc_info()[0]}'
        #    warn(msg, UserWarning)

    def process_all(self, proc_funcs=[], max_workers=1, max_size=None, **proc_funcs_kwargs):
        """`process_all` runs `process_one` in parallel using the number of workers defined
        by `max_workers` and passes the `proc_funcs` list to `process_one` method"""
        process_one = partial(self.process_one, proc_funcs=proc_funcs, max_size=max_size,
                              **proc_funcs_kwargs)
        with ThreadPoolExecutor(max_workers) as e:
            list(tqdm(e.map(process_one, self.times), total=len(self.times)))

    def __repr__(self):
        return '\n'.join([f'{i}: {o}' for i, o in self.__dict__.items()]) + '\n'

# Cell
def mask2nan(data:dict, *args, **kwargs) -> dict:
    for k in data:
        k_list = data[k]
        for i, d in enumerate(k_list):
            mask = d.mask
            data[k][i] = d.data
            data[k][i][mask] = np.nan
    return data

def nan2zero(data:dict, *args, **kwargs) -> dict:
    for k in data:
        data[k][np.isnan(data[k])] = 0
    return data

def group_files(files:list):
    return pd.DataFrame({'files': files, 'ids':['.'.join(f.stem.split('.')[-4:-2]) for f in files]}
                        ).groupby('ids').agg(lambda x : list(x)).files.values.tolist()

class ViirsDataset(BaseDataset):
    "Subclass of `BaseDataset` to process VIIRS 750-meter bands."
    _name = None
    def __init__(self, paths:InOutPath, region:Region,
                 times:pd.DatetimeIndex=None, bands:list=None):
        super().__init__(self._name, paths, region, times, bands)
        self.times = self.check_files()

    def list_files(self, time:pd.Timestamp)-> list:
        if time in self.times:
            dayOfYear = str(time.dayofyear).zfill(3)
            files = self.paths.src.ls(include=[f'.A{time.year}{dayOfYear}.', '.nc'])
        return files

    def check_files(self):
        not_missing = []
        for i, t in tqdm(enumerate(self.times), total=len(self.times)):
            files = self.list_files(t)
            if len(files)==0:
                print(f'Missing files for {t}')
            else: not_missing.append(i)
        return self.times[not_missing]

    def find_dates(self, first:pd.Timestamp=None, last:pd.Timestamp=None):
        pattern = r'^\w+.A(20[0-9][0-9])([0-3][0-9][0-9])..*$'
        times = []
        for f in self.paths.src.ls():
            x = re.search(pattern, f.stem)
            if x is not None:
                year, doy = map(x.group, [1,2])
                times.append(pd.Timestamp(f'{year}-01-01') + pd.Timedelta(days=int(doy)-1))
        self.times = pd.DatetimeIndex(sorted(set(times)))
        if first is not None:
            self.times = self.times[self.times>=first]
        if last is not None:
            self.times = self.times[self.times<=last]
        return self.times

    def open_netcdf4(self, files:list, dtype=np.float32) -> dict:
        data_dict = {}
        for f0 in files:
            nc_data = netCDF4.Dataset(str(f0), mode='r')
            if 'observation_data' in nc_data.groups:
                nc_data = nc_data.groups['observation_data']
            elif 'geolocation_data' in nc_data.groups:
                nc_data = nc_data.groups['geolocation_data']
            bands = [k for k in nc_data.variables]
            for s in list(set(bands).intersection(set(self.bands))):
                data = nc_data[s][:].astype(dtype)
                if s in data_dict: raise Exception('...')
                data_dict[s] = data
        mask = data_dict['I01'].mask | data_dict['I02'].mask | data_dict['I03'].mask | data_dict['I04'].mask | data_dict['I05'].mask
        for k in data_dict:
            data_dict[k].mask = mask
        return data_dict

    def resample(self, data: dict, epsg=4326, max_distance_meter=1000, num_workers=8):
        swath_def = SwathDefinition(lons=data['longitude'], lats=data['latitude'])
        crs = CRS(f'EPSG:{epsg}')
        area_def = create_area_def(crs.name, crs.to_dict(), area_extent=self.region.bbox, resolution=self.region.pixel_size)
        valid_input_index, valid_output_index, index_array, distance_array = \
            kd_tree.get_neighbour_info(swath_def, area_def, max_distance_meter, neighbours=1, nprocs=num_workers)
        for k in self.bands:
            data[k] = kd_tree.get_sample_from_neighbour_info('nn', output_shape=self.region.shape, data=data[k],
                                                   valid_input_index=valid_input_index, valid_output_index=valid_output_index,
                                                   index_array=index_array, distance_array=distance_array, fill_value=None)
        return data

    def process_one(self, time:pd.Timestamp, proc_funcs:list=[], save=True, replace=False,
                    **proc_funcs_kwargs):
        """This method defines a processing pipeline consisting of opening the file
        using the `open` method, applying each of the `proc_funcs` to the output of the previous
        and `save` the processed data using save method."""
        tstr = time.strftime('%Y%m%d')
        files = self.list_files(time)
        files = group_files(files)
        tstr = time.strftime('%Y%m%d')
        filename = f'{self.paths.dst}/{self.name}{self.region.name}_{tstr}.nc'
        if not Path(filename).is_file() or replace:
            data_dict = {v:[] for v in self.bands}
            for file_group in files:
                data = self.open_netcdf4(file_group)
                data = self.resample(data)
                for v in self.bands:
                    data_dict[v].append(data[v])
            kwargs = {'cls': self, **proc_funcs_kwargs}
            for f in proc_funcs:
                data_dict = f(data_dict, time, **kwargs)
            if save:
                self.save(time, data_dict)
            else: return data_dict
        else: warn(f'Skipping files for {time}. File already exists.')

class Viirs375Dataset(ViirsDataset):
    _name = 'VIIRS375'

class Viirs750Dataset(ViirsDataset):
    _name = 'VIIRS750'

# Cell
class MCD64Dataset(BaseDataset):
    "Subclass of `BaseDataset` to process MCD64A1 Collection 6 burned areas product."
    def __init__(self, paths:InOutPath, region:Region, times:pd.DatetimeIndex=None,
                 bands=['bafrac']):
        super().__init__('MCD64A1C6', paths, region, times, bands)

    def list_files(self, time:pd.Timestamp) -> list:
        out = []
        if time in self.times:
            time = pd.Timestamp(f'{time.year}-{time.month}-01')
            time_pattern = f'.A{time.year}{time.dayofyear:03d}.'
            files = self.paths.src.ls(recursive=True, include=['burndate.tif', time_pattern],
                                exclude=['.xml'])
            # Find windows joint with region bounding box
            for f in files:
                data = open_tif(f)
                if not disjoint_bounds(data.bounds, self.region.bbox):
                    out.append(f)
        return out

    def find_dates(self, first:pd.Timestamp=None, last:pd.Timestamp=None):
        pattern = r'^\w+.A(20[0-9][0-9])([0-3][0-9][0-9])..*$'
        times = []
        for f in self.paths.src.ls(recursive=True):
            x = re.search(pattern, f.stem)
            if x is not None:
                year, doy = map(x.group, [1,2])
                times.append(pd.Timestamp(f'{year}-01-01') + pd.Timedelta(days=int(doy)-1))
        self.times = pd.DatetimeIndex(sorted(set(times)))
        if first is not None:
            self.times = self.times[self.times>=first]
        if last is not None:
            self.times = self.times[self.times<=last]
        return self.times

    def file_time_range(self, file) -> pd.DatetimeIndex:
        pattern = r'^\w+.A(20[0-9][0-9])([0-3][0-9][0-9])..*$'
        x = re.search(pattern, file.stem)
        year, doy = map(x.group, [1,2])
        t0 = pd.Timestamp(f'{year}-01-01') + pd.Timedelta(days=int(doy)-1)
        return pd.date_range(t0, periods=monthlen(t0.year, t0.month), freq='D')

    def open(self, files:list) -> dict:
        times = self.file_time_range(files[0])
        data_dict = {'times': times}
        out = np.zeros((len(times), *self.region.shape))
        data = [open_tif(f) for f in files]
        data, tfm = crop(data, self.region.bbox)
        for i, time in enumerate(times):
            x = (data == time.dayofyear).astype(np.int8)
            out[i] += downsample(x, tfm, self.region.transform, self.region.shape)
        data_dict[self.bands[0]] = out
        return data_dict

    def save(self, time:pd.Timestamp, data:dict):
        v = self.bands[0]
        for i, t in enumerate(data['times']):
            super().save(t, {v: data[v][i]})

# Cell
class FireCCI51Dataset(BaseDataset):
    "Subclass of `BaseDataset` to process FireCCI51 burned areas product."
    def __init__(self, paths:InOutPath, region:Region, times:pd.DatetimeIndex=None,
                 bands=['bafrac']):
        super().__init__('FireCCI51', paths, region, times, bands)

    def list_files(self, time:pd.Timestamp) -> list:
        out = []
        if time in self.times:
            time = pd.Timestamp(f'{time.year}-{time.month}-01')
            time_pattern = time.strftime('%Y%m%d')
            files = self.paths.src.ls(recursive=True, include=['JD.tif', time_pattern],
                                exclude=['.xml'])
            # Find windows joint with region bounding box
            for f in files:
                data = open_tif(f)
                if not disjoint_bounds(data.bounds, self.region.bbox):
                    out.append(f)
        return out

    def find_dates(self, first:pd.Timestamp=None, last:pd.Timestamp=None):
        files = self.paths.src.ls(recursive=True, include=['JD.tif'], exclude=['.xml'])
        self.times = pd.DatetimeIndex(sorted(set([pd.Timestamp(o.name[:8]) for o in files])))
        if first is not None:
            self.times = self.times[self.times>=first]
        if last is not None:
            self.times = self.times[self.times<=last]
        return self.times

    def file_time_range(self, file) -> pd.DatetimeIndex:
        t0 = pd.Timestamp(file.name[:8])
        return pd.date_range(t0, periods=monthlen(t0.year, t0.month), freq='D')

    def open(self, files:list) -> dict:
        times = self.file_time_range(files[0])
        data_dict = {'times': times}
        out = np.zeros((len(times), *self.region.shape))
        data = [open_tif(f) for f in files]
        data, tfm = crop(data, self.region.bbox)
        for i, time in enumerate(times):
            x = (data == time.dayofyear).astype(np.int8)
            out[i] += downsample(x, tfm, self.region.transform, self.region.shape)
        data_dict[self.bands[0]] = out
        return data_dict

    def save(self, time:pd.Timestamp, data:dict, do_compression=True):
        v = self.bands[0]
        for i, t in enumerate(data['times']):
            super().save(t, {v: data[v][i]}, do_compression=do_compression)

# Cell
class AusCoverDataset(BaseDataset):
    "Subclass of `BaseDataset` to process AusCover burned areas product."
    def __init__(self, paths:InOutPath, region:Region, times:pd.DatetimeIndex=None,
                bands=['bafrac']):
        super().__init__('AusCover', paths, region, times, bands)

    def list_files(self, time:pd.Timestamp) -> list:
        out = []
        if time.year in self.times.year:
            time = pd.Timestamp(f'{time.year}-01-01')
            time_pattern = time.strftime('_%Y_')
            files = self.paths.src.ls(recursive=True, include=['.tif', time_pattern],
                                exclude=['.xml'])
        return files

    def find_dates(self, first:pd.Timestamp=None, last:pd.Timestamp=None):
        files = self.paths.src.ls(recursive=True, include=['.tif'], exclude=['.xml'])
        self.times = pd.DatetimeIndex(sorted(set([pd.Timestamp(f'{o.stem[-10:-6]}-01-01')
                                                  for o in files])))
        if first is not None:
            self.times = self.times[self.times>=first]
        if last is not None:
            self.times = self.times[self.times<=last]
        return self.times

    def file_time_range(self, file) -> pd.DatetimeIndex:
        t0 = pd.Timestamp(f'{file.stem[-10:-6]}-01-01')
        return pd.date_range(t0, periods=12, freq='MS')

    def open(self, files:list) -> dict:
        times = self.file_time_range(files[0])
        data_dict = {'times': times}
        out = np.zeros((len(times), *self.region.shape))
        data = [open_tif(f) for f in files]
        data = data[0]
        crs = data.crs
        tfm = data.transform
        data = data.read(1)
        for i, time in enumerate(times):
            x = (data == time.month).astype(np.int8)
            out[i] += downsample(x, tfm, self.region.transform,
                                 self.region.shape, src_crs=crs)
        data_dict[self.bands[0]] = out
        return data_dict

    def save(self, time:pd.Timestamp, data:dict, do_compression=True):
        v = self.bands[0]
        for i, t in enumerate(data['times']):
            super().save(t, {v: data[v][i]}, do_compression=do_compression)

# Cell
class MTBSDataset(BaseDataset):
    "Subclass of `BaseDataset` to process MTBS burned areas product."
    def __init__(self, paths:InOutPath, region:Region, times:pd.DatetimeIndex=None,
                bands=['bafrac']):
        super().__init__('MTBS', paths, region, times, bands)

    def list_files(self, *args) -> list:
        files = self.paths.src.ls(recursive=True, include=['.shp'], exclude=['.xml'])
        return files

    def find_dates(self, first:pd.Timestamp=None, last:pd.Timestamp=None):
        files = self.list_files()
        df = open_shp(files[0])
        self.times = pd.date_range(f'{df.Year.min()}-01-01',
                                   f'{df.Year.max()}-12-01', freq='MS')
        if first is not None:
            self.times = self.times[self.times>=first]
        if last is not None:
            self.times = self.times[self.times<=last]
        return self.times

    def open(self, files:list) -> dict:
        data_dict = {'times': self.times}
        data = open_shp(files[0]).to_crs({'init': 'EPSG:4326'})
        out = np.zeros((len(self.times), *self.region.shape))
        R = Region(self.region.name, self.region.bbox, pixel_size=0.0003)
        for i, time in enumerate(self.times):
            x = data.loc[(data.Year==time.year) & (data.StartMonth==time.month)].copy()
            x_raster = rasterize(x, region=R)
            out[i] += downsample(x_raster, R.transform, self.region.transform,
                                self.region.shape)
        data_dict[self.bands[0]] = out
        return data_dict

    def save(self, time:pd.Timestamp, data:dict, do_compression=True):
        v = self.bands[0]
        for i, t in enumerate(data['times']):
            super().save(t, {v: data[v][i]}, do_compression=do_compression)

    def process_all(self, *args):
        self.process_one(self.times[0])

# Cell
class ICNFDataset(BaseDataset):
    "Subclass of `BaseDataset` to process ICNF burned areas product."
    def __init__(self, paths:InOutPath, region:Region, times:pd.DatetimeIndex=None,
                bands=['bafrac']):
        super().__init__('ICNF', paths, region, times, bands)

    def list_files(self, *args) -> list:
        files = self.paths.src.ls(recursive=True, include=['.shp'], exclude=['.xml'])
        return files

    def find_dates(self, first:pd.Timestamp=None, last:pd.Timestamp=None):
        files = self.list_files()
        df = open_shp(files[0])
        self.times = sorted(set([pd.Timestamp(f'{o[:-2]}01')
                            for o in df.FIREDATE if o is not None]))
        if first is not None:
            self.times = self.times[self.times>=first]
        if last is not None:
            self.times = self.times[self.times<=last]
        return self.times

    def open(self, files:list) -> dict:
        data_dict = {'times': self.times}
        data = open_shp(files[0]).to_crs({'init': 'EPSG:4326'})
        data = data.loc[~data.FIREDATE.isna()]
        times = pd.DatetimeIndex([pd.Timestamp(o) for o in data.FIREDATE])
        data['times'] = times
        out = np.zeros((len(self.times), *self.region.shape))
        R = Region(self.region.name, self.region.bbox, pixel_size=0.0003)
        for i, time in enumerate(self.times):
            x = data.loc[(times.year==time.year) &
                         (times.month==time.month)].copy()
            x_raster = rasterize(x, region=R)
            out[i] += downsample(x_raster, R.transform, self.region.transform,
                                self.region.shape)
        data_dict[self.bands[0]] = out
        return data_dict

    def save(self, time:pd.Timestamp, data:dict, do_compression=True):
        v = self.bands[0]
        for i, t in enumerate(data['times']):
            super().save(t, {v: data[v][i]}, do_compression=do_compression)

    def process_all(self, *args):
        self.process_one(self.times[0])

# Cell
class Region2Tiles():
    """Crops processed dataset into 128x128 tiles to be used to train the model"""
    def __init__(self, paths:InOutPath, input_name:str, target_name:str,
                 regions:list=None, bands:list=None, size=128, step=100):
        self.paths = paths
        self.input_name = input_name
        self.target_name = target_name
        self.bands = bands
        self.size = size
        self.step = step
        if regions is None:
            self.regions = [o.name for o in self.paths.src.ls()]
        else:
            self.regions = regions

        for folder in ['images', 'masks']:
            (self.paths.dst/folder).mkdir(exist_ok=True)

    def open(self, file, bands:list):
        "Open .mat file and select `bands`."
        f = sio.loadmat(file)
        return np.array([f[k] for k in bands]).transpose(1,2,0)

    def process_one(self, file, bands, folder):
        """Create tiles for a `file` saving the results in `folder` for each
         crop using `save` method"""
        try:
            data = self.open(file, bands)
            rr, cc, _ = data.shape
            for c in range(0, cc-1, self.step):
                for r in range(0, rr-1, self.step):
                    img = self.crop(data, r, c)
                    if np.nansum(~np.isnan(img)) > 0:
                        self.save(img, file, r, c, folder, bands)
        except:
            warn(f'Unable to process {file}.')

    def process_all(self, max_workers=8, include=[]):
        "Run `process_one` in parallel using the number of workers given by `max_workers`"
        for r in self.regions:
            print(f'Creating tiles for {r}')
            for i, s in enumerate([self.input_name, self.target_name]):
                files_list = self.paths.src.ls(recursive=True, include=[*include, *['.mat', r, s]])
                folder = 'images' if s == self.input_name else 'masks'
                bands = self.bands[i]
                process_one = partial(self.process_one, bands=bands, folder=folder, )
                with ThreadPoolExecutor(max_workers) as e:
                    list(tqdm(e.map(process_one, files_list), total=len(files_list)))

    def crop(self, im, r, c):
        "crop image into a square of size sz"
        sz = self.size
        out_sz = (sz, sz, im.shape[-1])
        rs,cs,hs = im.shape
        tile = np.zeros(out_sz)
        if (r+sz > rs) and (c+sz > cs):
            tile[:rs-r, :cs-c, :] = im[r:, c:, :]
        elif (r+sz > rs):
            tile[:rs-r, :, :] = im[r:, c:c+sz, :]
        elif (c+sz > cs):
            tile[:, :cs-c, :] = im[r:r+sz ,c:, :]
        else:
            tile[...] = im[r:r+sz, c:c+sz, :]
        return tile

    def save(self, data, file, r, c, folder, bands):
        """Save `data` to `file` on `folder` selecting only the given `bands`.
           The row and column index are included in the filename so that the large
           scene can be reconstructed afterwards after generating the model predictions."""
        sio.savemat(self.paths.dst/f'{folder}/{file.stem}_{r}_{c}.mat',
            {v: data[...,i] for i, v in enumerate(bands)}, do_compression=True)