'''
This should handle transforming targets, currently only implemented for lung variation
'''
import pytorch_lightning as pl
from torch.optim.lr_scheduler import StepLR
from medseg.utils import get_optimizer, DICELoss, DICEMetric, itk_snap_spawner
from medseg.unet_v2 import UNet


class PolySeg3DModule(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.save_hyperparameters(hparams)

        dropout = getattr(self.hparams, "dropout", None)
        self.weight_decay = getattr(self.hparams, "weight_decay", None)
        self.scheduling_factor = getattr(self.hparams, "scheduling_factor", None)
        if dropout == True:
            print("WARNING: Replacing old hparams true dropout by full dropout")
            dropout = "full"
        
        self.model = UNet(self.hparams.nin, self.hparams.seg_nout, "instance", "3d", self.hparams.init_channel)

        self.lossfn = DICELoss(volumetric=False, per_channel=True, check_bounds=False)
        self.dicer = DICEMetric(per_channel_metric=True, check_bounds=False)

    def forward(self, x, get_bg=False):
        '''
        Testing with 2 channel output for both lungs
        '''
        logits = self.model(x)  # 3 canais, bg, ll, rl
        if get_bg:
            y_hat = logits.softmax(dim=1)
        else:
            y_hat = logits.softmax(dim=1)[:, 1:]
        y_hat_lung = y_hat.sum(dim=1, keepdim=True)

        return y_hat, y_hat_lung
        
    def polymorphic_loss_metrics(self, x, y, y_hat, y_hat_lung, meta, mode=None):
        target_format = meta["target_format"][0]
        assert target_format in ["simple", "has_left_right", "has_ggo_con", "full_poly"]
        loss, metrics = None, None

        if target_format in ["simple", "has_ggo_con"]:
            # Format returns only lung binary mask, extract lung target
            y = y[:, 0:1]
            if mode == "val":
                metrics = self.dicer(y_hat_lung, y)[0]
            elif mode == "train":
                loss = self.lossfn(y_hat_lung, y)
        elif target_format in ["has_left_right", "full_poly"]:
            # Extract left right lung
            y = y[:, 0:2]
            if mode == "val":
                metrics = self.dicer(y_hat, y)
            elif mode == "train":
                loss = self.lossfn(y_hat, y)

        return y, loss, metrics

    def compute_loss(self, x, y, y_hat, y_hat_lung, meta, prestr):
        _, loss, _ = self.polymorphic_loss_metrics(x, y, y_hat, y_hat_lung, meta, mode="train")
        
        self.log(f"{prestr}loss", loss, on_step=True, on_epoch=True)

        return loss

    def compute_metrics(self, x, y, y_hat, y_hat_lung, meta):
        _, _, metrics = self.polymorphic_loss_metrics(x, y, y_hat, y_hat_lung, meta, mode="val")

        if isinstance(metrics, list):
            left_lung_dice, right_lung_dice = metrics
            self.log("left_lung_dice", left_lung_dice, on_epoch=True, on_step=True, prog_bar=True)
            self.log("right_lung_dice", right_lung_dice, on_epoch=True, on_step=True, prog_bar=True)
        else:
            self.log("lung_dice", metrics, on_epoch=True, on_step=True, prog_bar=True)

        
    def training_step(self, train_batch, batch_idx):
        x, y, meta = train_batch

        y_hat, y_hat_lung = self.forward(x)

        loss = self.compute_loss(x, y, y_hat, y_hat_lung, meta, prestr='')

        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y, meta = val_batch
        
        y_hat, y_hat_lung = self.forward(x)
        
        self.compute_loss(x, y, y_hat, y_hat_lung, meta, prestr="val_")
        self.compute_metrics(x, y, y_hat, y_hat_lung, meta)

    def configure_optimizers(self):
        '''
        Select optimizer and scheduling strategy according to hparams.
        '''
        opt = getattr(self.hparams, "opt", "Adam")
        optimizer = get_optimizer(opt, self.model.parameters(), self.hparams.lr, wd=self.weight_decay)
        print(f"Weight decay: {self.weight_decay}")

        if self.scheduling_factor is not None:
            print(f"Using step LR {self.scheduling_factor}!")
            scheduler = StepLR(optimizer, 1, self.scheduling_factor)
            return [optimizer], [scheduler]
        else:
            return optimizer