from argparse import Namespace

from io import BytesIO

import os

from typing import NamedTuple, Union, Dict, List, Optional, Tuple

from urllib.request import urlopen

from zipfile import ZipFile

import torch

import numpy as np

from sybil.serie import Serie

from sybil.models.sybil import SybilNet

from sybil.models.calibrator import SimpleClassifierGroup

from sybil.utils.logging_utils import get_logger

from sybil.utils.device_utils import get_default_device, get_most_free_gpu, get_device_mem_info

# Leaving this here for a bit; these are IDs to download the models from Google Drive

NAME_TO_FILE = {

    "sybil_base": {

        "checkpoint": ["28a7cd44f5bcd3e6cc760b65c7e0d54d"],

        "google_checkpoint_id": ["1ftYbav_BbUBkyR3HFCGnsp-h4uH1yhoz"],

        "google_calibrator_id": "1F5TOtzueR-ZUvwl8Yv9Svs2NPP5El3HY",

    },

    "sybil_1": {

        "checkpoint": ["28a7cd44f5bcd3e6cc760b65c7e0d54d"],

        "google_checkpoint_id": ["1ftYbav_BbUBkyR3HFCGnsp-h4uH1yhoz"],

        "google_calibrator_id": "1F5TOtzueR-ZUvwl8Yv9Svs2NPP5El3HY",

    },

    "sybil_2": {

        "checkpoint": ["56ce1a7d241dc342982f5466c4a9d7ef"],

        "google_checkpoint_id": ["1rscGi1grSxaVGzn-tqKtuAR3ipo0DWgA"],

        "google_calibrator_id": "1zKLVYBaiuMOx7p--e2zabs1LbQ-XXxcZ",

    },

    "sybil_3": {

        "checkpoint": ["624407ef8e3a2a009f9fa51f9846fe9a"],

        "google_checkpoint_id": ["1DV0Ge7n9r8WAvBXyoNRPwyA7VL43csAr"],

        "google_calibrator_id": "1qh4nawgE2Kjf_H97XuuTpL7XUIX7JOJn",

    },

    "sybil_4": {

        "checkpoint": ["64a91b25f84141d32852e75a3aec7305"],

        "google_checkpoint_id": ["1Acz_yzdJMpkz3PRrjXy526CjAboMEIHX"],

        "google_calibrator_id": "1QIvvCYLaesPGMEiE2Up77pKL3ygDdGU2",

    },

    "sybil_5": {

        "checkpoint": ["65fd1f04cb4c5847d86a9ed8ba31ac1a"],

        "google_checkpoint_id": ["1uV58SD-Qtb6xElTzWPDWWnloH1KB_zrP"],

        "google_calibrator_id": "1yDq1_A5w-fSdxzq4K2YSBRNcQQkDnH0K",

    },

    "sybil_ensemble": {

        "checkpoint": [

            "28a7cd44f5bcd3e6cc760b65c7e0d54d",

            "56ce1a7d241dc342982f5466c4a9d7ef",

            "624407ef8e3a2a009f9fa51f9846fe9a",

            "64a91b25f84141d32852e75a3aec7305",

            "65fd1f04cb4c5847d86a9ed8ba31ac1a",

        ],

        "google_checkpoint_id": [

            "1ftYbav_BbUBkyR3HFCGnsp-h4uH1yhoz",

            "1rscGi1grSxaVGzn-tqKtuAR3ipo0DWgA",

            "1DV0Ge7n9r8WAvBXyoNRPwyA7VL43csAr",

            "1Acz_yzdJMpkz3PRrjXy526CjAboMEIHX",

            "1uV58SD-Qtb6xElTzWPDWWnloH1KB_zrP",

        ],

        "google_calibrator_id": "1FxHNo0HqXYyiUKE_k2bjatVt9e64J9Li",

    },

}

CHECKPOINT_URL = os.getenv("SYBIL_CHECKPOINT_URL", "https://github.com/reginabarzilaygroup/Sybil/releases/download/v1.5.0/sybil_checkpoints.zip")

class Prediction(NamedTuple):

    scores: List[List[float]]

    attentions: List[Dict[str, np.ndarray]] = None

class Evaluation(NamedTuple):

    auc: List[float]

    c_index: float

    scores: List[List[float]]

    attentions: List[Dict[str, np.ndarray]] = None

def download_sybil(name, cache) -> Tuple[List[str], str]:

    """Download trained models and calibrator"""

    # Create cache folder if not exists

    cache = os.path.expanduser(cache)

    os.makedirs(cache, exist_ok=True)

    # Download models

    model_files = NAME_TO_FILE[name]

    checkpoints = model_files["checkpoint"]

    download_calib_path = os.path.join(cache, f"{name}_simple_calibrator.json")

    have_all_files = os.path.exists(download_calib_path)

    download_model_paths = []

    for checkpoint in checkpoints:

        cur_checkpoint_path = os.path.join(cache, f"{checkpoint}.ckpt")

        have_all_files &= os.path.exists(cur_checkpoint_path)

        download_model_paths.append(cur_checkpoint_path)

    if not have_all_files:

        print(f"Downloading models to {cache}")

        download_and_extract(CHECKPOINT_URL, cache)

    return download_model_paths, download_calib_path

def download_and_extract(remote_url: str, local_dir: str) -> List[str]:

    os.makedirs(local_dir, exist_ok=True)

    resp = urlopen(remote_url)

    with ZipFile(BytesIO(resp.read())) as zip_file:

        all_files_and_dirs = zip_file.namelist()

        zip_file.extractall(local_dir)

    return all_files_and_dirs

def _torch_set_num_threads(threads) -> int:

    """

    Set the number of CPU threads for torch to use.

    Set to a negative number for no-op.

    Set to 0 for the number of CPUs.

    """

    if threads < 0:

        return torch.get_num_threads()

    if threads is None or threads == 0:

        # I've never seen a benefit to going higher than 8 and sometimes there is a big slowdown

        threads = min(8, os.cpu_count())

    torch.set_num_threads(threads)

    return torch.get_num_threads()

class Sybil:

    def __init__(

        self,

        name_or_path: Union[List[str], str] = "sybil_ensemble",

        cache: str = "~/.sybil/",

        calibrator_path: Optional[str] = None,

        device: Optional[str] = None,

    ):

        """Initialize a trained Sybil model for inference.

        Parameters

        ----------

        name_or_path: list or str

            Alias to a provided pretrained Sybil model or path

            to a sybil checkpoint.

        cache: str

            Directory to download model checkpoints to

        calibrator_path: str

            Path to calibrator pickle file corresponding with model

        device: str

            If provided, will run inference using this device.

            By default, uses GPU with the most free memory, if available.

        """

        self._logger = get_logger()

        # Download if needed

        if isinstance(name_or_path, str) and (name_or_path in NAME_TO_FILE):

            name_or_path, calibrator_path = download_sybil(name_or_path, cache)

        elif not all(os.path.exists(p) for p in name_or_path):

            raise ValueError(

                "No saved model or local path: {}".format(

                    [p for p in name_or_path if not os.path.exists(p)]

                )

            )

        # Check calibrator path before continuing

        if (calibrator_path is not None) and (not os.path.exists(calibrator_path)):

            raise ValueError(f"Path not found for calibrator {calibrator_path}")

        # Set device.

        # If set manually, use it and stay there.

        # Otherwise, pick the most free GPU now and at predict time.

        self._device_flexible = True

        if device is not None:

            self.device = device

            self._device_flexible = False

        else:

            self.device = get_default_device()

        self.ensemble = torch.nn.ModuleList()

        for path in name_or_path:

            self.ensemble.append(self.load_model(path))

        self.to(self.device)

        if calibrator_path is not None:

            self.calibrator = SimpleClassifierGroup.from_json_grouped(calibrator_path)

        else:

            self.calibrator = None

    def load_model(self, path):

        """Load model from path.

        Parameters

        ----------

        path : str

            Path to a sybil checkpoint.

        Returns

        -------

        model

            Pretrained Sybil model

        """

        # Load checkpoint

        checkpoint = torch.load(path, map_location="cpu")

        args = checkpoint["args"]

        self._max_followup = args.max_followup

        self._censoring_dist = args.censoring_distribution

        model = SybilNet(args)

        # Remove model from param names

        state_dict = {k[6:]: v for k, v in checkpoint["state_dict"].items()}

        model.load_state_dict(state_dict)  # type: ignore

        if self.device is not None:

            model.to(self.device)

        # Set eval

        model.eval()

        self._logger.info(f"Loaded model from {path}")

        return model

    def _calibrate(self, scores: np.ndarray) -> np.ndarray:

        """Calibrate raw predictions

        Parameters

        ----------

        scores: np.ndarray

            risk scores as numpy array

        Returns

        -------

            np.ndarray: calibrated risk scores as numpy array

        """

        if self.calibrator is None:

            return scores

        calibrated_scores = []

        for YEAR in range(scores.shape[1]):

            probs = scores[:, YEAR].reshape(-1, 1)

            probs = self.calibrator["Year{}".format(YEAR + 1)].predict_proba(probs)[:, -1]

            calibrated_scores.append(probs)

        return np.stack(calibrated_scores, axis=1)

    def _predict(

        self,

        model: SybilNet,

        series: Union[Serie, List[Serie]],

        return_attentions: bool = False,

    ) -> Prediction:

        """Run predictions over the given serie(s).

        Parameters

        ----------

        model: SybilNet

            Instance of SybilNet

        series : Union[Serie, Iterable[Serie]]

            One or multiple series to run predictions for.

        return_attentions : bool

            If True, returns attention scores for each serie. See README for details.

        Returns

        -------

        Prediction

            Output prediction as risk scores.

        """

        if isinstance(series, Serie):

            series = [series]

        elif not isinstance(series, list):

            raise ValueError("Expected either a Serie object or list of Serie objects.")

        scores: List[List[float]] = []

        attentions: List[Dict[str, np.ndarray]] = [] if return_attentions else None

        for serie in series:

            if not isinstance(serie, Serie):

                raise ValueError("Expected a list of Serie objects.")

            volume = serie.get_volume()

            if self.device is not None:

                volume = volume.to(self.device)

            with torch.no_grad():

                out = model(volume)

                score = out["logit"].sigmoid().squeeze(0).cpu().numpy()

                scores.append(score.tolist())

                if return_attentions:

                    attentions.append(

                        {

                            "image_attention_1": out["image_attention_1"]

                            .detach()

                            .cpu(),

                            "volume_attention_1": out["volume_attention_1"]

                            .detach()

                            .cpu(),

                            "hidden": out["hidden"]

                            .detach()

                            .cpu(),

                        }

                    )

        return Prediction(scores=scores, attentions=attentions)

    def predict(

        self, series: Union[Serie, List[Serie]], return_attentions: bool = False, threads=0,

    ) -> Prediction:

        """Run predictions over the given serie(s) and ensemble

        Parameters

        ----------

        series : Union[Serie, Iterable[Serie]]

            One or multiple series to run predictions for.

        return_attentions : bool

            If True, returns attention scores for each serie. See README for details.

        threads : int

            Number of CPU threads to use for PyTorch inference.

        Returns

        -------

        Prediction

            Output prediction. See details for :class:`~sybil.model.Prediction`".

        """

        # Set CPU threads available to torch

        num_threads = _torch_set_num_threads(threads)

        self._logger.debug(f"Using {num_threads} threads for PyTorch inference")

        if self._device_flexible:

            self.device = self._pick_device()

            self.to(self.device)

        self._logger.debug(f"Beginning prediction on device: {self.device}")

        scores = []

        attentions_ = [] if return_attentions else None

        attention_keys = None

        for sybil in self.ensemble:

            pred = self._predict(sybil, series, return_attentions)

            scores.append(pred.scores)

            if return_attentions:

                attentions_.append(pred.attentions)

                if attention_keys is None:

                    attention_keys = pred.attentions[0].keys()

        scores = np.mean(np.array(scores), axis=0)

        calib_scores = self._calibrate(scores).tolist()

        attentions = None

        if return_attentions:

            attentions = []

            for i in range(len(series)):

                att = {}

                for key in attention_keys:

                    att[key] = np.stack([

                        attentions_[j][i][key] for j in range(len(self.ensemble))

                    ])

                attentions.append(att)

        return Prediction(scores=calib_scores, attentions=attentions)

    def evaluate(

        self, series: Union[Serie, List[Serie]], return_attentions: bool = False

    ) -> Evaluation:

        """Run evaluation over the given serie(s).

        Parameters

        ----------

        series : Union[Serie, List[Serie]]

            One or multiple series to run evaluation for.

        return_attentions : bool

            If True, returns attention scores for each serie. See README for details.

        Returns

        -------

        Evaluation

            Output evaluation. See details for :class:`~sybil.model.Evaluation`.

        """

        from sybil.utils.metrics import get_survival_metrics

        if isinstance(series, Serie):

            series = [series]

        elif not isinstance(series, list):

            raise ValueError(

                "Expected either a Serie object or an iterable over Serie objects."

            )

        # Check all have labels

        if not all(serie.has_label() for serie in series):

            raise ValueError("All series must have a label for evaluation")

        # Get scores and labels

        predictions = self.predict(series, return_attentions)

        scores = predictions.scores

        labels = [serie.get_label(self._max_followup) for serie in series]

        # Convert to format for survival metrics

        input_dict = {

            "probs": torch.tensor(scores),

            "censors": torch.tensor([label.censor_time for label in labels]),

            "golds": torch.tensor([label.y for label in labels]),

        }

        args = Namespace(

            max_followup=self._max_followup, censoring_distribution=self._censoring_dist

        )

        out = get_survival_metrics(input_dict, args)

        auc = [float(out[f"{i + 1}_year_auc"]) for i in range(self._max_followup)]

        c_index = float(out["c_index"])

        return Evaluation(auc=auc, c_index=c_index, scores=scores, attentions=predictions.attentions)

    def to(self, device: str):

        """Move model to device.

        Parameters

        ----------

        device : str

            Device to move model to.

        """

        self.device = device

        self.ensemble.to(device)

    def _pick_device(self):

        """

        Pick the device to run inference on.

        This is based on the device with the most free memory, with a preference for remaining

        on the current device.

        Motivation is to enable multiprocessing without the processes needed to communicate.

        """

        if not torch.cuda.is_available():

            return get_default_device()

        # Get size of the model in memory (approximate)

        model_mem = 9*sum(p.numel() * p.element_size() for p in self.ensemble.parameters())

        # Check memory available on current device.

        # If it seems like we're the only thing on this GPU, stay.

        free_mem, total_mem = get_device_mem_info(self.device)

        cur_allocated = total_mem - free_mem

        min_to_move = int(1.01 * model_mem)

        if cur_allocated < min_to_move:

            return self.device

        else:

            # Otherwise, get the most free GPU

            return get_most_free_gpu()