from __future__ import annotations

import os
import torch
import onnx
import numpy as np

from pathlib import Path
from torch.utils.data import DataLoader
from typing import Union, Optional, Dict, Iterable, Tuple, Any, List

from ..losses import get_loss
from ..models import get_model
from ..protocols import Protocol
from ..optimizers import get_optimizer
from ..output_files import InferenceOutputManager
from ..datasets import get_dataset, get_embeddings_collate_function
from ..solvers import get_solver, get_mean_and_confidence_bounds, MetricsCalculator
from ..utilities import seed_all, EmbeddingDatasetSample, MASK_AND_LABELS_PAD_VALUE, revert_mappings


class Inferencer:

    def __init__(
            self,
            iom: InferenceOutputManager,
    ):
        self.protocol = iom.protocol()
        self.embedder_name = iom.embedder_name()
        self.embedding_dimension = iom.n_features()
        self.class_int2str = iom.class_int2str()
        self.class_str2int = iom.class_str2int()
        self.device = iom.device()
        self.collate_function = get_embeddings_collate_function(self.protocol)

        self.solvers_and_loaders_by_split = self._create_solvers_and_loaders_by_split(iom)
        print(f"Got {len(self.solvers_and_loaders_by_split.keys())} split(s): "
              f"{', '.join(self.solvers_and_loaders_by_split.keys())}")

    @classmethod
    def create_from_out_file(cls, out_file_path: str,
                             automatic_path_correction: bool = True,
                             ) -> Tuple[Inferencer, InferenceOutputManager]:
        """
        Create the inferencer object from the out.yml file generated by biotrainer.
        Reads the out.yml file without the split ids which would blow up the file unnecessarily.

        :param out_file_path: Path to out.yml file generated by biotrainer
        :param automatic_path_correction: If True, the method tries to correct the log_dir path if it is not found.
                                          In this case, checkpoints are searched at
                                          out_file_path/model_choice/embedder_name

        :return: Tuple with Inferencer object configured with the output variables from the out.yml file,
                 and the output variables as dict
        """
        inference_output_manager = InferenceOutputManager(Path(out_file_path),
                                                          automatic_path_correction=automatic_path_correction)
        return cls(inference_output_manager), inference_output_manager

    def _create_solvers_and_loaders_by_split(self, iom: InferenceOutputManager) -> Dict[str, Tuple[Any, Any]]:
        result_dict = {}
        splits = iom.training_results().keys()
        log_dir = iom.log_dir()
        split_checkpoints = {file.split("_checkpoint.")[0]: file for file in os.listdir(log_dir) if
                             (Path(log_dir) / Path(file)).is_file() and file.endswith(".safetensors")}
        for split_name in splits:
            # Ignore average or best result
            if "average" in split_name or "best" in split_name:
                continue
            split_config = iom.split_config(split_name=split_name)
            split_config["protocol"] = self.protocol
            split_config["device"] = self.device
            # Positional arguments
            #model_choice = split_config.pop("model_choice")
            #n_classes = split_config.pop("n_classes")
            #loss_choice = split_config.pop("loss_choice")
            #optimizer_choice = split_config.pop("optimizer_choice")
            #learning_rate = split_config.pop("learning_rate")
            #log_dir = split_config.pop("log_dir")
            #disable_pytorch_compile = split_config.pop("disable_pytorch_compile")
            checkpoint_path = Path(log_dir) / Path(split_checkpoints[split_name])

            model = get_model(**split_config)
            loss_function = get_loss(**split_config)
            optimizer = get_optimizer(model_parameters=model.parameters(),
                                      **split_config
                                      )

            solver = get_solver(name=split_name, network=model, optimizer=optimizer,
                                loss_function=loss_function,
                                **split_config)
            solver.load_checkpoint(checkpoint_path=checkpoint_path, resume_training=False,
                                   disable_pytorch_compile=iom.disable_pytorch_compile())

            def dataloader_function(dataset):
                return DataLoader(dataset=dataset, batch_size=split_config["batch_size"],
                                  shuffle=False, drop_last=False,
                                  collate_fn=self.collate_function)

            result_dict[split_name] = (solver, dataloader_function)
        return result_dict

    def _convert_class_str2int(self, to_convert: str):
        if type(to_convert) is str:
            if self.protocol in Protocol.per_residue_protocols():
                return [self.class_str2int[t] for t in to_convert]
            else:
                return self.class_str2int[to_convert]
        else:
            return to_convert

    @staticmethod
    def _pad_tensor(protocol: Protocol, target: Union[Any, torch.Tensor], length_to_pad: int, device):
        target_tensor = torch.as_tensor(target, device=device)
        if protocol in Protocol.per_residue_protocols():
            if target_tensor.shape[0] < length_to_pad:
                padding_size = length_to_pad - target_tensor.shape[0]
                padding = torch.full((padding_size,), MASK_AND_LABELS_PAD_VALUE, dtype=target_tensor.dtype,
                                     device=device)
                return torch.cat([target_tensor, padding])
            else:
                return target_tensor
        else:
            return target_tensor

    def _convert_target_dict(self, target_dict: Dict[str, str]):
        if self.protocol in Protocol.classification_protocols():
            if self.protocol in Protocol.per_residue_protocols():
                max_prediction_length = len(max(target_dict.values(), key=len))
                return {seq_id: self._pad_tensor(protocol=self.protocol, target=self._convert_class_str2int(prediction),
                                                 length_to_pad=max_prediction_length, device=self.device)
                        for seq_id, prediction in target_dict.items()}
            else:
                return {seq_id: torch.tensor(self._convert_class_str2int(prediction),
                                             device=self.device)
                        for seq_id, prediction in target_dict.items()}
        return {seq_id: torch.tensor(prediction,
                                     device=self.device)
                for seq_id, prediction in target_dict.items()}

    def _load_solver_and_dataloader(self, embeddings: Union[Iterable, Dict],
                                    split_name, targets: Optional[List] = None):
        if split_name not in self.solvers_and_loaders_by_split.keys():
            raise ValueError(f"Unknown split_name {split_name} for given configuration!")

        if isinstance(embeddings, Dict):
            embeddings_dict = embeddings
        else:
            embeddings_dict = {str(idx): embedding for idx, embedding in enumerate(embeddings)}

        if targets and self.protocol in Protocol.classification_protocols():
            targets = [self._convert_class_str2int(target) for target in targets]

        solver, loader = self.solvers_and_loaders_by_split[split_name]
        dataset = get_dataset(samples=[
            EmbeddingDatasetSample(seq_id, torch.tensor(np.array(embedding)),
                                   torch.empty(1) if not targets else torch.tensor(np.array(targets[idx])))
            for idx, (seq_id, embedding) in enumerate(embeddings_dict.items())
        ], finetuning=False, random_masking=False, mask_token=None, class_str2int=None)
        dataloader = loader(dataset)
        return solver, dataloader

    def convert_all_checkpoints_to_safetensors(self) -> None:
        """
        Converts all checkpoint files for the splits from .pt to .safetensors, if not already stored as .safetensors
        """
        for split_name, (solver, _) in self.solvers_and_loaders_by_split.items():
            if "pt" in solver.checkpoint_type:
                solver.save_checkpoint(solver.start_epoch)

    def convert_to_onnx(self, output_dir: Optional[str] = None) -> List[str]:
        """
        Converts the model to ONNX format for the given embedding dimension.

        :param output_dir: The directory to save the ONNX files. If not provided, the ONNX files will be saved in the
            solver's log directory. Defaults to None.

        :return: A list of file paths where the ONNX files are saved.
        """
        result_file_paths = []
        for split_name, (solver, _) in self.solvers_and_loaders_by_split.items():
            onnx_save_path = solver.save_as_onnx(embedding_dimension=self.embedding_dimension, output_dir=output_dir)
            result_file_paths.append(onnx_save_path)
        return result_file_paths

    def from_embeddings(self, embeddings: Union[Iterable, Dict], targets: Optional[List] = None,
                        split_name: str = "hold_out",
                        include_probabilities: bool = False) -> Dict[str, Union[Dict, str, int, float]]:
        """
        Calculate predictions from embeddings.

        :param embeddings: Iterable or dictionary containing the input embeddings to predict on.
        :param targets: Iterable that contains the targets to calculate metrics
        :param split_name: Name of the split to use for prediction. Default is "hold_out".
        :param include_probabilities: If True, the probabilities used to predict classes are also reported.
                                      Is only useful for classification tasks, otherwise the "probabilities" are the
                                      same as the predictions.
        :return: Dictionary containing the following sub-dictionaries:
                 - 'metrics': Calculated metrics if 'targets' are given, otherwise 'None'.
                 - 'mapped_predictions': Class or value prediction from the given embeddings.
                 - 'mapped_probabilities': Probabilities for classification tasks if include_probabilities is True.
                 Predictions and probabilities are either 'mapped' to keys from an embeddings dict or indexes if
                 embeddings are given as a list.
        """
        solver, dataloader = self._load_solver_and_dataloader(embeddings, split_name, targets)

        inference_dict = solver.inference(dataloader, calculate_test_metrics=targets is not None)
        predictions = inference_dict["mapped_predictions"]

        # For class predictions, revert from int (model output) to str (class name)
        inference_dict["mapped_predictions"] = revert_mappings(protocol=self.protocol, test_predictions=predictions,
                                                               class_int2str=self.class_int2str)
        inference_dict["mapped_probabilities"] = {k: v.cpu().tolist() if v is torch.tensor else v for k, v in
                                                  inference_dict["mapped_probabilities"].items()}

        if not include_probabilities:
            return {k: v for k, v in inference_dict.items() if k != "mapped_probabilities"}
        else:
            return inference_dict

    def from_embeddings_with_bootstrapping(self, embeddings: Union[Iterable, Dict], targets: List,
                                           split_name: str = "hold_out",
                                           iterations: int = 30,
                                           sample_size: int = -1,
                                           confidence_level: float = 0.05,
                                           seed: int = 42) -> Dict[str, Dict[str, float]]:
        """
        Calculate predictions from embeddings.

        :param embeddings: Iterable or dictionary containing the input embeddings to predict on.
        :param targets: Iterable that contains the targets to calculate metrics
        :param split_name: Name of the split to use for prediction. Default is "hold_out".
        :param iterations: Number of iterations to perform bootstrapping
        :param sample_size: Sample size to use for bootstrapping. -1 defaults to all embeddings which is recommended.
                            It is possible, but not recommended to use a sample size larger or smaller
                            than the number of embeddings, because this might render the variance estimate unreliable.
                            See: https://math.mit.edu/~dav/05.dir/class24-prep-a.pdf (6.2)
        :param confidence_level: Confidence level for result error intervals (0.05 => 95% percentile)
        :param seed: Seed to use for the bootstrapping algorithm
        :return: Dictionary containing the following sub-dictionaries:
                 - 'metrics': Calculated metrics if 'targets' are given, otherwise 'None'.
                 - 'mapped_predictions': Class or value prediction from the given embeddings.
                 - 'mapped_probabilities': Probabilities for classification tasks if include_probabilities is True.
                 Predictions and probabilities are either 'mapped' to keys from an embeddings dict or indexes if
                 embeddings are given as a list.
        """
        if not 0 < confidence_level < 1:
            raise Exception(f"Confidence level must be between 0 and 1, given: {confidence_level}!")

        seed_all(seed)

        if isinstance(embeddings, Dict):
            embeddings_dict = embeddings
        else:
            embeddings_dict = {str(idx): embedding for idx, embedding in enumerate(embeddings)}

        seq_ids = list(embeddings_dict.keys())

        all_predictions = self.from_embeddings(embeddings_dict, targets)["mapped_predictions"]
        all_predictions_dict = self._convert_target_dict(all_predictions)

        all_targets_dict = {seq_id: targets[idx] for idx, seq_id in enumerate(seq_ids)}
        all_targets_dict = self._convert_target_dict(all_targets_dict)

        solver, _ = self.solvers_and_loaders_by_split[split_name]

        return self._do_bootstrapping(iterations=iterations, sample_size=sample_size, confidence_level=confidence_level,
                                      seq_ids=seq_ids, all_predictions_dict=all_predictions_dict,
                                      all_targets_dict=all_targets_dict, metrics_calculator=solver.metrics_calculator)

    @staticmethod
    def _do_bootstrapping(iterations: int,
                          sample_size: int,
                          confidence_level: float,
                          seq_ids: List[str],
                          all_predictions_dict: Dict,
                          all_targets_dict: Dict,
                          metrics_calculator: MetricsCalculator):
        """

        :param iterations: Number of iterations to perform bootstrapping
        :param sample_size: Sample size to use for bootstrapping. -1 defaults to all embeddings which is recommended.
                            It is possible, but not recommended to use a sample size larger or smaller
                            than the number of embeddings, because this might render the variance estimate unreliable.
                            See: https://math.mit.edu/~dav/05.dir/class24-prep-a.pdf (6.2)
        :param confidence_level: Confidence level for result error intervals (0.05 => 95% percentile)
        :param seq_ids: List of sequence IDs
        :param all_predictions_dict: Dictionary of all predictions
        :param all_targets_dict: Dictionary of all targets
        :param metrics_calculator: Metrics calculator object
        :return:
        """
        if sample_size == -1:
            sample_size = len(seq_ids)

        # Convert dictionaries to tensors
        all_predictions = torch.stack([all_predictions_dict[seq_id] for seq_id in seq_ids])
        all_targets = torch.stack([all_targets_dict[seq_id] for seq_id in seq_ids])

        # Set random seed
        seed = np.random.get_state()[1][0] if np.random.get_state() else 42
        rng = np.random.RandomState(seed)

        # Generate all random indices at once
        all_indices = rng.choice(len(seq_ids), size=(iterations, sample_size), replace=True)

        iteration_results = []
        for indices in all_indices:
            # Use integer indexing instead of string keys
            sampled_predictions = all_predictions[indices]
            sampled_targets = all_targets[indices]

            iteration_result = metrics_calculator.compute_metrics(
                predicted=sampled_predictions,
                labels=sampled_targets
            )
            iteration_results.append(iteration_result)

        # Process results
        metrics = list(iteration_results[0].keys())
        result_dict = {}
        for metric in metrics:
            all_metric_values = torch.tensor([res[metric] for res in iteration_results], dtype=torch.float16)
            mean, lower_bound, upper_bound = get_mean_and_confidence_bounds(
                values=all_metric_values,
                dimension=0,
                confidence_level=confidence_level
            )
            result_dict[metric] = {"mean": mean.item(), "lower": lower_bound.item(), "upper": upper_bound.item()}

        return result_dict

    def from_embeddings_with_monte_carlo_dropout(self, embeddings: Union[Iterable, Dict],
                                                 split_name: str = "hold_out",
                                                 n_forward_passes: int = 30,
                                                 confidence_level: float = 0.05,
                                                 seed: int = 42) -> Dict:
        """
        Calculate predictions by using Monte Carlo dropout.
        Only works if the model has at least one dropout layer employed.
        Method to quantify the uncertainty within the model.

        :param embeddings: Iterable or dictionary containing the input embeddings to predict on.
        :param split_name: Name of the split to use for prediction. Default is "hold_out".
        :param n_forward_passes: Number of times to repeat the prediction calculation
                                with different dropout nodes enabled. Must be > 1.
        :param confidence_level: Confidence level for the result confidence intervals. Default is 0.05,
                                which corresponds to a 95% percentile.
        :param seed: Seed to use for the dropout predictions

        :return: Dictionary containing with keys that will either be taken from the embeddings dict or
         represent the indexes if embeddings are given as a list. Contains the following values for each key:
                 - 'prediction': Class or value prediction based on the mean over `n_forward_passes` forward passes.
                 - 'all_predictions': All raw predictions, i.e. one prediction per forward pass.
                 - 'mcd_mean': Average over `n_forward_passes` forward passes for each class.
                 - 'mcd_lower_bound': Lower bound of the confidence interval using a normal distribution with the given
                                      confidence level.
                 - 'mcd_upper_bound': Upper bound of the confidence interval using a normal distribution with the given
                                      confidence level.
                 - 'confidence_range': Confidence range (== mcd_mean - mcd_lower_bound)
        """
        if n_forward_passes <= 1:
            raise ValueError(f"n_forward_passes must be > 1, given: {n_forward_passes}")

        if not 0 < confidence_level < 1:
            raise ValueError(f"Confidence level must be between 0 and 1, given: {confidence_level}!")

        # Necessary because dropout layer have a random part by design
        seed_all(seed)

        solver, dataloader = self._load_solver_and_dataloader(embeddings, split_name)

        predictions = solver.inference_monte_carlo_dropout(dataloader=dataloader,
                                                           n_forward_passes=n_forward_passes,
                                                           confidence_level=confidence_level)["mapped_predictions"]

        # For class predictions, revert from int (model output) to str (class name)
        if self.protocol in Protocol.per_residue_protocols():
            for seq_id, prediction_list in predictions.items():
                for prediction_dict in prediction_list:
                    prediction_dict["prediction"] = list(revert_mappings(protocol=self.protocol,
                                                                         test_predictions={
                                                                             seq_id: prediction_dict["prediction"]
                                                                         },
                                                                         class_int2str=self.class_int2str).values())[0]
        else:
            for seq_id, prediction_dict in predictions.items():
                prediction_dict["prediction"] = list(revert_mappings(protocol=self.protocol,
                                                                     test_predictions={
                                                                         seq_id: prediction_dict["prediction"]
                                                                     },
                                                                     class_int2str=self.class_int2str).values())[0]

        return predictions

    @staticmethod
    def from_onnx_with_embeddings(model_path: str, embeddings: Union[Iterable, Dict],
                                  protocol: Optional[Protocol] = None):
        try:
            import onnxruntime
        except ImportError:
            raise Exception("No onnxruntime in current environment found! Please install one first, e.g. "
                            "uv pip install -e '.[onnx-cpu]'!")

        if isinstance(embeddings, Dict):
            embeddings_dict = embeddings
        else:
            embeddings_dict = {str(idx): embedding for idx, embedding in enumerate(embeddings)}

        onnx_model = onnx.load(model_path)
        onnx.checker.check_model(onnx_model)

        ep_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        ort_session = onnxruntime.InferenceSession(model_path, providers=ep_list)

        result_dict = {}
        for seq_id, embedding in embeddings_dict.items():
            input_feed = {ort_session.get_inputs()[0].name: np.expand_dims(embedding, axis=0)}
            ort_outs = ort_session.run(None, input_feed=input_feed)
            if protocol is not None and protocol in Protocol.classification_protocols():
                result_dict[seq_id] = torch.softmax(torch.tensor(ort_outs[0][0]), dim=0).tolist()
            else:
                result_dict[seq_id] = ort_outs[0][0].tolist()[0]
        return result_dict
