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

__all__ = ['open_nc', 'crop', 'image2tiles', 'tiles2image', 'get_preds', 'predict_one', 'save_netcdf', 'inspect_netcdf',
           'load_netcdf', 'predict_time', 'predict_month', 'predict_nrt', 'split_mask']

# Cell
import scipy.io as sio
import sys
from tqdm import tqdm
import netCDF4
import pdb
import gc
import calendar
import scipy.ndimage as ndimage
from fastai.vision.all import *

from .core import *
from .geo import Region
from .models import BA_Net

# Cell
def open_nc(fn, slice_idx=None, *args, **kwargs):
    data = netCDF4.Dataset(fn, mode='r')
    if slice_idx is None:
        data = np.array([data[r][:] for r in ['Red', 'NIR', 'MIR', 'FRP']], dtype=np.float32)
    else:
        data = np.array([data[r][slice_idx[0]:slice_idx[1], slice_idx[2]:slice_idx[3]]
                         for r in ['Red', 'NIR', 'MIR', 'FRP']], dtype=np.float32)
    data[data == 65535] = 0
    data = data/100
    return data.astype(np.float16)

def crop(im, r, c, size=128):
    '''
    crop image into a square of size sz,
    '''
    sz = size
    out_sz = (sz, sz, im.shape[-1])
    rs,cs,hs = im.shape
    tile = np.zeros(out_sz, dtype=np.float16)
    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 image2tiles(x, step=100, size=128):
    tiles = []
    rr, cc, _ = x.shape
    for c in range(0, cc-1, step):
        for r in range(0, rr-1, step):
            img = crop(x, r, c, size=size)
            tiles.append(img)
    return np.array(tiles, dtype=np.float16)

def tiles2image(tiles, image_size, size=128, step=100):
    rr, cc, = image_size
    sz = size
    im = np.zeros(image_size, dtype=np.float16)
    indicator = np.zeros_like(im)
    k = 0
    for c in range(0, cc-1, step):
        for r in range(0, rr-1, step):
            if (r+sz > rr) and (c+sz > cc):
                im[r:, c:] += tiles[k][:rr-r, :cc-c]
                indicator[r:, c:] += 1
            elif (r+sz > rr):
                im[r:, c:c+sz] += tiles[k][:rr-r, :]
                indicator[r:, c:c+sz] += 1
            elif (c+sz > cc):
                im[r:r+sz ,c:] += tiles[k][:, :cc-c]
                indicator[r:r+sz, c:] += 1
            else:
                im[r:r+sz, c:c+sz] += tiles[k]
                indicator[r:r+sz, c:c+sz] += 1
            k += 1
    im /= indicator
    return im

def get_preds(tiles, model, weights=None, verbose=False):
    if weights is not None:
        model.load_state_dict(weights)
    mu = tensor([0.2349, 0.3548, 0.1128, 0.0016]).view(1,4,1,1,1)
    std = tensor([0.1879, 0.1660, 0.0547, 0.0776]).view(1,4,1,1,1)
    with torch.no_grad():
        data = []
        iterator = tqdm(tiles) if verbose else tiles
        for x in iterator:
            model.eval()
            if torch.cuda.is_available():
                model.cuda()
            x = (x[None]-mu)/std
            if torch.cuda.is_available():
                x = x.cuda()
            out = model(x).sigmoid().float()
            data.append(out.cpu().squeeze().numpy())
    return np.array(data, np.float16)

def predict_one(iop:InOutPath, times:list, weights_files:list, region:str, threshold=0.5,
                slice_idx=None, product='VIIRS750', verbose=False):
    fname = lambda t : iop.src/f'{product}{region}_{t.strftime("%Y%m%d")}.nc'
    files = [fname(t) for t in times]
    im_size = open_nc(files[0], slice_idx=slice_idx).shape[1:]
    tiles = []
    print('Loading data and generating tiles:')
    iterator = tqdm(files) if verbose else files
    for file in iterator:
        try:
            s = image2tiles(open_nc(file, slice_idx=slice_idx).transpose((1,2,0))).transpose((0,3,1,2))
        except:
            warn(f'No data for {file}')
            s = np.zeros_like(s)
        tiles.append(s)
    del s
    tiles = np.array(tiles, dtype=np.float16).transpose((1, 2, 0, 3, 4))
    tiles = torch.from_numpy(tiles).float()
    preds_ens = []
    for wf in weights_files:
        if torch.cuda.is_available():
            weights = torch.load(wf)
        else: weights = torch.load(wf, map_location=torch.device('cpu'))
        if 'model' in weights:
            weights = weights['model']
        print(f'Generating model predictions for {wf}:')
        preds = get_preds(tiles, model=BA_Net(4, 1, 64), weights=weights, verbose=verbose)
        preds = np.array([tiles2image(preds[:,i], im_size) for i in range(preds.shape[1])])
        preds_ens.append(preds)
    del tiles
    preds = np.array(preds_ens, dtype=np.float16).mean(0)
    return preds.astype(np.float16)

def save_netcdf(file, data:dict, region:Region, crop=None):
    mode = "r+" if Path(file).is_file() and crop is not None else "w"
    with netCDF4.Dataset(file, mode, format="NETCDF4") as rootgrp:
        if mode == "w":
            timeD = rootgrp.createDimension("time", None)
            latD = rootgrp.createDimension("lat", region.height)
            lonD = rootgrp.createDimension("lon", region.width)
        for k in data:
            if mode == "w":
                a = rootgrp.createVariable(k,"f4",("time","lat","lon",), zlib=True)
            else: a = rootgrp[k]
            if crop is None:
                a[:] = data[k]
            else:
                a[:,crop[0]:crop[1],crop[2]:crop[3]] = data[k]

def inspect_netcdf(file, var='preds', crop=None):
    if file.is_file():
        with netCDF4.Dataset(file, 'r') as ncdata:
            if crop is not None:
                is_missing = np.sum(ncdata[var][:,crop[0]:crop[1],crop[2]:crop[3]].mask)>0
            else:
                is_missing = np.sum(ncdata[var][:].mask)>0
    else: is_missing = True
    return is_missing

def load_netcdf(file, crop=None):
    with netCDF4.Dataset(file, 'r', format="NETCDF4") as ncdata:
        if crop is None:
            return ncdata['preds'][:].astype(np.float16)
        else:
            return ncdata['preds'][:, crop[0]:crop[1], crop[2]:crop[3]].astype(np.float16)

def predict_time(path:InOutPath, times, weight_files:list, region:Region,
                 threshold=0.05, save=True, max_size=2000, buffer=128,
                 product='VIIRS750', output='data', check_file=False, verbose=False):
    tstart, tend = times.min(), times.max()
    tstart = pd.Timestamp(f'{tstart.year}-{tstart.month}-01')
    tend = pd.Timestamp(f'{tend.year}-{tend.month}-01')
    ptimes = pd.date_range(tstart, tend, freq='MS')[1:-1]
    ptimes_eom = pd.date_range(tstart, tend, freq='M')[1:]
    si = [[max(0,j*max_size-buffer), (j+1)*max_size+buffer,
           max(0,i*max_size-buffer), (i+1)*max_size+buffer]
          for i in range(region.shape[1]//max_size+1) for j in range(region.shape[0]//max_size+1)]
    for time in progress_bar(ptimes):
        tstr = time.strftime('%Y%m')
        file_save = path.dst/f'{output}_{region.name}{tstr}.nc'
        if not file_save.is_file() or check_file:
            if verbose: print(f'Processing time {time}')
            time_start = pd.Timestamp((time - pd.Timedelta(days=15)).strftime('%Y-%m-15')) # Day 15, previous month
            times = pd.date_range(time_start, periods=64, freq='D')
            preds_all = []
            for i, split in enumerate(si):
                if verbose: print(f'Split {split}')
                if check_file:
                    if verbose: print(f'Checking for saved data in the existing netcdf file')
                    is_missing = inspect_netcdf(file_save, crop=split)
                else: is_missing = True
                if is_missing:
                    preds = predict_one(path, times, weight_files, region.name, slice_idx=split,
                                        product=product, verbose=verbose)
                    preds = preds[times.month == time.month]
                    save_netcdf(file_save, {'preds': preds}, region, crop=split)
                    del preds
                    gc.collect()
        else: warn(f'File {file_save} already exists, skipping to next month.')

    # Merging all months and computing confidence level and date of burning
    si = [[max(0,j*max_size-buffer), (j+1)*max_size+buffer,
           max(0,i*max_size-buffer), (i+1)*max_size+buffer]
           for i in range(region.shape[1]//max_size+1) for j in range(region.shape[0]//max_size+1)]
    if verbose: print('Merging monthly files')
    ylen = 366 if calendar.isleap(ptimes[0].year) else 365
    bd_all, ba_all = np.zeros(region.shape, dtype=np.float16), np.zeros(region.shape, dtype=np.float16)
    for i, split in enumerate(si):
        print(f'Split {split}')
        preds_all = []
        for time in progress_bar(ptimes):
            tstr = time.strftime('%Y%m')
            file = path.dst/f'{output}_{region.name}{tstr}.nc'
            data = load_netcdf(file, crop=split)
            preds_all.append(data)
        preds_all = np.concatenate(preds_all, axis=0)
        ba = preds_all.sum(0)
        ba[ba>1] = 1
        ba[ba<threshold] = np.nan
        bd = preds_all.argmax(0)
        bd = bd.astype(np.float16)
        bd[np.isnan(ba)] = np.nan
        ba_all[split[0]:split[1], split[2]:split[3]] = ba
        bd_all[split[0]:split[1], split[2]:split[3]] = bd
        del preds_all, ba, bd
        gc.collect()
    times = pd.date_range(ptimes[0], ptimes_eom[-1], freq='D')
    tstr_start = times[0].strftime('%Y')
    file_save = path.dst/f'{output}_{region.name}{tstr_start}.mat'
    sio.savemat(file_save, {'burned': ba_all, 'date': bd_all, 'times': times.astype(str).tolist()},
                do_compression=True)

def predict_month(iop, time, weight_files, region, threshold=0.5, save=True, slice_idx=None):
    time_start = pd.Timestamp((time - pd.Timedelta(days=30)).strftime('%Y-%m-15')) # Day 15, previous month
    times = pd.date_range(time_start, periods=64, freq='D')
    preds = predict_one(iop, times, weight_files, region, threshold=threshold, slice_idx=slice_idx)
    assert preds.shape[0] == len(times)
    preds = preds[times.month == time.month]
    ba = preds.sum(0)
    bd = preds.argmax(0)
    doy = np.asarray(pd.DatetimeIndex(times).dayofyear)
    bd = doy[bd].astype(float)
    bd[bd==doy[0]] = np.nan
    bd[ba<threshold] = np.nan
    ba[ba<threshold] = np.nan
    ba[ba>1] = 1
    if not save: return ba, bd
    tstr = time.strftime('%Y%m')
    sio.savemat(iop.dst/f'ba_{region}_{tstr}.mat', {'burned': ba, 'date': bd}, do_compression=True)

def predict_nrt(path:InOutPath, time, weight_files:list, region:Region,
                 threshold=0.1, save=True, max_size=2000, buffer=128,
                 product='VIIRS750', verbose=False):
    times = pd.date_range(time-pd.Timedelta(days=63), time, freq='D')
    si = [[max(0,j*max_size-buffer), (j+1)*max_size+buffer,
           max(0,i*max_size-buffer), (i+1)*max_size+buffer]
          for i in range(region.shape[1]//max_size+1) for j in range(region.shape[0]//max_size+1)]

    bas, bds = [], []
    for i, split in progress_bar(enumerate(si), total=len(si)):
        print(f'Split {split}')
        preds_all = predict_one(path, times, weight_files, region.name, slice_idx=split,
                                product=product, verbose=verbose)
        ba = preds_all.sum(0)
        ba[ba>1] = 1
        ba[ba<threshold] = np.nan
        bd = preds_all.argmax(0)
        bd = bd.astype(np.float16)
        bd[np.isnan(ba)] = np.nan
        bas.append(ba.astype(np.float16))
        bds.append(bd.astype(np.float16))
        gc.collect()
    ba_all = np.zeros(region.shape, dtype=np.float16)
    bd_all = np.zeros_like(ba_all)
    for i, split_idx in enumerate(si):
        ba_all[split_idx[0]:split_idx[1], split_idx[2]:split_idx[3]] = bas[i]
        bd_all[split_idx[0]:split_idx[1], split_idx[2]:split_idx[3]] = bds[i]
    if not save: return ba_all, bd_all
    tstr = time.strftime('%Y%m%d')
    sio.savemat(path.dst/f'ba_{tstr}.mat', {'burned': ba_all, 'date': bd_all, 'times': np.array(times).astype(str)}, do_compression=True)

def split_mask(mask, thr=0.5, thr_obj=1):
    labled, n_objs = ndimage.label(mask > thr)
    result = []
    for i in range(n_objs):
        obj = (labled == i + 1).astype(int)
        if (obj.sum() > thr_obj): result.append(obj)
    return result