import numpy as np
import pandas as pd

from sklearn.base import BaseEstimator
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.linear_model import Ridge, RidgeCV
from sklearn.multioutput import MultiOutputClassifier

from mrna_bench.data_splitter.data_splitter import DataSplitter
from mrna_bench.linear_probe.evaluator import LinearProbeEvaluator
from mrna_bench.linear_probe.persister import LinearProbePersister


class LinearProbe:
    """Linear Probe Module.

    Performs linear probing on embeddings for sequences from a benchmark
    dataset generated by an embedding model. The linear probing task can
    be selected at initialization on a target column from the benchmarking
    dataset dataframe.

    The train and test splits for linear probing can be recomputed from random
    seeds. This module supports generating splits based on sequence gene name
    homology from homologene.

    Trained linear probe models for a specific data split can be accessed
    through the models attribute. The linear probe results can be persisted
    to a disk.
    """

    linear_models = {
        "regression": RidgeCV(alphas=[1e-3, 1e-2, 1e-1, 1, 10]),
        "reg_lin": LinearRegression(n_jobs=-1),
        "reg_ridge": RidgeCV(alphas=[1e-3, 1e-2, 1e-1, 1, 10]),
        "classification": LogisticRegression(max_iter=5000, n_jobs=-1),
        "multilabel": MultiOutputClassifier(
            LogisticRegression(max_iter=5000, n_jobs=-1),
            n_jobs=-1
        )
    }

    def __init__(
        self,
        data_df: pd.DataFrame,
        embeddings: np.ndarray,
        target_col: str,
        task: str,
        splitter: DataSplitter,
        evaluator: LinearProbeEvaluator,
        eval_all_splits: bool,
        persister: LinearProbePersister | None = None
    ):
        """Initialize LinearProbe.

        Args:
            dataset: Dataframe containing sequences to be probed. Dataframe
                should contain target columns and sequences should have the
                same order as the array of embeddings.
            model_short_name: Name of model used to generate embedding.
            embeddings: Embeddings of dataset sequences to be probed.
            target_col: Target column in dataframe for linear probing task.
            task: Linear probing task to perform.
            evaluator: Evaluator object for linear probing results.
            eval_all_splits: Evaluate metrics on all splits. Only evaluates
                validation split otherwise.
            persister: Persister object for linear probing results.
                If not provided, results will not be saved.
        """
        if task not in self.linear_models.keys():
            raise ValueError("Invalid task name.")

        if target_col not in data_df.columns:
            raise ValueError("Target column not found in dataframe.")

        self.data_df = data_df
        self.data_df["embeddings"] = list(embeddings)

        self.task = task
        self.target_col = target_col

        self.splitter = splitter
        self.persister = persister
        self.evaluator = evaluator
        self.eval_all_splits = eval_all_splits

        self.models: dict[int, BaseEstimator] = {}

    def get_df_splits(
        self,
        random_seed: int,
        dropna: bool = True
    ) -> dict[str, np.ndarray]:
        """Get train, validation, and test splits for data and labels.

        Args:
            random_seed: Random seed used for generating data splits.
            dropna: Drop rows with NaN values in target column.

        Returns:
            Dictionary containing per-split data and labels as numpy arrays.
        """
        data_df_copy = self.data_df.copy()

        if dropna:
            data_df_copy = data_df_copy.dropna(subset=[self.target_col])

        train_df, val_df, test_df = self.splitter.get_all_splits_df(
            data_df_copy,
            random_seed=random_seed
        )

        splits = {
            "train_X": np.array(train_df["embeddings"].tolist()),
            "train_y": train_df[self.target_col].to_numpy(),
        }

        if not val_df.empty:
            splits["val_X"] = np.array(val_df["embeddings"].tolist())
            splits["val_y"] = val_df[self.target_col].to_numpy()

        if not test_df.empty:
            splits["test_X"] = np.array(test_df["embeddings"].tolist())
            splits["test_y"] = test_df[self.target_col].to_numpy()

        if isinstance(splits["train_y"][0], np.ndarray):
            splits["train_y"] = np.vstack(list(splits["train_y"]))

            if "val_y" in splits:
                splits["val_y"] = np.vstack(list(splits["val_y"]))
            if "test_y" in splits:
                splits["test_y"] = np.vstack(list(splits["test_y"]))

        return splits

    def run_linear_probe(
        self,
        random_seed: int = 2541,
        persist: bool = False,
        dropna: bool = True,
    ) -> dict[str, float]:
        """Perform data split and run linear probe.

        Args:
            random_seed: Random seed used for data split.
            persist: Save results to disk.
            dropna: Drop rows with NaN values in target column

        Returns:
            Dictionary of linear probing metrics per split.
        """
        model = self.linear_models[self.task]
        splits = self.get_df_splits(random_seed, dropna)

        np.random.seed(random_seed)

        try:
            model.fit(splits["train_X"], splits["train_y"])
        except ValueError:
            if self.task in ["regression", "reg_ridge"]:
                model = Ridge(solver="sag", alpha=1e-3)
                model.fit(splits["train_X"], splits["train_y"])
            else:
                raise

        self.models[random_seed] = model

        if self.eval_all_splits:
            eval_splits = splits
        else:
            eval_splits = {
                "val_X": splits["val_X"],
                "val_y": splits["val_y"]
            }

        metrics = self.evaluator.evaluate_linear_probe(model, eval_splits)

        if self.persister and persist:
            self.persister.persist_run_results(metrics, random_seed)
        elif persist and not self.persister:
            raise RuntimeError("Must provide persister to persist results.")

        return metrics

    def linear_probe_multirun(
        self,
        random_seeds: list[int],
        persist: bool = False
    ) -> dict[int, dict[str, float]]:
        """Run multiple linear probes with distinct data split randomization.

        Args:
            random_seeds: Trandom seeds used per individual linear probe run.
            persist: Save results to data directory.

        Returns:
            Dictionary of metrics per random seed used to generate data splits
            for each individual linear probing run.
        """
        metrics = {}

        for random_seed in random_seeds:
            metric = self.run_linear_probe(random_seed, persist)
            metrics[random_seed] = metric

        return metrics

    def compute_multirun_results(
        self,
        metrics: dict[int, dict[str, float]],
        print_output: bool = False,
        ci_multiplier: float = 1.96,
        persist: bool = False
    ) -> dict[str, str]:
        """Aggregate multi-run linear probing results.

        Prints mean metric and confidence interval at desired level.

        Args:
            metrics: Result of multi-run linear probing.
            print_output: Prints aggregated metrics.
            ci_multiplier: Constant multiplied to standard error to get CI.
            persist: Save results to data directory.

        Returns:
            Dictionary of mean metric and CI per data split.
        """
        metric_vals: dict[str, list[float]] = {}

        for metric_dict in metrics.values():
            for metric_name, metric_val in metric_dict.items():
                metric_vals.setdefault(metric_name, []).append(metric_val)

        metric_mean = {k: np.mean(v) for k, v in metric_vals.items()}
        metric_std = {k: np.std(v) for k, v in metric_vals.items()}

        metric_out = {}

        for k in metric_vals.keys():
            se = ci_multiplier * (metric_std[k] / (np.sqrt(len(metrics))))
            metric_out[k] = "{} ± {}".format(metric_mean[k], se)

            if print_output:
                print("{} ± {}".format(metric_mean[k], se))

        if self.persister and persist:
            self.persister.persist_run_results(metric_out, "all")
        elif persist and not self.persister:
            raise RuntimeError("Must provide persister to persist results.")

        return metric_out

    def get_fit_model(self, random_seed: int) -> BaseEstimator:
        """Get model trained on specific data split.

        Args:
            random_seed: Random seed used to generate data split.

        Returns:
            Trained model for specific data split.
        """
        return self.models[random_seed]
