import pytorch_lightning as pl

# from pl_bolts.models.self_supervised import SimCLR

from pl_bolts.optimizers.lars_scheduling import LARSWrapper

from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR

from torch.optim import Adam

import torch

import re

import pdb 

import math

from argparse import ArgumentParser

from typing import Callable, Optional

import numpy as np

import torch

import torch.distributed as dist

import torch.nn.functional as F

from pytorch_lightning.utilities import AMPType

from torch import nn

from torch.optim.optimizer import Optimizer

from models.resnet_simclr import ResNetSimCLR

import re

import time

import pickle

import yaml

import logging

import os

from clinical_ts.simclr_dataset_wrapper import SimCLRDataSetWrapper

from clinical_ts.create_logger import create_logger

import pickle

from pytorch_lightning import Trainer, seed_everything

from torch import nn

from torch.nn import functional as F

from online_evaluator import SSLOnlineEvaluator

from ecg_datamodule import ECGDataModule

from pytorch_lightning.loggers import TensorBoardLogger

from pl_bolts.models.self_supervised.evaluator import Flatten

import pdb

method="simclr"

logger = create_logger(__name__)

def _accuracy(zis, zjs, batch_size):

    with torch.no_grad():

        representations = torch.cat([zjs, zis], dim=0)

        similarity_matrix = torch.mm(

            representations, representations.t().contiguous())

        corrected_similarity_matrix = similarity_matrix - \

            torch.eye(2*batch_size).type_as(similarity_matrix)

        pred_similarities, pred_indices = torch.max(

            corrected_similarity_matrix[:batch_size], dim=1)

        correct_indices = torch.arange(batch_size)+batch_size

        correct_preds = (

            pred_indices == correct_indices.type_as(pred_indices)).sum()

    return correct_preds.float()/batch_size

def mean(res, key1, key2=None):

    if key2 is not None:

        return torch.stack([x[key1][key2] for x in res]).mean()

    return torch.stack([x[key1] for x in res if type(x) == dict and key1 in x.keys()]).mean()

class Projection(nn.Module):

    def __init__(self, input_dim=2048, hidden_dim=2048, output_dim=128):

        super().__init__()

        self.output_dim = output_dim

        self.input_dim = input_dim

        self.hidden_dim = hidden_dim

        self.model = nn.Sequential(

            # nn.AdaptiveAvgPool2d((1, 1)),

            Flatten(),

            nn.Linear(self.input_dim, self.hidden_dim, bias=True),

            # nn.BatchNorm1d(self.hidden_dim),

            nn.ReLU(),

            nn.Linear(self.hidden_dim, self.output_dim, bias=True))

    def forward(self, x):

        x = self.model(x)

        return F.normalize(x, dim=1)

class SyncFunction(torch.autograd.Function):

    @staticmethod

    def forward(ctx, tensor):

        ctx.batch_size = tensor.shape[0]

        gathered_tensor = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())]

        torch.distributed.all_gather(gathered_tensor, tensor)

        gathered_tensor = torch.cat(gathered_tensor, 0)

        return gathered_tensor

    @staticmethod

    def backward(ctx, grad_output):

        grad_input = grad_output.clone()

        torch.distributed.all_reduce(grad_input, op=torch.distributed.ReduceOp.SUM, async_op=False)

        return grad_input[torch.distributed.get_rank() * ctx.batch_size:(torch.distributed.get_rank() + 1) *

                          ctx.batch_size]

class CustomSimCLR(pl.LightningModule):

    def __init__(self,

                 batch_size,

                 num_samples,

                 warmup_epochs=10,

                 lr=1e-4,

                 opt_weight_decay=1e-6,

                 loss_temperature=0.5,

                 config=None,

                 transformations=None,

                 **kwargs):

        """

        Args:

            batch_size: the batch size

            num_samples: num samples in the dataset

            warmup_epochs: epochs to warmup the lr for

            lr: the optimizer learning rate

            opt_weight_decay: the optimizer weight decay

            loss_temperature: the loss temperature

        """

        super(CustomSimCLR, self).__init__()

        self.config = config

        self.transformations = transformations

        self.epoch = 0

        self.batch_size = batch_size

        self.num_samples = num_samples

        self.save_hyperparameters()

        # pdb.set_trace()

    def configure_optimizers(self):

        global_batch_size = self.trainer.world_size * self.hparams.batch_size

        self.train_iters_per_epoch = self.hparams.num_samples // global_batch_size

        # TRICK 1 (Use lars + filter weights)

        # exclude certain parameters

        parameters = self.exclude_from_wt_decay(

            self.named_parameters(),

            weight_decay=self.hparams.opt_weight_decay

        )

        # optimizer = LARSWrapper(Adam(parameters, lr=self.hparams.lr))

        optimizer = Adam(parameters, lr=self.hparams.lr)

        # Trick 2 (after each step)

        self.hparams.warmup_epochs = self.hparams.warmup_epochs * self.train_iters_per_epoch

        max_epochs = self.trainer.max_epochs * self.train_iters_per_epoch

        linear_warmup_cosine_decay = LinearWarmupCosineAnnealingLR(

            optimizer,

            warmup_epochs=self.hparams.warmup_epochs,

            max_epochs=max_epochs,

            warmup_start_lr=0,

            eta_min=0

        )

        scheduler = {

            'scheduler': linear_warmup_cosine_decay,

            'interval': 'step',

            'frequency': 1

        }

        return [optimizer], [scheduler]

    def exclude_from_wt_decay(self, named_params, weight_decay, skip_list=['bias', 'bn']):

        params = []

        excluded_params = []

        for name, param in named_params:

            if not param.requires_grad:

                continue

            elif any(layer_name in name for layer_name in skip_list):

                excluded_params.append(param)

            else:

                params.append(param)

        return [

            {'params': params, 'weight_decay': weight_decay},

            {'params': excluded_params, 'weight_decay': 0.}

        ]

    def shared_forward(self, batch, batch_idx):

        (x1, y1), (x2, y2) = batch

        # ENCODE

        # encode -> representations

        # (b, 3, 32, 32) -> (b, 2048, 2, 2)

        x1 = self.to_device(x1)

        x2 = self.to_device(x2)

        h1 = self.encoder(x1)[0]

        h2 = self.encoder(x2)[0]

        # the bolts resnets return a list of feature maps

        if isinstance(h1, list):

            h1 = h1[-1]

            h2 = h2[-1]

        # PROJECT

        # img -> E -> h -> || -> z

        # (b, 2048, 2, 2) -> (b, 128)

        z1 = self.projection(h1.squeeze())

        z2 = self.projection(h2.squeeze())

        return z1, z2

    def nt_xent_loss(self, out_1, out_2, temperature, eps=1e-6):

        """

            assume out_1 and out_2 are normalized

            out_1: [batch_size, dim]

            out_2: [batch_size, dim]

        """

        # gather representations in case of distributed training

        # out_1_dist: [batch_size * world_size, dim]

        # out_2_dist: [batch_size * world_size, dim]

        if torch.distributed.is_available() and torch.distributed.is_initialized():

            out_1_dist = SyncFunction.apply(out_1)

            out_2_dist = SyncFunction.apply(out_2)

            print("out dist shape: ", out_1_dist.shape)

        else:

            out_1_dist = out_1

            out_2_dist = out_2

        # out: [2 * batch_size, dim]

        # out_dist: [2 * batch_size * world_size, dim]

        out = torch.cat([out_1, out_2], dim=0)

        out_dist = torch.cat([out_1_dist, out_2_dist], dim=0)

        # cov and sim: [2 * batch_size, 2 * batch_size * world_size]

        # neg: [2 * batch_size]

        cov = torch.mm(out, out_dist.t().contiguous())

        sim = torch.exp(cov / temperature)

        neg = sim.sum(dim=-1)

        # from each row, subtract e^1 to remove similarity measure for x1.x1

        row_sub = torch.Tensor(neg.shape).fill_(math.e).to(neg.device)

        neg = torch.clamp(neg - row_sub, min=eps)  # clamp for numerical stability

        # Positive similarity, pos becomes [2 * batch_size]

        pos = torch.exp(torch.sum(out_1 * out_2, dim=-1) / temperature)

        pos = torch.cat([pos, pos], dim=0)

        loss = -torch.log(pos / (neg + eps)).mean()

        return loss

    def training_step(self, batch, batch_idx):

        z1, z2 = self.shared_forward(batch, batch_idx)

        loss = self.nt_xent_loss(z1, z2, self.hparams.loss_temperature)

        # result = pl.TrainResult(minimize=loss)

        # result.log('train/train_loss', loss, on_epoch=True)

        acc = _accuracy(z1, z2, z1.shape[0])

        # result.log('train/train_acc', acc, on_epoch=True)

        result = {

            "train/train_loss": loss, 

            "minimize":loss,

            "train/train_acc" : acc,

        }

        return loss

    def validation_step(self, batch, batch_idx, dataloader_idx):

        if dataloader_idx != 0:

            return {}

        z1, z2 = self.shared_forward(batch, batch_idx)

        loss = self.nt_xent_loss(z1, z2, self.hparams.loss_temperature)

        acc = _accuracy(z1, z2, z1.shape[0])

        results = {

            'val_loss': loss,

            'val_acc': torch.tensor(acc)

        }

        return results

    def validation_epoch_end(self, outputs):

        # outputs[0] because we are using multiple datasets!

        val_loss = mean(outputs[0], 'val_loss')

        val_acc = mean(outputs[0], 'val_acc')

        log = {

            'val/val_loss': val_loss,

            'val/val_acc': val_acc

        }

        return {'val_loss': val_loss, 'log': log, 'progress_bar': log}

    def on_train_start(self):

        # log configuration

        config_str = re.sub(r"[,\}\{]", "<br/>", str(self.config))

        config_str = re.sub(r"[\[\]\']", "", config_str)

        transformation_str = re.sub(r"[\}]", "<br/>", str(["<br>" + str(

            t) + ":<br/>" + str(t.get_params()) for t in self.transformations]))

        transformation_str = re.sub(r"[,\"\{\'\[\]]", "", transformation_str)

        self.logger.experiment.add_text(

            "configuration", str(config_str), global_step=0)

        self.logger.experiment.add_text("transformations", str(

            transformation_str), global_step=0)

        self.epoch = 0

    def on_epoch_end(self):

        self.epoch += 1

    def type(self):

        return self.encoder.features[0][0].weight.type()

    def get_representations(self, x):

        return self.encoder(x)[0]

    def get_model(self):

        return self.encoder

    def get_device(self):

        return self.encoder.features[0][0].weight.device

    def to_device(self, x):

        return x.type(self.type()).to(self.get_device())

def parse_args(parent_parser):

    parser = ArgumentParser(parents=[parent_parser], add_help=False)

    parser.add_argument('-t', '--trafos', nargs='+', help='add transformation to data augmentation pipeline',

                        default=["GaussianNoise", "ChannelResize", "RandomResizedCrop"])

    # GaussianNoise

    parser.add_argument(

            '--gaussian_scale', help='std param for gaussian noise transformation', default=0.005, type=float)

    # RandomResizedCrop

    parser.add_argument('--rr_crop_ratio_range',

                            help='ratio range for random resized crop transformation', default=[0.5, 1.0], type=float)

    parser.add_argument(

            '--output_size', help='output size for random resized crop transformation', default=250, type=int)

    # DynamicTimeWarp

    parser.add_argument(

            '--warps', help='number of warps for dynamic time warp transformation', default=3, type=int)

    parser.add_argument(

            '--radius', help='radius of warps of dynamic time warp transformation', default=10, type=int)

    # TimeWarp

    parser.add_argument(

            '--epsilon', help='epsilon param for time warp', default=10, type=float)

    # ChannelResize

    parser.add_argument('--magnitude_range', nargs='+',

                            help='range for scale param for ChannelResize transformation', default=[0.5, 2], type=float)

    # Downsample

    parser.add_argument(

            '--downsample_ratio', help='downsample ratio for Downsample transformation', default=0.2, type=float)

    # TimeOut

    parser.add_argument('--to_crop_ratio_range', nargs='+',

                            help='ratio range for timeout transformation', default=[0.2, 0.4], type=float)

    # resume training

    parser.add_argument('--resume', action='store_true')

    parser.add_argument(

            '--gpus', help='number of gpus to use; use cpu if gpu=0', type=int, default=1)

    parser.add_argument(

            '--num_nodes', default=1, help='number of cluster nodes', type=int)

    parser.add_argument(

            '--distributed_backend', help='sets backend type')

    parser.add_argument('--batch_size', type=int)

    parser.add_argument('--epochs', type=int)

    parser.add_argument('--debug', action='store_true')

    parser.add_argument('--warm_up', default=1, type=int, help="number of warm up epochs")

    parser.add_argument('--precision', type=int)

    parser.add_argument('--datasets', dest="target_folders",

                            nargs='+', help='used datasets for pretraining')

    parser.add_argument('--log_dir', default="./experiment_logs")

    parser.add_argument(

            '--percentage', help='determines how much of the dataset shall be used during the pretraining', type=float, default=1.0)

    parser.add_argument('--lr', type=float, help="learning rate")

    parser.add_argument('--out_dim', type=int, help="output dimension of model")

    parser.add_argument('--filter_cinc', default=False, action="store_true", help="only valid if cinc is selected: filter out the ptb data")

    parser.add_argument('--base_model')

    parser.add_argument('--widen',type=int, help="use wide xresnet1d50")

    parser.add_argument('--run_callbacks', default=False, action="store_true", help="run callbacks which asses linear evaluaton and finetuning metrics during pretraining")

    parser.add_argument('--checkpoint_path', default="")

    return parser

def init_logger(config):

    level = logging.INFO

    if config['debug']:

        level = logging.DEBUG

    # remove all handlers to change basic configuration

    for handler in logging.root.handlers[:]:

        logging.root.removeHandler(handler)

    if not os.path.isdir(config['log_dir']):

        os.mkdir(config['log_dir'])

    logging.basicConfig(filename=os.path.join(config['log_dir'], 'info.log'), level=level,

                        format='%(asctime)s %(name)s:%(lineno)s %(levelname)s:  %(message)s  ')

    return logging.getLogger(__name__)

def pretrain_routine(args):

    t_params = {"gaussian_scale": args.gaussian_scale, "rr_crop_ratio_range": args.rr_crop_ratio_range, "output_size": args.output_size, "warps": args.warps, "radius": args.radius,

                "epsilon": args.epsilon, "magnitude_range": args.magnitude_range, "downsample_ratio": args.downsample_ratio, "to_crop_ratio_range": args.to_crop_ratio_range,

                "bw_cmax":0.1, "em_cmax":0.5, "pl_cmax":0.2, "bs_cmax":1}

    transformations = args.trafos

    checkpoint_config = os.path.join("checkpoints", "bolts_config.yaml")

    config_file = checkpoint_config if args.resume and os.path.isfile(

        checkpoint_config) else "bolts_config.yaml"

    config = yaml.load(open(config_file, "r"), Loader=yaml.FullLoader)

    args_dict = vars(args)

    for key in set(config.keys()).union(set(args_dict.keys())):

        config[key] = config[key] if (key not in args_dict.keys() or key in args_dict.keys(

        ) and key in config.keys() and args_dict[key] is None) else args_dict[key]

    if args.target_folders is not None:

        config["dataset"]["target_folders"] = args.target_folders

    config["dataset"]["percentage"] = args.percentage if args.percentage is not None else config["dataset"]["percentage"]

    config["dataset"]["filter_cinc"] = args.filter_cinc if args.filter_cinc is not None else config["dataset"]["filter_cinc"]

    config["model"]["base_model"] = args.base_model if args.base_model is not None else config["model"]["base_model"]

    config["model"]["widen"] = args.widen if args.widen is not None else config["model"]["widen"]

    if args.out_dim is not None:

        config["model"]["out_dim"] = args.out_dim

    init_logger(config)

    dataset = SimCLRDataSetWrapper(

        config['batch_size'], **config['dataset'], transformations=transformations, t_params=t_params)

    for i, t in enumerate(dataset.transformations):

        logger.info(str(i) + ". Transformation: " +

                    str(t) + ": " + str(t.get_params()))

    date = time.asctime()

    label_to_num_classes = {"label_all": 71, "label_diag": 44, "label_form": 19,

                            "label_rhythm": 12, "label_diag_subclass": 23, "label_diag_superclass": 5}

    ptb_num_classes = label_to_num_classes[config["eval_dataset"]

                                           ["ptb_xl_label"]]

    abr = {"Transpose": "Tr", "TimeOut": "TO", "DynamicTimeWarp": "DTW", "RandomResizedCrop": "RRC", "ChannelResize": "ChR", "GaussianNoise": "GN",

           "TimeWarp": "TW", "ToTensor": "TT", "GaussianBlur": "GB", "BaselineWander": "BlW", "PowerlineNoise": "PlN", "EMNoise": "EM", "BaselineShift": "BlS"}

    trs = re.sub(r"[,'\]\[]", "", str([abr[str(tr)] if abr[str(tr)] not in [

                 "TT", "Tr"] else '' for tr in dataset.transformations]))

    name = str(date) + "_" + method + "_" + str(

        time.time_ns())[-3:] + "_" + trs[1:]

    tb_logger = TensorBoardLogger(args.log_dir, name=name, version='')

    config["log_dir"] = os.path.join(args.log_dir, name)

    print(config)

    return config, dataset, date, transformations, t_params, ptb_num_classes, tb_logger

def aftertrain_routine(config, args, trainer, pl_model, datamodule, callbacks):

    scores = {}

    for ca in callbacks:

        if isinstance(ca, SSLOnlineEvaluator):

            scores[str(ca)] = {"macro": ca.best_macro}

    results = {"config": config, "trafos": args.trafos, "scores": scores}

    with open(os.path.join(config["log_dir"], "results.pkl"), 'wb') as handle:

        pickle.dump(results, handle)

    trainer.save_checkpoint(os.path.join(config["log_dir"], "checkpoints", "model.ckpt"))

    with open(os.path.join(config["log_dir"], "config.txt"), "w") as text_file:

        print(config, file=text_file)

def cli_main():

    from pytorch_lightning import Trainer

    from online_evaluator import SSLOnlineEvaluator

    from ecg_datamodule import ECGDataModule

    from clinical_ts.create_logger import create_logger

    from os.path import exists

    parser = ArgumentParser()

    parser = parse_args(parser)

    logger.info("parse arguments")

    args = parser.parse_args()

    config, dataset, date, transformations, t_params, ptb_num_classes, tb_logger = pretrain_routine(args)

    # data

    ecg_datamodule = ECGDataModule(config, transformations, t_params)

    callbacks = []

    if args.run_callbacks:

            # callback for online linear evaluation/fine-tuning

        linear_evaluator = SSLOnlineEvaluator(drop_p=0,

                                          z_dim=512, num_classes=ptb_num_classes, hidden_dim=None, lin_eval_epochs=config["eval_epochs"], eval_every=config["eval_every"], mode="linear_evaluation", verbose=False)

        fine_tuner = SSLOnlineEvaluator(drop_p=0,

                                          z_dim=512, num_classes=ptb_num_classes, hidden_dim=None, lin_eval_epochs=config["eval_epochs"], eval_every=config["eval_every"], mode="fine_tuning", verbose=False)

        callbacks.append(linear_evaluator)

        callbacks.append(fine_tuner)

    # configure trainer

    trainer = Trainer(logger=tb_logger, max_epochs=config["epochs"], gpus=args.gpus,

                      distributed_backend=args.distributed_backend, auto_lr_find=False, num_nodes=args.num_nodes, precision=config["precision"], callbacks=callbacks)

    # pytorch lightning module

    model = ResNetSimCLR(**config["model"])

    pl_model = CustomSimCLR(

            config["batch_size"], ecg_datamodule.num_samples, warmup_epochs=config["warm_up"], lr=config["lr"],

            out_dim=config["model"]["out_dim"], config=config,

            transformations=ecg_datamodule.transformations, loss_temperature=config["loss"]["temperature"], weight_decay=eval(config["weight_decay"]))

    pl_model.encoder = model

    pl_model.projection = Projection(

            input_dim=model.l1.in_features, hidden_dim=512, output_dim=config["model"]["out_dim"])

    # load checkpoint

    if args.checkpoint_path != "":

        if exists(args.checkpoint_path):

            logger.info("Retrieve checkpoint from " + args.checkpoint_path)

            pl_model.load_from_checkpoint(args.checkpoint_path)

        else:

            raise("checkpoint does not exist")

    # start training

    trainer.fit(pl_model, ecg_datamodule)

    aftertrain_routine(config, args, trainer, pl_model, ecg_datamodule, callbacks)

if __name__ == "__main__":  

    cli_main()