# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/013_callback.TSBERT.ipynb (unless otherwise specified).

__all__ = ['create_subsequence_mask', 'create_variable_mask', 'TSBERT_Loss', 'TSBERT']

# Cell
from fastai.callback.all import *
from ..imports import *
from ..utils import *
from ..models.utils import *
from ..models.layers import *

# Cell
from torch.distributions.beta import Beta

# Cell
def create_subsequence_mask(o, r=.15, lm=3, stateful=True, sync=False):
    if o.ndim == 2: o = o[None]
    n_masks, mask_dims, mask_len = o.shape
    if sync == 'random': sync = random.random() > .5
    dims = 1 if sync else mask_dims
    numels = n_masks * dims * mask_len
    pm = 1 / lm
    pu = pm * (r / (1 - r))
    a, b, proba_a, proba_b = ([1], [0], pu, pm) if random.random() > pm else ([0], [1], pm, pu)
    if stateful:
        max_len = max(1, 2 * math.ceil(numels // (1/pm + 1/pu)))
        while True:
            dist_a = np.clip(np.random.geometric(proba_a, max_len), 1, mask_len)
            dist_b = np.clip(np.random.geometric(proba_b, max_len), 1, mask_len)
            if (dist_a + dist_b).sum() >= numels:
                dist_len = np.argmax((dist_a + dist_b).cumsum() >= numels) + 1
                break
        l = [a*ax + b*bx for (ax, bx) in zip(dist_a[:dist_len], dist_b[:dist_len])]
        _mask = list(itertools.chain.from_iterable(l))[:numels]
    else:
        _mask = np.random.binomial(1, 1 - r, numels)
    mask = torch.Tensor(_mask).reshape(n_masks, dims, mask_len)
    if sync: mask = mask.repeat(1, mask_dims, 1)
    return mask.to(o.device)

def create_variable_mask(o, r=.15):
    n_masks, mask_dims, mask_len = o.shape
    sel_dims = np.random.choice(n_masks * mask_dims, int(n_masks * mask_dims * r), False)
    _mask = np.ones((n_masks * mask_dims, mask_len))
    _mask[sel_dims] = 0
    mask = torch.Tensor(_mask).reshape(*o.shape)
    return mask.to(o.device)

# Cell
class TSBERT_Loss(Module):
    def __init__(self, crit=None):
        self.crit = ifnone(crit, MSELossFlat())
        self.mask = slice(None)

    def forward(self, preds, target):
        return self.crit(preds[self.mask], target[self.mask])

# Cell
import matplotlib.colors as mcolors

class TSBERT(Callback):
    def __init__(self, r:float=.15, subsequence_mask:bool=True, lm:float=3., stateful:bool=True, sync:bool=False, variable_mask:bool=False,
                 dropout:float=.1, crit:callable=None, target_dir:str='./data/TSBERT', fname:str='model', verbose:bool=True):
        r"""
        Callback used to perform the autoregressive task of denoising the input after a binary mask has been applied.

        Args:
            r: proba of masking.
            subsequence_mask: apply a mask to random subsequences.
            lm: average mask len when using stateful (geometric) masking.
            stateful: geometric distribution is applied so that average mask length is lm.
            sync: all variables have the same masking.
            variable_mask: apply a mask to random variables.
            dropout: dropout applied to the head of the model during pretraining.
            crit: loss function that will be used. If None MSELossFlat().
            target_dir : directory where trained model will be stored.
            fname : file name that will be used to save the pretrained model.
    """
        assert subsequence_mask or variable_mask, 'you must set subsequence_mask and/or variable_mask to True'
        store_attr("subsequence_mask,variable_mask,dropout,r,lm,stateful,sync,crit,fname,verbose")
        self.target_dir = Path(target_dir)


    def before_fit(self):
        # modify loss for denoising task
        self.old_loss_func = self.learn.loss_func
        self.learn.loss_func = TSBERT_Loss(self.crit)
        self.learn.TSBERT = self

        #remove and store metrics
        self.learn.metrics = L([])

        # save initial model head
        assert hasattr(self.learn.model, "head"), f"you can only use {cls_name(self)} with models that have .head attribute"
        backbone = split_model(self.learn.model)[0]
        self.learn.model = backbone

        # prepare model for denoising task
        with torch.no_grad():
            b = self.learn.dls.train.one_batch()
            out = self.learn.model(b[0])
            assert out.ndim == 3, "make sure the backbone (model.head = Noop) produces a 3d output [bs x nf x seq_len]"
            ni = b[0].shape[1]
            if 1 / ni > self.r: self.variable_mask = False
            no = out.shape[1]
            self.learn.model.head = nn.Sequential(nn.Dropout(self.dropout),
                                                  nn.Conv1d(no, ni, 1)).to(b[0].device) # equivalent to linear layer applied to dim=1
            assert self.learn.model(b[0]).shape == b[0].shape, f"{cls_name(self)} cannot recreate a tensor with the input shape"

    def before_batch(self):
        if self.subsequence_mask and self.variable_mask:
            random_thr = 1/3 if self.sync == 'random' else 1/2
            if random.random() > random_thr:
                mask = create_subsequence_mask(self.x, r=self.r, lm=self.lm, stateful=self.stateful, sync=self.sync)
            else:
                mask = create_variable_mask(self.x, r=self.r)
        elif self.subsequence_mask:
            mask = create_subsequence_mask(self.x, r=self.r, lm=self.lm, stateful=self.stateful, sync=self.sync)
        elif self.variable_mask:
            mask = create_variable_mask(self.x, r=self.r)
        else:
            raise ValueError('You need to set subsequence_mask and/ or variable_mask to True in TSBERT.')

        self.learn.yb = (self.x,)
        self.learn.xb = (self.x * mask,)
        self.learn.loss_func.mask = (mask == 0) # boolean mask
        self.mask = mask

    def after_fit(self):
        if self.epoch == self.n_epoch - 1 and not "LRFinder" in [cls_name(cb) for cb in self.learn.cbs]:
            PATH = Path(f'{self.target_dir/self.fname}.pth')
            if not os.path.exists(PATH.parent): os.makedirs(PATH.parent)
            torch.save(self.learn.model.state_dict(), PATH)
            pv(f"\npre-trained model weights_path='{PATH}'\n", self.verbose)

    def show_preds(self, max_n=9, nrows=3, ncols=3, figsize=None, sharex=True, **kwargs):
        b = self.learn.dls.valid.one_batch()
        self.learn._split(b)
        xb = self.xb[0].detach().cpu().numpy()
        bs, nvars, seq_len = xb.shape
        self.learn('before_batch')
        pred = learn.model(*self.learn.xb).detach().cpu().numpy()
        mask = self.mask.cpu().numpy()
        masked_pred = np.ma.masked_where(mask, pred)
        ncols = min(ncols, math.ceil(bs / ncols))
        nrows = min(nrows, math.ceil(bs / ncols))
        max_n = min(max_n, bs, nrows*ncols)
        if figsize is None: figsize = (ncols*6, math.ceil(max_n/ncols)*4)
        fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize, sharex=sharex, **kwargs)
        idxs = np.random.permutation(np.arange(bs))
        colors = list(mcolors.TABLEAU_COLORS.keys()) + random_shuffle(list(mcolors.CSS4_COLORS.keys()))
        i = 0
        for row in ax:
            for col in row:
                color_iter = iter(colors)
                for j in range(nvars):
                    try:
                        color = next(color_iter)
                    except:
                        color_iter = iter(colors)
                        color = next(color_iter)
                    col.plot(xb[idxs[i]][j], alpha=.5, color=color)
                    col.plot(masked_pred[idxs[i]][j], marker='o', markersize=4, color=color)
                i += 1
        plt.tight_layout()
        plt.show()