###############
#generic
import torch
from torch import nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, ConcatDataset
from torchvision import transforms
import torch.nn.functional as F

import torchvision
import os
import argparse
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
import copy

#################
#specific
from clinical_ts.timeseries_utils import *
from clinical_ts.ecg_utils import *

from functools import partial
from pathlib import Path
import pandas as pd
import numpy as np

from clinical_ts.xresnet1d import xresnet1d50,xresnet1d101
from clinical_ts.basic_conv1d import weight_init
from clinical_ts.eval_utils_cafa import eval_scores, eval_scores_bootstrap
from clinical_ts.cpc import *

def _freeze_bn_stats(model, freeze=True):
    for m in model.modules():
        if(isinstance(m,nn.BatchNorm1d)):
            if(freeze):
                m.eval()
            else:
                m.train()
                
def sanity_check(model, state_dict_pre):
    """
    Linear classifier should not change any weights other than the linear layer.
    This sanity check asserts nothing wrong happens (e.g., BN stats updated).
    """
    print("=> loading state dict for sanity check")
    state_dict = model.state_dict()

    for k in list(state_dict.keys()):
        print(k)
        # only ignore fc layer
        if 'head.1.weight' in k or 'head.1.bias' in k:
            continue


        assert ((state_dict[k].cpu() == state_dict_pre[k].cpu()).all()), \
            '{} is changed in linear classifier training.'.format(k)

    print("=> sanity check passed.")

class LightningCPC(pl.LightningModule):

    def __init__(self, hparams):
        super(LightningCPC, self).__init__()
        
        self.hparams = hparams
        self.lr = self.hparams.lr
        
        #these coincide with the adapted wav2vec2 params
        if(self.hparams.fc_encoder):
            strides=[1]*4 
            kss = [1]*4 
            features = [512]*4
        else: #strided conv encoder
            strides=[2,2,2,2] #original wav2vec2 [5,2,2,2,2,2] original cpc [5,4,2,2,2]
            kss = [10,4,4,4] #original wav2vec2 [10,3,3,3,3,2] original cpc [18,8,4,4,4]
            features = [512]*4 #wav2vec2 [512]*6 original cpc [512]*5
        
        if(self.hparams.finetune):
            self.criterion = F.cross_entropy if self.hparams.finetune_dataset == "thew" else F.binary_cross_entropy_with_logits
            if(self.hparams.finetune_dataset == "thew"):
                num_classes = 5
            elif(self.hparams.finetune_dataset == "ptbxl_super"):
                num_classes = 5
            if(self.hparams.finetune_dataset == "ptbxl_all"):
                num_classes = 71
        else:
            num_classes = None

        self.model_cpc = CPCModel(input_channels=self.hparams.input_channels, strides=strides,kss=kss,features=features,n_hidden=self.hparams.n_hidden,n_layers=self.hparams.n_layers,mlp=self.hparams.mlp,lstm=not(self.hparams.gru),bias_proj=self.hparams.bias,num_classes=num_classes,skip_encoder=self.hparams.skip_encoder,bn_encoder=not(self.hparams.no_bn_encoder),lin_ftrs_head=[] if self.hparams.linear_eval else eval(self.hparams.lin_ftrs_head),ps_head=0 if self.hparams.linear_eval else self.hparams.dropout_head,bn_head=False if self.hparams.linear_eval else not(self.hparams.no_bn_head))
        
        target_fs=100
        if(not(self.hparams.finetune)):
            print("CPC pretraining:\ndownsampling factor:",self.model_cpc.encoder_downsampling_factor,"\nchunk length(s)",self.model_cpc.encoder_downsampling_factor/target_fs,"\npixels predicted ahead:",self.model_cpc.encoder_downsampling_factor*self.hparams.steps_predicted,"\nseconds predicted ahead:",self.model_cpc.encoder_downsampling_factor*self.hparams.steps_predicted/target_fs,"\nRNN input size:",self.hparams.input_size//self.model_cpc.encoder_downsampling_factor)

    def forward(self, x):
        return self.model_cpc(x)
        
    def _step(self,data_batch, batch_idx, train):       
        if(self.hparams.finetune):
            preds = self.forward(data_batch[0])
            loss = self.criterion(preds,data_batch[1])
            self.log("train_loss" if train else "val_loss", loss)
            return {'loss':loss, "preds":preds.detach(), "targs": data_batch[1]}
        else:
            loss, acc = self.model_cpc.cpc_loss(data_batch[0],steps_predicted=self.hparams.steps_predicted,n_false_negatives=self.hparams.n_false_negatives, negatives_from_same_seq_only=self.hparams.negatives_from_same_seq_only, eval_acc=True)
            self.log("loss" if train else "val_loss", loss)
            self.log("acc" if train else "val_acc", acc)
            return loss
      
    def training_step(self, train_batch, batch_idx):
        if(self.hparams.linear_eval):
            _freeze_bn_stats(self)
        return self._step(train_batch,batch_idx,True)
        
    def validation_step(self, val_batch, batch_idx, dataloader_idx=0):
        return self._step(val_batch,batch_idx,False)
        
    def validation_epoch_end(self, outputs_all):
        if(self.hparams.finetune):
            for dataloader_idx,outputs in enumerate(outputs_all): #multiple val dataloaders
                preds_all = torch.cat([x['preds'] for x in outputs])
                targs_all = torch.cat([x['targs'] for x in outputs])
                if(self.hparams.finetune_dataset=="thew"):
                    preds_all = F.softmax(preds_all,dim=-1)
                    targs_all = torch.eye(len(self.lbl_itos))[targs_all].to(preds.device) 
                else:
                    preds_all = torch.sigmoid(preds_all)
                preds_all = preds_all.cpu().numpy()
                targs_all = targs_all.cpu().numpy()
                #instance level score
                res = eval_scores(targs_all,preds_all,classes=self.lbl_itos)
                
                idmap = self.val_dataset.get_id_mapping()
                preds_all_agg,targs_all_agg = aggregate_predictions(preds_all,targs_all,idmap,aggregate_fn=np.mean)
                res_agg = eval_scores(targs_all_agg,preds_all_agg,classes=self.lbl_itos)
                self.log_dict({"macro_auc_agg"+str(dataloader_idx):res_agg["label_AUC"]["macro"], "macro_auc_noagg"+str(dataloader_idx):res["label_AUC"]["macro"]})
                print("epoch",self.current_epoch,"macro_auc_agg"+str(dataloader_idx)+":",res_agg["label_AUC"]["macro"],"macro_auc_noagg"+str(dataloader_idx)+":",res["label_AUC"]["macro"])


    def on_fit_start(self):
        if(self.hparams.linear_eval):
            print("copying state dict before training for sanity check after training")   
            self.state_dict_pre = copy.deepcopy(self.state_dict().copy())

    
    def on_fit_end(self):
        if(self.hparams.linear_eval):
            sanity_check(self,self.state_dict_pre)
            
            
    def setup(self, stage):
        # configure dataset params
        chunkify_train = False
        chunk_length_train = self.hparams.input_size if chunkify_train else 0
        stride_train = self.hparams.input_size
        
        chunkify_valtest = True
        chunk_length_valtest = self.hparams.input_size if chunkify_valtest else 0
        stride_valtest = self.hparams.input_size//2

        train_datasets = []
        val_datasets = []
        test_datasets = []
        
        for i,target_folder in enumerate(self.hparams.data):
            target_folder = Path(target_folder)           
            
            df_mapped, lbl_itos,  mean, std = load_dataset(target_folder)
            # always use PTB-XL stats
            mean = np.array([-0.00184586, -0.00130277,  0.00017031, -0.00091313, -0.00148835,  -0.00174687, -0.00077071, -0.00207407,  0.00054329,  0.00155546,  -0.00114379, -0.00035649])
            std = np.array([0.16401004, 0.1647168 , 0.23374124, 0.33767231, 0.33362807,  0.30583013, 0.2731171 , 0.27554379, 0.17128962, 0.14030828,   0.14606956, 0.14656108])
            
            #specific for PTB-XL
            if(self.hparams.finetune and self.hparams.finetune_dataset.startswith("ptbxl")):
                if(self.hparams.finetune_dataset=="ptbxl_super"):
                    ptb_xl_label = "label_diag_superclass"
                elif(self.hparams.finetune_dataset=="ptbxl_all"):
                    ptb_xl_label = "label_all"
                    
                lbl_itos= np.array(lbl_itos[ptb_xl_label])
                
                def multihot_encode(x, num_classes):
                    res = np.zeros(num_classes,dtype=np.float32)
                    for y in x:
                        res[y]=1
                    return res
                    
                df_mapped["label"]= df_mapped[ptb_xl_label+"_filtered_numeric"].apply(lambda x: multihot_encode(x,len(lbl_itos)))
                    
            
            self.lbl_itos = lbl_itos
            tfms_ptb_xl_cpc = ToTensor() if self.hparams.normalize is False else transforms.Compose([Normalize(mean,std),ToTensor()])
            
            max_fold_id = df_mapped.strat_fold.max() #unfortunately 1-based for PTB-XL; sometimes 100 (Ribeiro)
            df_train = df_mapped[df_mapped.strat_fold<(max_fold_id-1 if self.hparams.finetune else max_fold_id)]
            df_val = df_mapped[df_mapped.strat_fold==(max_fold_id-1 if self.hparams.finetune else max_fold_id)]
            if(self.hparams.finetune):
                df_test = df_mapped[df_mapped.strat_fold==max_fold_id]
            train_datasets.append(TimeseriesDatasetCrops(df_train,self.hparams.input_size,num_classes=len(lbl_itos),data_folder=target_folder,chunk_length=chunk_length_train,min_chunk_length=self.hparams.input_size, stride=stride_train,transforms=tfms_ptb_xl_cpc,annotation=False,col_lbl ="label" if self.hparams.finetune else None,memmap_filename=target_folder/("memmap.npy")))
            val_datasets.append(TimeseriesDatasetCrops(df_val,self.hparams.input_size,num_classes=len(lbl_itos),data_folder=target_folder,chunk_length=chunk_length_valtest,min_chunk_length=self.hparams.input_size, stride=stride_valtest,transforms=tfms_ptb_xl_cpc,annotation=False,col_lbl ="label" if self.hparams.finetune else None,memmap_filename=target_folder/("memmap.npy")))
            if(self.hparams.finetune):
                test_datasets.append(TimeseriesDatasetCrops(df_test,self.hparams.input_size,num_classes=len(lbl_itos),data_folder=target_folder,chunk_length=chunk_length_valtest,min_chunk_length=self.hparams.input_size, stride=stride_valtest,transforms=tfms_ptb_xl_cpc,annotation=False,col_lbl ="label",memmap_filename=target_folder/("memmap.npy")))
            
            print("\n",target_folder)
            print("train dataset:",len(train_datasets[-1]),"samples")
            print("val dataset:",len(val_datasets[-1]),"samples")
            if(self.hparams.finetune):
                print("test dataset:",len(test_datasets[-1]),"samples")

        if(len(train_datasets)>1): #multiple data folders
            print("\nCombined:")
            self.train_dataset = ConcatDataset(train_datasets)
            self.val_dataset = ConcatDataset(val_datasets)
            print("train dataset:",len(self.train_dataset),"samples")
            print("val dataset:",len(self.val_dataset),"samples")
            if(self.hparams.finetune):
                self.test_dataset = ConcatDataset(test_datasets)
                print("test dataset:",len(self.test_dataset),"samples")
        else: #just a single data folder
            self.train_dataset = train_datasets[0]
            self.val_dataset = val_datasets[0]
            if(self.hparams.finetune):
                self.test_dataset = test_datasets[0]
    
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.hparams.batch_size, num_workers=4, shuffle=True, drop_last = True)
        
    def val_dataloader(self):
        if(self.hparams.finetune):#multiple val dataloaders
            return [DataLoader(self.val_dataset, batch_size=self.hparams.batch_size, num_workers=4),DataLoader(self.test_dataset, batch_size=self.hparams.batch_size, num_workers=4)]
        else:
            return DataLoader(self.val_dataset, batch_size=self.hparams.batch_size, num_workers=4)

    def configure_optimizers(self):
        if(self.hparams.optimizer == "sgd"):
            opt = torch.optim.SGD
        elif(self.hparams.optimizer == "adam"):
            opt = torch.optim.AdamW
        else:
            raise NotImplementedError("Unknown Optimizer.")
        
        if(self.hparams.finetune and (self.hparams.linear_eval or self.hparams.train_head_only)):
            optimizer = opt(self.model_cpc.head.parameters(), self.lr, weight_decay=self.hparams.weight_decay)
        elif(self.hparams.finetune and self.hparams.discriminative_lr_factor != 1.):#discrimative lrs
            optimizer = opt([{"params":self.model_cpc.encoder.parameters(), "lr":self.lr*self.hparams.discriminative_lr_factor*self.hparams.discriminative_lr_factor},{"params":self.model_cpc.rnn.parameters(), "lr":self.lr*self.hparams.discriminative_lr_factor},{"params":self.model_cpc.head.parameters(), "lr":self.lr}],self.hparams.lr, weight_decay=self.hparams.weight_decay)
        else:
            optimizer = opt(self.parameters(), self.lr, weight_decay=self.hparams.weight_decay)

        return optimizer
        
    def load_weights_from_checkpoint(self, checkpoint):
        """ Function that loads the weights from a given checkpoint file. 
        based on https://github.com/PyTorchLightning/pytorch-lightning/issues/525
        """
        checkpoint = torch.load(checkpoint, map_location=lambda storage, loc: storage,)
        pretrained_dict = checkpoint["state_dict"]
        model_dict = self.state_dict()
            
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        self.load_state_dict(model_dict)

#####################################################################################################
#ARGPARSERS
#####################################################################################################
def add_model_specific_args(parser):
    parser.add_argument("--input-channels", type=int, default=12)
    parser.add_argument("--normalize", action='store_true', help='Normalize input using PTB-XL stats')
    parser.add_argument('--mlp', action='store_true', help="False: original CPC True: as in SimCLR")
    parser.add_argument('--bias', action='store_true', help="original CPC: no bias")
    parser.add_argument("--n-hidden", type=int, default=512)
    parser.add_argument("--gru", action="store_true")
    parser.add_argument("--n-layers", type=int, default=2)
    parser.add_argument("--steps-predicted", dest="steps_predicted", type=int, default=12)
    parser.add_argument("--n-false-negatives", dest="n_false_negatives", type=int, default=128)
    parser.add_argument("--skip-encoder", action="store_true", help="disable the convolutional encoder i.e. just RNN; for testing")
    parser.add_argument("--fc-encoder", action="store_true", help="use a fully connected encoder (as opposed to an encoder with strided convs)")
    parser.add_argument("--negatives-from-same-seq-only", action="store_true", help="only draw false negatives from same sequence (as opposed to drawing from everywhere)")
    parser.add_argument("--no-bn-encoder", action="store_true", help="switch off batch normalization in encoder")
    parser.add_argument("--dropout-head", type=float, default=0.5)
    parser.add_argument("--train-head-only", action="store_true", help="freeze everything except classification head (note: --linear-eval defaults to no hidden layer in classification head)")
    parser.add_argument("--lin-ftrs-head", type=str, default="[512]", help="hidden layers in the classification head")
    parser.add_argument('--no-bn-head', action='store_true', help="use no batch normalization in classification head")
    return parser

def add_default_args():
    parser = argparse.ArgumentParser(description='PyTorch Lightning CPC Training')
    parser.add_argument('--data', metavar='DIR',type=str,
                        help='path(s) to dataset',action='append')
    parser.add_argument('--epochs', default=30, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('--batch-size', default=64, type=int,
                        metavar='N',
                        help='mini-batch size (default: 256), this is the total '
                             'batch size of all GPUs on the current node when '
                             'using Data Parallel or Distributed Data Parallel')
    parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float,
                        metavar='LR', help='initial learning rate', dest='lr')
    parser.add_argument('--wd', '--weight-decay', default=1e-3, type=float,
                        metavar='W', help='weight decay (default: 0.)',
                        dest='weight_decay')

    parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')

    parser.add_argument('--pretrained', default='', type=str, metavar='PATH',
                    help='path to pretrained checkpoint (default: none)')
    parser.add_argument('--optimizer', default='adam', help='sgd/adam')#was sgd
    parser.add_argument('--output-path', default='.', type=str,dest="output_path",
                        help='output path')
    parser.add_argument('--metadata', default='', type=str,
                        help='metadata for output')
    
    parser.add_argument("--gpus", type=int, default=1, help="number of gpus")
    parser.add_argument("--num-nodes", dest="num_nodes", type=int, default=1, help="number of compute nodes")
    parser.add_argument("--precision", type=int, default=16, help="16/32")
    parser.add_argument("--distributed-backend", dest="distributed_backend", type=str, default=None, help="None/ddp")
    parser.add_argument("--accumulate", type=int, default=1, help="accumulate grad batches (total-bs=accumulate-batches*bs)")
        
    parser.add_argument("--input-size", dest="input_size", type=int, default=16000)
    
    parser.add_argument("--finetune", action="store_true", help="finetuning (downstream classification task)",  default=False )
    parser.add_argument("--linear-eval", action="store_true", help="linear evaluation instead of full finetuning",  default=False )
    
    parser.add_argument(
        "--finetune-dataset",
        type=str,
        help="thew/ptbxl_super/ptbxl_all",
        default="thew"
    )
    
    parser.add_argument(
        "--discriminative-lr-factor",
        type=float,
        help="factor by which the lr decreases per layer group during finetuning",
        default=0.1
    )
    
    
    parser.add_argument(
        "--lr-find",
        action="store_true",
        help="run lr finder before training run",
        default=False
    )
    
    return parser
             
###################################################################################################
#MAIN
###################################################################################################
if __name__ == '__main__':
    parser = add_default_args()
    parser = add_model_specific_args(parser)
    hparams = parser.parse_args()
    hparams.executable = "cpc"

    if not os.path.exists(hparams.output_path):
        os.makedirs(hparams.output_path)
        
    model = LightningCPC(hparams)
    
    if(hparams.pretrained!=""):
        print("Loading pretrained weights from",hparams.pretrained)
        model.load_weights_from_checkpoint(hparams.pretrained)


    logger = TensorBoardLogger(
        save_dir=hparams.output_path,
        #version="",#hparams.metadata.split(":")[0],
        name="")
    print("Output directory:",logger.log_dir)    
    checkpoint_callback = ModelCheckpoint(
        filepath=os.path.join(logger.log_dir,"best_model"),#hparams.output_path
        save_top_k=1,
		save_last=True,
        verbose=True,
        monitor='macro_auc_agg0' if hparams.finetune else 'val_loss',#val_loss/dataloader_idx_0
        mode='max' if hparams.finetune else 'min',
        prefix='')
    lr_monitor = LearningRateMonitor()

    trainer = pl.Trainer(
        #overfit_batches=0.01,
        auto_lr_find = hparams.lr_find,
        accumulate_grad_batches=hparams.accumulate,
        max_epochs=hparams.epochs,
        min_epochs=hparams.epochs,
        
        default_root_dir=hparams.output_path,
        
        num_sanity_val_steps=0,
        
        logger=logger,
        checkpoint_callback=checkpoint_callback,
        callbacks = [],#lr_monitor],
        benchmark=True,
    
        gpus=hparams.gpus,
        num_nodes=hparams.num_nodes,
        precision=hparams.precision,
        distributed_backend=hparams.distributed_backend,
        
        progress_bar_refresh_rate=0,
        weights_summary='top',
        resume_from_checkpoint= None if hparams.resume=="" else hparams.resume)
        
    if(hparams.lr_find):#lr find
        trainer.tune(model)
        
    trainer.fit(model)