import numpy as np
from fastai.vision import *
import torch

data_path = Path('../data')
sag_path = data_path/'sagittal'
cor_path = data_path/'coronal'
ax_path = data_path/'axial'

weights = torch.load('loss_weights.pt')

class MRNet(nn.Module):
    def __init__(self, pretrained=True):
        super().__init__()
        self.model = models.alexnet(pretrained=pretrained)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(256, 1)

    def forward(self, x):
        # in the original code, the input was squeezed here, but this won't work with fastai
        x = self.model.features(x)
        x = self.gap(x).view(x.size(0), -1)
        x = torch.max(x, 0, keepdim=True)[0]
        return torch.sigmoid(self.classifier(x))

    def __call__(self, x): return self.forward(x)


class WtBCELoss(nn.Module):
    def __init__(self, wts):
        super().__init__()
        self.wts = wts.float()

    def forward(self, output, target):
        loss = self.wts[0]*(target.float() * torch.log(output).float()) + self.wts[1]*((1-target).float() * torch.log(1-output).float())
        return torch.neg(torch.mean(loss))


class MR3DImDataBunch(ImageDataBunch):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def one_batch(self, ds_type:DatasetType=DatasetType.Train, detach:bool=True, denorm:bool=True, cpu:bool=True)->Collection[Tensor]:
        "Get one batch from the data loader of `ds_type`. Optionally `detach` and `denorm`."
        dl = self.dl(ds_type)
        w = self.num_workers
        self.num_workers = 0
        try:     x,y = next(iter(dl))
        finally: self.num_workers = w
        if detach: x,y = to_detach(x,cpu=cpu),to_detach(y,cpu=cpu)
        norm = getattr(self,'norm',False)
        if denorm and norm:
            x = self.denorm(x)
            if norm.keywords.get('do_y',False): y = self.denorm(y, do_x=True)
        x = torch.squeeze(x, dim=0) # squeeze needed here for learn.summary()
        return x,y


class MR3DImageList(ImageList):
    _bunch = MR3DImDataBunch # necessary for Data Block API functionality
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.max_slc = 51 # optimized for sagittal image stacks only...TODO rewrite for any max stack size
        self.c = 1

    # pads on both sides of image stack with zero arrays to equal max_slc
    def open(self, fn):
        x = np.load(fn)
        if x.shape[0] < self.max_slc:
            x_pad = np.zeros((self.max_slc, 256, 256))
            mid = x_pad.shape[0] // 2
            up = x.shape[0] // 2
            if x.shape[0] % 2 == 1: x_pad[mid-up:mid+up+1] = x
            else: x_pad[mid-up:mid+up] = x
        else:
            x_pad = x
        return self.arr2image(np.stack([x_pad]*3, axis=1))

    # converts np.ndarray to fastai Image class
    @staticmethod
    def arr2image(arr:np.ndarray, div:bool=True, cls:type=Image):
        x = Tensor(arr)
        if div == True: x.div_(255)
        return cls(x)


# squeeze input prior to loss calculation
class MRNetCallback(Callback):
    def on_batch_begin(self, last_input, **kwargs):
        x = torch.squeeze(last_input, dim=0)
        return dict(last_input=x)


class MRNetLearner(Learner):
    # redefine specifically for MRNet layer groups
    def freeze_to(self, n:int)->None:
        "Freeze layers up to layer group `n`."
        for g in self.layer_groups[:n]:
            if is_listy(g):
                for l in g:
                    if not self.train_bn or not isinstance(l, bn_types): requires_grad(l, False)
            else: requires_grad(g, False)
        for g in self.layer_groups[n:]: requires_grad(g, True)
        self.create_opt(defaults.lr)

    def freeze(self)->None:
        "Freeze up to the last layer group."
        self.freeze_to(-1)
        self.create_opt(defaults.lr)

    def unfreeze(self):
        "Unfreeze entire model."
        self.freeze_to(0)
        self.create_opt(defaults.lr)


def mrnet_learner(data:DataBunch, model:Callable=MRNet(), pretrained:bool=True, init=nn.init.kaiming_normal_, **kwargs:Any)->Learner:
    _layer_groups = [model.model.features, model.model.avgpool, model.model.classifier, model.gap, model.classifier]
    learn = MRNetLearner(data, model, layer_groups=_layer_groups, **kwargs)
    if pretrained: learn.freeze()
    if init: apply_init(model.classifier, init)
    return learn