"""Contrastive learning experiment training script.

Copyright (c) 2022 University of Michigan. All rights reserved.

Licensed under the MIT License. See LICENSE for license information.

"""

import yaml

import logging

from functools import partial

from typing import Dict, Any

import torch

import pytorch_lightning as pl

import torchmetrics

from opensrh.models import MLP, resnet_backbone, ContrastiveLearningNetwork, vit_backbone

from opensrh.train.common import (setup_output_dirs, parse_args, get_exp_name,

                                  get_contrastive_dataloaders, config_loggers,

                                  get_optimizer_func, get_scheduler_func)

from opensrh.losses.supcon import SupConLoss

class ContrastiveSystem(pl.LightningModule):

    """Lightning system for contrastive learning experiments."""

    def __init__(self, cf: Dict[str, Any], num_it_per_ep: int):

        super().__init__()

        self.cf_ = cf

        if cf["model"]["backbone"] == "resnet50":

            bb = partial(resnet_backbone, arch=cf["model"]["backbone"])

        elif cf["model"]["backbone"] == "vit":

            bb = partial(vit_backbone, cf["model"]["backbone_params"])

        else:

            raise NotImplementedError()

        mlp = partial(MLP,

                      n_in=bb().num_out,

                      hidden_layers=cf["model"]["mlp_hidden"],

                      n_out=cf["model"]["num_embedding_out"])

        self.model = ContrastiveLearningNetwork(bb, mlp)

        self.criterion = SupConLoss()

        self.train_loss = torchmetrics.MeanMetric()

        self.val_loss = torchmetrics.MeanMetric()

        self.num_it_per_ep_ = num_it_per_ep

    def predict_step(self, batch, batch_idx):

        out = self.model.bb(batch["image"])

        return {

            "path": batch["path"],

            "label": batch["label"],

            "embeddings": out

        }

    def on_train_epoch_end(self):

        train_loss = self.train_loss.compute()

        self.log("train/contrastive_manualepoch",

                 train_loss,

                 on_epoch=True,

                 sync_dist=False,

                 rank_zero_only=True)

        logging.info(f"train/contrastive_manualepoch {train_loss}")

        self.train_loss.reset()

    def on_validation_epoch_end(self):

        val_loss = self.val_loss.compute()

        self.log("val/contrastive_manualepoch",

                 val_loss,

                 on_epoch=True,

                 sync_dist=False,

                 rank_zero_only=True)

        logging.info(f"val/contrastive_manualepoch {val_loss}")

        self.val_loss.reset()

    def configure_optimizers(self):

        # if not training, no optimizer

        if "training" not in self.cf_:

            return None

        # get optimizer

        opt = get_optimizer_func(self.cf_)(self.model.parameters())

        # check if use a learn rate scheduler

        sched_func = get_scheduler_func(self.cf_, self.num_it_per_ep_)

        if not sched_func:

            return opt

        # get learn rate scheduler

        lr_scheduler_config = {

            "scheduler": sched_func(opt),

            "interval": "step",

            "frequency": 1,

            "name": "lr"

        }

        return [opt], lr_scheduler_config

    def configure_ddp(self, *args, **kwargs):

        logging.basicConfig(level=logging.INFO)

        return super().configure_ddp(*args, **kwargs)

class SimCLRSystem(ContrastiveSystem):

    """Lightning system for SimCLR experiment"""

    def __init__(self, cf, num_it_per_ep):

        super().__init__(cf, num_it_per_ep)

    def forward(self, data):

        return torch.cat([self.model(x) for x in data["image"]], dim=1)

    def training_step(self, batch, batch_idx):

        pred = torch.cat([self.model(x) for x in batch["image"]], dim=1)

        pred_gather = self.all_gather(pred, sync_grads=True)

        pred_gather = pred_gather.reshape(-1, *pred_gather.shape[-2:])

        loss = self.criterion(pred_gather)

        bs = batch["image"][0].shape[0]

        self.log("train/contrastive",

                 loss,

                 on_step=True,

                 on_epoch=True,

                 batch_size=bs)

        self.train_loss.update(loss, weight=bs)

        return loss

    def validation_step(self, batch, batch_idx):

        bs = batch["image"][0].shape[0]

        pred = torch.cat([self.model(x) for x in batch["image"]], dim=1)

        pred_gather = self.all_gather(pred, sync_grads=True)

        pred_gather = pred_gather.reshape(-1, *pred_gather.shape[-2:])

        loss = self.criterion(pred_gather)

        self.val_loss.update(loss, weight=bs)

class SupConSystem(ContrastiveSystem):

    """Lightning system for SupCon experiment"""

    def __init__(self, cf, num_it_per_ep):

        super().__init__(cf, num_it_per_ep)

    def forward(self, data):

        return torch.cat([self.model(x) for x in data["image"]], dim=1)

    def training_step(self, batch, batch_idx):

        pred = torch.cat([self.model(x) for x in batch["image"]], dim=1)

        pred_gather = self.all_gather(pred, sync_grads=True)

        pred_gather = pred_gather.reshape(-1, *pred_gather.shape[-2:])

        label_gather = self.all_gather(batch["label"]).reshape(-1, 1)

        loss = self.criterion(pred_gather, label_gather)

        bs = batch["image"][0].shape[0]

        self.log("train/contrastive",

                 loss,

                 on_step=True,

                 on_epoch=True,

                 batch_size=bs)

        self.train_loss.update(loss, weight=bs)

        return loss

    def validation_step(self, batch, batch_idx):

        bs = batch["image"][0].shape[0]

        pred = torch.cat([self.model(x) for x in batch["image"]], dim=1)

        pred_gather = self.all_gather(pred, sync_grads=True)

        pred_gather = pred_gather.reshape(-1, *pred_gather.shape[-2:])

        label_gather = self.all_gather(batch["label"]).reshape(-1, 1)

        loss = self.criterion(pred_gather, label_gather)

        self.val_loss.update(loss, weight=bs)

def main():

    cf_fd = parse_args()

    cf = yaml.load(cf_fd, Loader=yaml.FullLoader)

    exp_root, model_dir, cp_config = setup_output_dirs(cf, get_exp_name, "")

    pl.seed_everything(cf["infra"]["seed"])

    # logging and copying config files

    cp_config(cf_fd.name)

    config_loggers(exp_root)

    # get dataloaders

    train_loader, valid_loader = get_contrastive_dataloaders(cf)

    logging.info(f"num devices: {torch.cuda.device_count()}")

    logging.info(f"num workers in dataloader: {train_loader.num_workers}")

    num_it_per_ep = len(train_loader)

    if torch.cuda.device_count() > 1:

        num_it_per_ep //= torch.cuda.device_count()

    if cf["training"]["objective"] == "supcon":

        system_func = SupConSystem

    elif cf["training"]["objective"] == "simclr":

        system_func = SimCLRSystem

    else:

        raise NotImplementedError()

    ce_exp = system_func(cf, num_it_per_ep)

    # config loggers

    logger = [

        pl.loggers.TensorBoardLogger(save_dir=exp_root, name="tb"),

        pl.loggers.CSVLogger(save_dir=exp_root, name="csv")

    ]

    # config callbacks

    epoch_ckpt = pl.callbacks.ModelCheckpoint(

        dirpath=model_dir,

        save_top_k=-1,

        save_on_train_epoch_end=True,

        filename="ckpt-epoch{epoch}-loss{val/contrastive_manualepoch:.2f}",

        auto_insert_metric_name=False)

    lr_monitor = pl.callbacks.LearningRateMonitor(logging_interval="step",

                                                  log_momentum=False)

    # create trainer

    trainer = pl.Trainer(accelerator="gpu",

                         devices=-1,

                         default_root_dir=exp_root,

                         strategy=pl.strategies.DDPStrategy(

                             find_unused_parameters=False, static_graph=True),

                         logger=logger,

                         log_every_n_steps=10,

                         callbacks=[epoch_ckpt, lr_monitor],

                         max_epochs=cf["training"]["num_epochs"])

    trainer.fit(ce_exp,

                train_dataloaders=train_loader,

                val_dataloaders=valid_loader)

if __name__ == '__main__':

    main()