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

__all__ = ['GBlendLoss', 'GBlend']

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

# Cell
class GBlendLoss(Module):
    "Wrapper loss used by the gradient blending callback to allow weights applied to each modality."

    def __init__(self, crit=None, w=None):
        self.crit = ifnone(crit, CrossEntropyLossFlat())
        self.w = w

    def forward(self, preds, target):
        # unweighted loss
        if not is_listy(preds): return self.crit(preds, target)

        # weighted loss
        if self.w is None: self.w = tensor([1.] * len(preds))
        loss = 0
        for i, pred in enumerate(preds): loss += self.crit(pred, target) * self.w[i]
        return loss / sum(self.w)

# Cell
class GBlend(Callback):
    r"""A callback to implement multi-modal gradient blending.

    This is an unofficial PyTorch implementation by Ignacio Oguiza of  - oguiza@gmail.com based on: Wang, W., Tran, D., & Feiszli, M. (2020).
    What Makes Training Multi-Modal Classification Networks Hard?. In Proceedings of the IEEE/CVF Conference on Computer Vision and
    Pattern Recognition (pp. 12695-12705).
    """

    def __init__(self, V_pct=.1, n:Union[None, int, tuple, list]=(10, 5), sel_metric:Optional[str]=None, show_plot:bool=False):

        r"""
        Args:
            V_pct      : subset of train where OGR will be measured (to estimate L*)
            n          : None: offline learning, int: super-epoch (online learning), tuple: (warmup super-epoch, super-epoch)(online learning with warm up)
            sel_metric : which metric will be used to calculate overfitting and generalization during training. If None, loss will be used.
            show_plot  : will show a plot with the wieghts at the end of training
        """
        assert V_pct < 1, 'V_pct must be < 1'
        self.V_pct, self.n, self.sel_metric, self.show_plot = V_pct, n, sel_metric, show_plot
        self.metric_idx = None

    def before_fit(self):

        # model
        self.M = self.model.M
        self.old_multi_output = self.learn.model.multi_output
        self.learn.model.multi_output = True

        #loss
        if cls_name(self.learn.loss_func) != 'GBlendLoss': self.learn.loss_func = GBlendLoss(crit=self.learn.loss_func)

        # calculate super_epochs
        if self.n is None:
            self.super_epochs = [0]
        else:
            if is_listy(self.n):
                self.wu_n = self.n[0]
                self.n = self.n[1]
            else:
                self.wu_n = self.n
            rng = range(int(max(0, self.n_epoch - self.wu_n) / self.n + 1))
            self.super_epochs = []
            for i in rng:
                self.super_epochs.append((i * self.wu_n) if i <= 1 else int((i + self.wu_n / self.n - 1) * self.n))
        self.super_epochs.append(self.n_epoch)

        # create T'(Tp) and V dataloaders
        n_out = len(self.learn.dls.train.dataset.ptls) - self.learn.dls.train.dataset.n_inp
        train_targets = self.learn.dls.train.dataset.ptls[-n_out]
        Tp_idx, V_idx = get_splits(train_targets, valid_size=self.V_pct)
        _Tp_train_dls = []
        _V_train_dls = []
        self.learn.new_dls = []
        for dl in self.learn.dls[0].loaders: # train MixedDataLoaders
            _Tp_dl = get_subset_dl(dl, Tp_idx)
            _V_dl = get_subset_dl(dl, V_idx)
            _Tp_train_dls.append(_Tp_dl)
            _V_train_dls.append(_V_dl)
            self.learn.new_dls.append(DataLoaders(_Tp_dl, _V_dl, device=self.learn.dls.device))
        self.learn.new_dls.append(MixedDataLoaders(MixedDataLoader(*_Tp_train_dls, shuffle=True),  # train - train
                                                   MixedDataLoader(*_V_train_dls, shuffle=False),  # train - valid
                                                   device=self.learn.dls.device))

        # prepare containers
        self.learn.LT = []
        self.learn.LV = []

    def before_train(self):
        if self.epoch in self.super_epochs[:-1] and not 'LRFinder' in [cls_name(cb) for cb in self.learn.cbs]:
            self.train_epochs = np.diff(self.super_epochs)[self.super_epochs.index(self.epoch)]

            #compute weights
            self.learn.save('gblend_learner')
            torch.save(self.learn.model, 'gblend_model')
            w = self.compute_weights()
            if self.epoch == 0: self.learn.ws = [w]
            else: self.learn.ws.append(w)
            self.learn = self.learn.load('gblend_learner')
            self.learn.loss_func.w = w

    def compute_weights(self):

        # _LT0 = []
        # _LV0 = []
        _LT = []
        _LV = []
        for i in range(self.M + 1):
            model = torch.load('gblend_model')
            learn = Learner(self.learn.new_dls[i], model.m[i], loss_func=GBlendLoss(),
                            opt_func=self.learn.opt_func, metrics=self.learn.metrics)
            learn.model.multi_output = False
            learn.remove_cbs(learn.cbs[1])
            learn.add_cb(Recorder(train_metrics=True))
            with learn.no_bar():
                with learn.no_logging():
                    learn.fit_one_cycle(self.train_epochs, pct_start=0)
            if self.metric_idx is None and self.sel_metric is not None:
                metric_names = learn.recorder.metric_names[1:-1]
                self.metric_idx = [i for i,m in enumerate(metric_names) if self.sel_metric in m]
            else: self.metric_idx = [0, 1]
            metric_values = learn.recorder.values[-1][self.metric_idx]
            _LT.append(metric_values[0])
            _LV.append(metric_values[1])

        # if self.epoch == 0: self.compute_previous_metrics()
        self.compute_previous_metrics()
        self.learn.LT.append(_LT)
        self.learn.LV.append(_LV)

        LT1 = array(self.learn.LT[-2])
        LT2 = array(self.learn.LT[-1])
        LV1 = array(self.learn.LV[-2])
        LV2 = array(self.learn.LV[-1])

        ΔG = (LV1 - LV2) if self.metric_idx[0] == 0 else (LV2 - LV1)
        O1 = (LV1 - LT1) if self.metric_idx[0] == 0 else (LT1 - LV1)
        O2 = (LV2 - LT2) if self.metric_idx[0] == 0 else (LT2 - LV2)

        ΔG = np.maximum(0, ΔG)

        ΔO = O2 - O1
        ΔO2 = np.maximum(1e-8, (O2 - O1)**2)
        w = np.maximum(1e-8, np.nan_to_num(ΔG / ΔO2))
        w = w / w.sum()
        w = w.tolist()
        return w

    def compute_previous_metrics(self):
        if self.metric_idx[0] == 0:  metric = self.loss_func
        else: metric = self.learn.metrics[(min(array(self.metric_idx) - 2) - 1) // 2]
        _LT = []
        _LV = []
        with torch.no_grad():
            for i in range(self.M + 1):
                model = torch.load('gblend_model')
                model.multi_output = False
                model = model.m[i]
                _train_metrics = []
                _valid_metrics = []
                for j,dl in enumerate(self.learn.new_dls[i]):
                    it = iter(dl)
                    _preds = []
                    _targets = []
                    for b in it:
                        _preds.extend(model(*b[:-1]))
                        _targets.extend(b[-1])
                    _preds, _targets = stack(_preds), stack(_targets)
                    try: _metric_values = metric(_preds, _targets).cpu().item()
                    except: _metric_values = metric(torch.argmax(_preds, 1), _targets).cpu().item()
                    if j == 0: _LT.append(_metric_values)
                    else: _LV.append(_metric_values)
            self.learn.LT.append(_LT)
            self.learn.LV.append(_LV)

    def after_fit(self):
        if hasattr(self.learn, "ws") and self.show_plot:
            widths = np.diff(self.super_epochs)
            cum_ws = 0
            for i in range(self.M + 1):
                plt.bar(self.super_epochs[:-1] + widths/2, stack(self.learn.ws)[:, i], bottom=cum_ws, width=widths,
                        label=f'k={i+1}' if i < self.M else f'fused')
                cum_ws += stack(self.learn.ws)[:, i]
            plt.xlim(0, self.super_epochs[-1])
            plt.ylim(0, 1)
            plt.xticks(self.super_epochs)
            plt.legend(loc='best')
            plt.title('Online G-Blend Weights by modality')
            plt.show()

        self.learn.model.multi_output = self.old_multi_output