from collections import OrderedDict
from argparse import Namespace
import pickle
import os
import sys

import pytorch_lightning as pl
import torch

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from sybil.utils.helpers import get_dataset
import sybil.utils.losses as losses
import sybil.utils.metrics as metrics
import sybil.utils.loading as loaders
import sybil.models.sybil as model
from sybil.parsing import parse_args



class SybilLightning(pl.LightningModule):
    """
    Lightning Module
    Methods:
        .log/.log_dict: log inputs to logger
    Notes:
        *_epoch_end method returns None
        self can log additional data structures to logger
        with self.logger.experiment.log_*
        (*= 'text', 'image', 'audio', 'confusion_matrix', 'histogram')
    """

    def __init__(self, args):
        super(SybilLightning, self).__init__()
        if isinstance(args, dict):
            args = Namespace(**args)
        self.args = args
        self.model = model.SybilNet(args)
        self.save_prefix = "default"
        self.save_hyperparameters(args)
        self._list_of_metrics = [
            metrics.get_classification_metrics,
            metrics.get_survival_metrics,
            metrics.get_risk_metrics
        ]

    def set_finetune(self, finetune_flag):
        return

    def forward(self, x):
        return self.model(x)

    def step(self, batch, batch_idx, optimizer_idx, log_key_prefix=""):
        model_output = self(batch["x"])
        logging_dict, predictions_dict = OrderedDict(), OrderedDict()

        if "exam" in batch:
            predictions_dict["exam"] = batch["exam"]
        if "y" in batch:
            predictions_dict["golds"] = batch["y"]

        if self.args.save_attention_scores:
            attentions = {k: v for k, v in model_output.items() if "attention" in k}
            predictions_dict.update(attentions)

        loss_fns = self.get_loss_functions(self.args)
        loss = 0
        for loss_fn in loss_fns:
            local_loss, local_log_dict, local_predictions_dict = loss_fn(
                model_output, batch, self, self.args
            )
            loss += local_loss
            logging_dict.update(local_log_dict)
            predictions_dict.update(local_predictions_dict)
        logging_dict = prefix_dict(logging_dict, log_key_prefix)
        predictions_dict = prefix_dict(predictions_dict, log_key_prefix)
        return loss, logging_dict, predictions_dict, model_output

    def training_step(self, batch, batch_idx, optimizer_idx=None):
        result = OrderedDict()
        loss, logging_dict, predictions_dict, _ = self.step(
            batch, batch_idx, optimizer_idx, log_key_prefix="train_"
        )
        logging_dict["train_loss"] = loss.detach()
        self.log_dict(logging_dict, prog_bar=False, on_step=True, on_epoch=True)
        result["logs"] = logging_dict
        self.log_tensor_dict(predictions_dict, prog_bar=False, logger=False)
        result.update(predictions_dict)
        # lightning expects 'loss' key in output dict. ow loss := None by default
        result["loss"] = loss
        return result

    def validation_step(self, batch, batch_idx, optimizer_idx=None):
        result = OrderedDict()
        loss, logging_dict, predictions_dict, _ = self.step(
            batch, batch_idx, optimizer_idx, log_key_prefix="val_"
        )
        logging_dict["val_loss"] = loss.detach()
        self.log_dict(logging_dict, prog_bar=True, sync_dist=True)
        result["logs"] = logging_dict
        if self.args.accelerator == "ddp":
            predictions_dict = gather_predictions_dict(predictions_dict)
        self.log_tensor_dict(predictions_dict, prog_bar=False, logger=False)
        result.update(predictions_dict)
        return result

    def test_step(self, batch, batch_idx, optimizer_idx=None):
        result = OrderedDict()
        loss, logging_dict, predictions_dict, model_output = self.step(
            batch, batch_idx, optimizer_idx, log_key_prefix="test_"
        )
        logging_dict["{}_loss".format(self.save_prefix)] = loss.detach()
        result["logs"] = logging_dict

        if self.args.accelerator == "ddp":
            predictions_dict = gather_predictions_dict(predictions_dict)

        self.log_tensor_dict(predictions_dict, prog_bar=False, logger=False)
        result.update(predictions_dict)
        return result

    def training_epoch_end(self, outputs):
        if len(outputs) == 0:
            return
        outputs = gather_step_outputs(outputs)
        # loss already logged in progress_bar_dict (get_progress_bar_dict()),
        # and logging twice creates issue
        del outputs["loss"]
        epoch_metrics = compute_epoch_metrics(
            self._list_of_metrics, outputs, self.args, self.device, key_prefix="train_"
        )
        for k, v in outputs["logs"].items():
            epoch_metrics[k] = v.mean()
        self.log_dict(epoch_metrics, prog_bar=True, logger=True)

    def validation_epoch_end(self, outputs):
        if len(outputs) == 0:
            return
        outputs = gather_step_outputs(outputs)
        epoch_metrics = compute_epoch_metrics(
            self._list_of_metrics, outputs, self.args, self.device, key_prefix="val_"
        )
        for k, v in outputs["logs"].items():
            epoch_metrics[k] = v.mean()
        self.log_dict(epoch_metrics, prog_bar=True, logger=True)

    def test_epoch_end(self, outputs):
        self.save_prefix= 'test'
        if len(outputs) == 0:
            return
        outputs = gather_step_outputs(outputs)
        epoch_metrics = compute_epoch_metrics(
            self._list_of_metrics, outputs, self.args, self.device, key_prefix="test_"
        )

        for k, v in outputs["logs"].items():
            epoch_metrics[k] = v.mean()

        self.log_dict(epoch_metrics, prog_bar=True, logger=True)

        # Dump metrics for use by dispatcher
        metrics_dict = {
            k[len(self.save_prefix) :]: v.mean().item()
            for k, v in outputs.items()
            if "loss" in k
        }
        metrics_dict.update(
            {
                k[len(self.save_prefix) :]: v.mean().item()
                for k, v in epoch_metrics.items()
            }
        )
        metrics_filename = "{}.{}.metrics".format(
            self.args.results_path, self.save_prefix
        )
        pickle.dump(metrics_dict, open(metrics_filename, "wb"))
        if self.args.save_predictions and self.global_rank == 0:
            predictions_dict = {
                k: v.cpu() if isinstance(v, torch.Tensor) else v
                for k, v in outputs.items()
            }
            predictions_filename = "{}.{}.predictions".format(
                self.args.results_path, self.save_prefix
            )
            pickle.dump(predictions_dict, open(predictions_filename, "wb"))

    def configure_optimizers(self):
        """
        Helper function to fetch optimizer based on args.
        """
        params = [param for param in self.model.parameters() if param.requires_grad]
        if self.args.optimizer == "adam":
            optimizer = torch.optim.Adam(
                params, lr=self.args.lr, weight_decay=self.args.weight_decay
            )
        elif self.args.optimizer == "adagrad":
            optimizer = torch.optim.Adagrad(
                params, lr=self.args.lr, weight_decay=self.args.weight_decay
            )
        elif self.args.optimizer == "sgd":
            optimizer = torch.optim.SGD(
                params,
                lr=self.args.lr,
                weight_decay=self.args.weight_decay,
                momentum=self.args.momentum,
            )
        else:
            raise Exception("Optimizer {} not supported!".format(self.args.optimizer))

        scheduler = {
            "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
                optimizer,
                patience=self.args.patience,
                factor=self.args.lr_decay,
                mode="min" if "loss" in self.args.tuning_metric else "max",
            ),
            "monitor": "val_{}".format(self.args.tuning_metric),
            "interval": "epoch",
            "frequency": 1,
        }
        return [optimizer], [scheduler]

    def log_tensor_dict(
        self,
        output,
        prog_bar=False,
        logger=True,
        on_step=None,
        on_epoch=None,
        sync_dist=False,
    ):
        dict_of_tensors = {
            k: v.float() for k, v in output.items() if isinstance(v, torch.Tensor)
        }
        self.log_dict(
            dict_of_tensors,
            prog_bar=prog_bar,
            logger=logger,
            on_step=on_step,
            on_epoch=on_epoch,
            sync_dist=sync_dist,
        )

    def get_loss_functions(self, args):
        loss_fns = [losses.get_survival_loss]
        if args.use_annotations:
            loss_fns.append(losses.get_annotation_loss)

        return loss_fns


def prefix_dict(d, prefix):
    r = OrderedDict()
    for k, v in d.items():
        r[prefix + k] = v
    return r


def gather_predictions_dict(predictions):
    gathered_preds = {
        k: concat_all_gather(v) if isinstance(v, torch.Tensor) else v
        for k, v in predictions.items()
    }
    return gathered_preds


def gather_step_outputs(outputs):
    output_dict = OrderedDict()
    if isinstance(outputs[-1], list):
        outputs = outputs[0]

    for k in outputs[-1].keys():
        if k == "logs":
            output_dict[k] = gather_step_outputs([output["logs"] for output in outputs])
        elif (
            isinstance(outputs[-1][k], torch.Tensor) and len(outputs[-1][k].shape) == 0
        ):
            output_dict[k] = torch.stack([output[k] for output in outputs])
        elif isinstance(outputs[-1][k], torch.Tensor):
            output_dict[k] = torch.cat([output[k] for output in outputs], dim=0)
        else:
            output_dict[k] = [output[k] for output in outputs]
    return output_dict


@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """

    tensors_gather = [
        torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
    ]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
    output = torch.cat(tensors_gather, dim=0)
    return output


def compute_epoch_metrics(list_of_metrics, result_dict, args, device, key_prefix=""):
    stats_dict = OrderedDict()

    """
        Remove prefix from keys. For instance, convert:
        val_probs -> probs for standard handling in the metric fucntions
    """
    result_dict_wo_key_prefix = {}

    for k, v in result_dict.items():
        if isinstance(v, list) and isinstance(v[-1], torch.Tensor):
            v = torch.cat(v, dim=-1)
        if isinstance(v, torch.Tensor):
            v = v.cpu().numpy()
        if k == "meta":
            continue
        if key_prefix != "" and k.startswith(key_prefix):
            k_wo_prefix = k[len(key_prefix) :]
            result_dict_wo_key_prefix[k_wo_prefix] = v
        else:
            result_dict_wo_key_prefix[k] = v

    for k, v in result_dict["logs"].items():
        if k.startswith(key_prefix):
            result_dict_wo_key_prefix[k[len(key_prefix) :]] = v

    for metric_func in list_of_metrics:
        stats_wo_prefix = metric_func(result_dict_wo_key_prefix, args)
        for k, v in stats_wo_prefix.items():
            stats_dict[key_prefix + k] = torch.tensor(v, device=device)

    return stats_dict


def train(args):
    if not args.turn_off_checkpointing:
        checkpoint_callback = pl.callbacks.ModelCheckpoint(
            dirpath=args.save_dir,
            save_top_k=1,
            verbose=True,
            monitor="val_{}".format(args.tuning_metric)
            if args.tuning_metric is not None
            else None,
            save_last=True,
            mode="min"
            if args.tuning_metric is not None and "loss" in args.tuning_metric
            else "max",
        )
        args.callbacks = [checkpoint_callback]
    trainer = pl.Trainer.from_argparse_args(args)
    # Remove callbacks from args for safe pickling later
    args.callbacks = None
    args.num_nodes = trainer.num_nodes
    args.num_processes = trainer.num_processes
    args.world_size = args.num_nodes * args.num_processes
    args.global_rank = trainer.global_rank
    args.local_rank = trainer.local_rank

    train_dataset = loaders.get_train_dataset_loader(
        args, get_dataset(args.dataset, "train", args)
    )
    dev_dataset = loaders.get_eval_dataset_loader(
        args, get_dataset(args.dataset, "dev", args), False
    )

    args.censoring_distribution = metrics.get_censoring_dist(train_dataset.dataset)
    module = SybilLightning(args)

    # print args
    for key, value in sorted(vars(args).items()):
        print("{} -- {}".format(key.upper(), value))

    if args.snapshot is not None:
        module = module.load_from_checkpoint(checkpoint_path= args.snapshot, strict=False)
        module.args = args
    
    trainer.fit(module, train_dataset, dev_dataset)
    args.model_path = trainer.checkpoint_callback.best_model_path
    print("Saving args to {}".format(args.results_path))
    pickle.dump(vars(args), open(args.results_path, "wb"))

def test(args):
    trainer = pl.Trainer.from_argparse_args(args)
    # Remove callbacks from args for safe pickling later
    args.callbacks = None
    args.num_nodes = trainer.num_nodes
    args.num_processes = trainer.num_processes
    args.world_size = args.num_nodes * args.num_processes
    args.global_rank = trainer.global_rank
    args.local_rank = trainer.local_rank

    train_dataset = loaders.get_train_dataset_loader(
        args, get_dataset(args.dataset, "train", args)
    )
    test_dataset = loaders.get_eval_dataset_loader(
        args, get_dataset(args.dataset, "test", args), False
    )

    args.censoring_distribution = metrics.get_censoring_dist(train_dataset.dataset)
    module = SybilLightning(args)
    module = module.load_from_checkpoint(checkpoint_path= args.snapshot, strict=False)
    module.args = args

    # print args
    for key, value in sorted(vars(args).items()):
        print("{} -- {}".format(key.upper(), value))

    trainer.test(module, test_dataset)

    print("Saving args to {}".format(args.results_path))
    pickle.dump(vars(args), open(args.results_path, "wb"))

if __name__ == "__main__":
    args = parse_args()
    if args.train:
        train(args)
    elif args.test:
        test(args)