# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ml.mcd.ipynb.

# %% auto 0
__all__ = ['discrepancy', 'DiscrepancyLoss', 'discrepancy_metric', 'Feature', 'Predictor', 'McdDataset', 'McdDataLoader',
           'McdModel', 'EnsembleMcdModel', 'McdCallback', 'mcd_learner']

# %% ../../nbs/ml.mcd.ipynb 3
import torch
import torch.nn as nn
import torch.nn.functional as F
from .layer import GradReverse
from fastai.callback.core import Callback
from fastai.data.core import DataLoaders
from fastai.learner import CancelBatchException, CancelStepException
from fastai.learner import Learner, Recorder
from fastai.losses import BaseLoss, CrossEntropyLossFlat
from fastai.torch_core import default_device
from fastcore.basics import GetAttr, store_attr
from fastcore.meta import delegates
from typing import Callable

# %% ../../nbs/ml.mcd.ipynb 4
def discrepancy(a, b):
    return torch.mean(torch.abs(F.softmax(a, dim=-1) - F.softmax(b, dim=-1)))

# %% ../../nbs/ml.mcd.ipynb 6
class DiscrepancyLoss:
    def __call__(self, outs, *targets, **kwargs):
        assert len(outs) == 2
        return -discrepancy(*outs)

def discrepancy_metric(pred, *targets):
    a, b = pred[-2], pred[-1]
    return discrepancy(a, b)

# %% ../../nbs/ml.mcd.ipynb 7
class Feature(nn.Module):
    "Image feature extractor"

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=5, stride=1, padding=2)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=5, stride=1, padding=2)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=5, stride=1, padding=2)
        self.bn3 = nn.BatchNorm2d(128)
        self.fc1 = nn.Linear(8192, 3072)
        self.bn1_fc = nn.BatchNorm1d(3072)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.bn1(self.conv1(x))), stride=2, kernel_size=3, padding=1)
        x = F.max_pool2d(F.relu(self.bn2(self.conv2(x))), stride=2, kernel_size=3, padding=1)
        x = F.relu(self.bn3(self.conv3(x)))
        x = x.view(x.size(0), 8192)
        x = F.relu(self.bn1_fc(self.fc1(x)))
        x = F.dropout(x, training=self.training)
        return x

class Predictor(nn.Module):
    "Classifier on image features"

    def __init__(self, prob=0.5, lambd=1.0):
        super().__init__()
        self.prob = prob
        self.lambd = lambd
        self.fc1 = nn.Linear(8192, 3072)
        self.bn1_fc = nn.BatchNorm1d(3072)
        self.fc2 = nn.Linear(3072, 2048)
        self.bn2_fc = nn.BatchNorm1d(2048)
        self.fc3 = nn.Linear(2048, 10)
        self.bn_fc3 = nn.BatchNorm1d(10)
        self.gr = GradReverse(lambd)

    def forward(self, x, grad_reverse=False):
        if grad_reverse:
            x = self.gr(x)
        x = F.relu(self.bn2_fc(self.fc2(x)))
        x = self.fc3(x)
        return x


# %% ../../nbs/ml.mcd.ipynb 8
class McdDataset:
    """Dataset for MCD. 
    A data point is a tuple of 4 tensors, in the following order:
    
    (source_domain_x, target_domain_x, source_domain_y, target_domain_y)
    """
    def __init__(self, source_ds, target_ds):
        store_attr()
    
    def __getitem__(self, idx):
        xs, ys = self.source_ds[idx]
        xt, yt = self.target_ds[idx]
        return xs, xt, ys, yt
    
    def __len__(self):
        return min(len(self.source_ds), len(self.target_ds))

# %% ../../nbs/ml.mcd.ipynb 9
class McdDataLoader(GetAttr):
    """Dataloader for MCD."""
    
    _default = 'target_dl'

    def __init__(self, source_dl, target_dl):
        store_attr()
        self.source_it = None
        self.target_it = None
    
    def __iter__(self):
        self.source_it = iter(self.source_dl)
        self.target_it = iter(self.target_dl)
        return self
    
    def __len__(self):
        return min(len(self.source_dl), len(self.target_dl))
    
    def __next__(self):
        xsb, ysb = next(self.source_it)
        xtb, ytb = next(self.target_it)
        return xsb, xtb, ysb, ytb
    
    def to(self, device):
        self.source_dl.to(device)
        self.target_dl.to(device)
    
    def new(self, dataset=None, cls=None, **kwargs):
        if cls is None: cls = type(self)
        return cls(self.source_dl.new(**kwargs), self.source_dl.new(**kwargs))

# %% ../../nbs/ml.mcd.ipynb 10
class McdModel(nn.Module):
    """Image classification model with a gradient reversal layer and two classifier heads."""

    def __init__(self, feature_extractor, classifier1, classifier2, lambd=1.0):
        super().__init__()
        store_attr()
        self.gr = GradReverse(lambd)
    
    def forward(self, img, grad_reverse=False):
        feat = self.feature_extractor(img)
        if grad_reverse:
            feat = self.gr(feat)
        output1 = self.classifier1(feat)
        output2 = self.classifier2(feat)
        return output1, output2

# %% ../../nbs/ml.mcd.ipynb 11
class EnsembleMcdModel(nn.Module):
    """Ensemble of two classifiers trained by MCD."""

    def __init__(self, feature_extractor, classifier1, classifier2):
        super().__init__()
        store_attr()
    
    def forward(self, img):
        feat = self.feature_extractor(img)
        output1 = self.classifier1(feat)
        output2 = self.classifier2(feat)
        return (output1 + output2) / 2

    @classmethod
    def from_mcd_model(cls, mcd_model):
        return cls(mcd_model.feature_extractor, mcd_model.classifier1, mcd_model.classifier2)


# %% ../../nbs/ml.mcd.ipynb 12
class McdCallback(Callback):
    """It expects data in the form of `McdDataset`."""
    order = Recorder.order + 10

    def __init__(self, classification_loss_func, discrepancy_loss_func):
        super().__init__()
        store_attr()
    
    def before_batch(self):
        self._do_one_batch()
        raise CancelBatchException

    def before_fit(self):
        "Set device for loss funcs"
        device = getattr(self.dls, 'device', default_device())
        if isinstance(self.classification_loss_func, (nn.Module, BaseLoss)): 
            self.classification_loss_func.to(device)
        if isinstance(self.discrepancy_loss_func, (nn.Module, BaseLoss)): 
            self.discrepancy_loss_func.to(device)
    
    def _do_one_batch(self):
        assert len(self.xb) == 2
        assert len(self.yb) == 2
        source_pred, source_loss = self._predict_source()
        target_pred, target_loss = self._predict_target()
        self.learn.pred = tuple([*source_pred, *target_pred])
        self.learn('after_pred')
        if source_loss is not None and target_loss is not None:
            self.learn.loss = source_loss.detach().cpu() + target_loss.detach().cpu()
        self.learn('after_loss')
        if not self.training or not len(self.yb): 
            return
        self._do_grad_opt()
    
    def _predict_source(self):
        img = self.xb[0]
        pred = self.model(img, grad_reverse=False)
        loss = None
        if len(self.yb):
            loss = self.classification_loss_func(pred, *self.yb)
            if self.training:
                loss.backward(retain_graph=True)
        return pred, loss
    
    def _predict_target(self):
        img = self.xb[1]
        pred = self.model(img, grad_reverse=True)
        loss = None
        if len(self.yb):
            loss = self.discrepancy_loss_func(pred, *self.yb)
            if self.training:
                loss.backward()
        return pred, loss

    def _do_grad_opt(self):
        self._with_events(self.learn.opt.step, 'step', CancelStepException)
        self.learn.opt.zero_grad()


# %% ../../nbs/ml.mcd.ipynb 13
@delegates(Learner.__init__)
def mcd_learner(
        dls:DataLoaders, # `DataLoaders` containing fastai or PyTorch `DataLoader`s
        model:Callable, # PyTorch model for training or inference
        **kwargs
    ):
    """Creates a Learner for MCD by arranging loss functions, metrics, and adding `McdCallback` to callbacks."""
    from fastmtl.loss import CombinedLoss, LossRouting
    from fastmtl.metric import route_to_metric
    from fastai.metrics import accuracy
    
    source_classification_loss_func = CombinedLoss(
        LossRouting(CrossEntropyLossFlat(), pred_idx=0, target_idx=0, weight=1.0),
        LossRouting(CrossEntropyLossFlat(), pred_idx=1, target_idx=0, weight=1.0),
    )
    discrepancy_loss_func = DiscrepancyLoss()
    mcd_callback = McdCallback(source_classification_loss_func, discrepancy_loss_func)
    cbs = [mcd_callback, *list(kwargs.pop('cbs', []))]

    classification_metrics = [route_to_metric(accuracy, 0, 0), route_to_metric(accuracy, 1, 0)]
    discrepancy_metrics = [route_to_metric(accuracy, 2, 1), route_to_metric(accuracy, 3, 1), discrepancy_metric] 

    learn = Learner(
        dls,
        model,
        loss_func=source_classification_loss_func,
        metrics=[*classification_metrics, *discrepancy_metrics],
        cbs=cbs,
        **kwargs,
    )
    for metric, name in zip(
        learn.metrics, 
        ['clf1_acc_s', 'clf2_acc_s', 'clf1_acc_t', 'clf2_acc_t', 'discrep_t']
    ):
        try:
            metric.name = name
        except AttributeError:
            metric.func.__name__ = name
    return learn

