import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import torch

from datasets.dataset_survival import Generic_MIL_Survival_Dataset
from lifelines.utils import concordance_index
from models.model_amil import AMIL
from models.model_mil import MIL_fc_Surv
from pycox.evaluation import EvalSurv
from utils.utils import *


def initiate_model(settings, ckpt_path):
    print('Initialize model ...', end=' ')
    model_dict = {"dropout": settings['drop_out']}
    
    if settings['model_size'] is not None and settings['model_type'] == 'amil':
        model_dict.update({"size_arg": settings['model_size']})
    
    if settings['model_type'] =='amil':
        model = AMIL(**model_dict)
    elif settings['model_type'] == 'mil':
        model = MIL_fc_Surv(**model_dict)
    else:
        raise NotImplementedError

    ckpt = torch.load(ckpt_path)
    ckpt_clean = {}
    for key in ckpt.keys():
        if 'instance_loss_fn' in key:
            continue
        ckpt_clean.update({key.replace('.module', ''):ckpt[key]})
    model.load_state_dict(ckpt_clean, strict=True)

    model.relocate()
    model.eval()
    print('Done.')

    if settings['print_model_info']:
        print_network(model)  

    return model



class _BaseEvaluationData:
    event_col = 'event'
    time_col = 'time'
    
    def __init__(self, settings):
        print('Initialize data ...', end=' ')
        self.dataset = Generic_MIL_Survival_Dataset(csv_path = settings['csv_path'],
            data_dir= os.path.join(settings['data_root_dir'], settings['feature_dir']),
            shuffle = False, 
            print_info = settings['print_data_info'],
            label_dict = {'lebt':0, 'tod':1},
            event_col = self.event_col,
            time_col = self.time_col,
            patient_strat=True,
            ignore=[])

        self.split_path = '{}/splits_{}.csv'.format(settings['split_dir'], settings['split_idx'])
        print('Done.')

    def _get_split_data(self, split):
        assert split in ['train', 'val', 'test', 'all'], 'Split {} not recognized, must be in [train, val, test, all]'.format(split)
        train, val, test = self.dataset.return_splits(from_id=False, csv_path=self.split_path)
        if split == 'train':
            loader = get_simple_loader(train, survival=True)
        elif split == 'val':
            loader = get_simple_loader(val, survival=True)
        elif split == 'test':
            loader = get_simple_loader(test, survival=True)
        elif split == 'all':
            loader = get_simple_loader(self.dataset, survival=True)
        return loader, loader.dataset.slide_data



class _BaseEvaluationAMIL(_BaseEvaluationData):
    def __init__(self, settings):
        super().__init__(settings)

        # init model
        ckpt_path = os.path.join(settings['models_dir'], 's_{}_checkpoint.pt'.format(settings['split_idx']))
        self.model = initiate_model(settings, ckpt_path)

        self.baseline_hazard = None
        self.baseline_cumulutative_hazard = None

        self.patient_predictions = None
        self.c_index = None
        self.c_index_td = None
        self.ibs = None
        self.inbll = None


    def _compute_risk(self, loader):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        risks = []
        events = []
        times = []
        # print('Collect patient predictions ...', end=' ')
        for batch_idx, (data, event, time) in enumerate(loader):
            with torch.no_grad():
                risk, _ , _ = self.model(data.to(device))
            risks.append(risk.item())
            events.append(event.item())
            times.append(time.item())
        # print('Done.')
        return np.asarray(times), np.asarray(events), np.asarray(risks)
            

    def _compute_baseline_harzards(self):
        """Computes the Breslow esimates from the training data.

        Modified from https://github.com/havakv/pycox/blob/0e9d6f9a1eff88a355ead11f0aa68bfb94647bf8/pycox/models/cox.py#L63        
        """
        loader, dataset = self._get_split_data('train')
        _, _, risk_scores = self._compute_risk(loader)
        return  (dataset 
                .assign(exp_risk=np.exp(risk_scores)) 
                .groupby(dataset.time) 
                .agg({'exp_risk': 'sum', 'event': 'sum'})
                .sort_index(ascending=False) 
                .assign(exp_risk=lambda x: x['exp_risk'].cumsum())  
                .pipe(lambda x: x['event']/x['exp_risk']) 
                .iloc[::-1]
                .rename('baseline_hazards'))

    def _compute_baseline_cumulative_hazards(self):
        """Computes baseline and baseline cumulative hazards and stores as class variable"""
        print('Estimate baseline cumulative hazard ...', end=' ')
        base_hazard = self._compute_baseline_harzards()
        self.baseline_hazard = base_hazard
        self.baseline_cumulutative_hazard = base_hazard.cumsum().rename('baseline_cumulative_hazards')
        print('Done.')


    def _predict_survival_function(self, loader):
        """Predicts survival function for given data loader."""
        if self.baseline_cumulutative_hazard is None:
            self._compute_baseline_cumulative_hazards()

        base_ch = self.baseline_cumulutative_hazard.values.reshape(-1, 1).astype(float)
        times, events, risks = self._compute_risk(loader)
        exp_risk = np.exp(risks).reshape(1, -1)
        surv = np.exp(-base_ch.dot(exp_risk))
        return times, events, torch.from_numpy(surv)

    def _predict_risk(self, loader):
        times, events, risks = self._compute_risk(loader)
        return times, events, risks


    def _collect_patient_ids(self, split):
        loader, dataset = self._get_split_data(split)
        return dataset.index


    def _unpack_data(self, data):
        times = [data[patient]['time'] for patient in data]
        events = [data[patient]['event'] for patient in data]
        predictions = [data[patient]['probabilities'] for patient in data]
        return times, events, predictions

    def _compute_c_index(self, data):
        times, events, predictions = self._unpack_data(data)
        probs_by_interval = torch.stack(predictions).permute(1, 0)
        c_index = [concordance_index(event_times=times,
                                     predicted_scores=interval_probs,
                                     event_observed=events)
                   for interval_probs in probs_by_interval]
        return c_index

    def _predictions_to_pycox(self, data, time_points=None):
        predictions = {k: v['probabilities'] for k, v in data.items()}
        df = pd.DataFrame.from_dict(predictions)
        return df



class EvaluationAMIL(_BaseEvaluationAMIL):
    def __init__(self, settings):
        super().__init__(settings)

        self.split = None

    def _check_split_data(self, split):
        if self.split is None:
            self.split = split
        elif self.split != split:
            self.patient_predictions = None
            self.c_index = None
            self.c_index_td = None
            self.ibs = None
            self.inbll = None

    def _collect_patient_predictions(self, split):
        patient_data = dict()
        loader, _ = self._get_split_data(split)
        pids = self._collect_patient_ids(split)
        times, events, surv = self._predict_survival_function(loader)
        for i, patient in enumerate(pids):
            patient_data[patient] = {'time': times[i],
                                        'event': events[i],
                                        'probabilities': surv[:, i]}
        return patient_data


    def _compute_pycox_metrics(self, data, time_points=None,
                               drop_last_times=0):
        times, events, _ = self._unpack_data(data)
        times, events = np.array(times), np.array(events)
        predictions = self._predictions_to_pycox(data, time_points)

        ev = EvalSurv(predictions, times, events, censor_surv='km')
        # Using "antolini" method instead of "adj_antolini" resulted in Ctd
        # values different from C-index for proportional hazard methods (for
        # CNV data); this is presumably due to the tie handling, since that is
        # what the pycox authors "adjust" (see code comments at:
        # https://github.com/havakv/pycox/blob/6ed3973954789f54453055bbeb85887ded2fb81c/pycox/evaluation/eval_surv.py#L171)
        # c_index_td = ev.concordance_td('antolini')
        c_index_td = ev.concordance_td('adj_antolini')

        # time_grid = np.array(predictions.index)
        # Use 100-point time grid based on data
        time_grid = np.linspace(times.min(), times.max(), 100)
        # Since the score becomes unstable for the highest times, drop the last
        # time points?
        if drop_last_times > 0:
            time_grid = time_grid[:-drop_last_times]
        ibs = ev.integrated_brier_score(time_grid)
        inbll = ev.integrated_nbll(time_grid)

        return c_index_td, ibs, inbll


    def compute_metrics(self, split, time_points=None):
        """Calculate evaluation metrics."""
        self._check_split_data(split)

        print('Compute evaluation metrics ... \n', end =' ')
        if self.patient_predictions is None:
            # Get all patient labels and predictions
            self.patient_predictions = self._collect_patient_predictions(split)

        if self.c_index is None:
            self.c_index = self._compute_c_index(self.patient_predictions)

        if self.c_index_td is None:
            td_metrics = self._compute_pycox_metrics(self.patient_predictions,
                                                     time_points)
            self.c_index_td, self.ibs, self.inbll = td_metrics
        print('Done.')


    def predict_risk(self, split):
        loader, _ = self._get_split_data(split)
        return self._predict_risk(loader)


    def return_results(self):
        assert all([
            self.c_index, 
            self.c_index_td,
            self.ibs,
            self.inbll
        ]),  'Results not available.' + \
                ' Please call "compute_metrics" or "run_bootstrap" first.'

        return (
            self.c_index, 
            self.c_index_td,
            self.ibs,
            self.inbll
            )