from typing import Any, Callable, Dict, List

import mlflow
import numpy as np
import tensorflow as tf
from tensorflow.keras.callbacks import Callback
from tqdm import tqdm

from .metrics import compute_retrieval_metrics, compute_retrieval_metrics_numpy


class TfMLflowEpochLogger(Callback):
    """A Keras Callback to log metrics to MLflow at the end of each epoch or evaluation."""

    def __init__(
        self,
        enable_mlflow: bool = False,
        experiment_name: str | None = None,
        run_name: str = "Training_Run",
        train_df: Any | None = None,
        test_df: Any | None = None,
        log_params: Dict[str, Any] | None = None,
    ):
        """Initialize."""
        super().__init__()
        self.enable_mlflow = enable_mlflow
        self.experiment_name = experiment_name
        self.run_name = run_name
        self.train_df = train_df
        self.test_df = test_df
        self.log_params = log_params if log_params is not None else {}

    def on_train_begin(self, logs: Dict[str, Any] | None = None) -> None:
        if self.enable_mlflow:
            if self.experiment_name:
                if not mlflow.get_experiment_by_name(self.experiment_name):
                    mlflow.create_experiment(self.experiment_name)
                mlflow.set_experiment(self.experiment_name)

            print("Starting training with MLflow tracking...")
            # End any existing active run to avoid nested run conflicts
            if mlflow.active_run():
                mlflow.end_run()

            mlflow.start_run(run_name=self.run_name)
            if self.log_params:
                mlflow.log_params(self.log_params)

            # Log Datasets
            try:
                import pandas as pd

                if self.train_df is not None and isinstance(self.train_df, pd.DataFrame):
                    dataset_train = mlflow.data.from_pandas(self.train_df, name="training_dataset")
                    mlflow.log_input(dataset_train, context="training")
                    print(" -> MLflow logged training dataset")

                if self.test_df is not None and isinstance(self.test_df, pd.DataFrame):
                    dataset_test = mlflow.data.from_pandas(self.test_df, name="testing_dataset")
                    mlflow.log_input(dataset_test, context="testing")
                    print(" -> MLflow logged testing dataset")
            except Exception as e:
                print(f" -> MLflow dataset logging failed: {e}")

    def on_train_end(self, logs: Dict[str, Any] | None = None) -> None:
        if self.enable_mlflow:
            mlflow.end_run()
            print(" -> MLflow tracking run ended.")

    def on_epoch_end(self, epoch: int, logs: Dict[str, Any] | None = None) -> None:
        if self.enable_mlflow and logs is not None:
            mlflow.log_metrics(logs, step=epoch)
            print(f" -> MLflow logged metrics for epoch {epoch + 1}")

    # def on_test_end(self, logs: Optional[Dict[str, Any]] = None) -> None:
    # on_test_end can be called during model.fit for validation, or independently model.evaluate
    # we will log if mlflow is enabled
    # if self.enable_mlflow and logs is not None:
    #     metrics = {f"validation_{k}": v for k, v in logs.items()}
    #     mlflow.log_metrics(metrics)
    #     print(" -> MLflow logged test/validation metrics")


class NumpyMLflowLogger:
    """A generic MLflow Logger to manually log metrics during training loops in PyTorch, Scikit-learn, etc."""

    def __init__(
        self,
        enable_mlflow: bool = False,
        experiment_name: str | None = None,
        run_name: str = "Training_Run_Numpy",
        train_df: Any | None = None,
        test_df: Any | None = None,
        log_params: Dict[str, Any] | None = None,
    ):
        """Initialize."""
        self.enable_mlflow = enable_mlflow
        self.experiment_name = experiment_name
        self.run_name = run_name
        self.train_df = train_df
        self.test_df = test_df
        self.log_params = log_params if log_params is not None else {}

    def on_train_begin(self) -> None:
        if self.enable_mlflow:
            if self.experiment_name:
                if not mlflow.get_experiment_by_name(self.experiment_name):
                    mlflow.create_experiment(self.experiment_name)
                mlflow.set_experiment(self.experiment_name)

            print("Starting generic generic training with MLflow tracking...")
            if mlflow.active_run():
                mlflow.end_run()

            mlflow.start_run(run_name=self.run_name)
            if self.log_params:
                mlflow.log_params(self.log_params)

            # Log Datasets
            try:
                import pandas as pd

                if self.train_df is not None and isinstance(self.train_df, pd.DataFrame):
                    dataset_train = mlflow.data.from_pandas(self.train_df, name="training_dataset")
                    mlflow.log_input(dataset_train, context="training")
                    print(" -> MLflow logged training dataset")

                if self.test_df is not None and isinstance(self.test_df, pd.DataFrame):
                    dataset_test = mlflow.data.from_pandas(self.test_df, name="testing_dataset")
                    mlflow.log_input(dataset_test, context="testing")
                    print(" -> MLflow logged testing dataset")
            except Exception as e:
                print(f" -> MLflow dataset logging failed: {e}")

    def on_train_end(self) -> None:
        if self.enable_mlflow:
            mlflow.end_run()
            print(" -> MLflow tracking run ended.")

    def on_epoch_end(self, epoch: int, metrics: Dict[str, float]) -> None:
        if self.enable_mlflow and metrics is not None:
            mlflow.log_metrics(metrics, step=epoch)
            print(f" -> MLflow logged metrics for epoch {epoch + 1}")

    # def on_test_end(self, metrics: Dict[str, float]) -> None:
    # if self.enable_mlflow and metrics is not None:
    #     prefixed_metrics = {f"validation_{k}": v for k, v in metrics.items()}
    #     mlflow.log_metrics(prefixed_metrics)
    #     print(" -> MLflow logged test/validation metrics")


def evaluate_model_tf(  # noqa: C901
    model: tf.keras.Model,
    dataset: tf.data.Dataset,
    k: int | List[int] = 5,
    enable_mlflow: bool = False,
    experiment_name: str | None = None,
    run_name: str = "Final_Evaluation",
    log_params: Dict[str, Any] | None = None,
) -> Dict[str, float]:
    """Evaluation loop that processes the dataset and logs to MLflow.

    Args:
        model (tf.keras.Model): Trained Keras model.
        dataset (tf.data.Dataset): Dataset for evaluation.
        k (Union[int, List[int]], optional): Top-K metrics to compute. Defaults to 5.
        enable_mlflow (bool, optional): Whether to track evaluation metrics. Defaults to False.
        experiment_name (str, optional): MLflow experiment name. Defaults to None.
        run_name (str, optional): MLflow run name. Defaults to "Final_Evaluation".
        log_params (Dict[str, Any], optional): Optional dict of params to log (e.g., {"epochs": 10}). Defaults to None.

    Returns:
        Dict[str, float]: Averaged evaluation metrics.
    """
    if enable_mlflow:
        if experiment_name:
            if not mlflow.get_experiment_by_name(experiment_name):
                mlflow.create_experiment(experiment_name)
            mlflow.set_experiment(experiment_name)

        print("Starting evaluation with MLflow tracking...")
        # Since we use context manager, everything underneath falls in run tracking
        mlflow.start_run(run_name=run_name)
        if log_params:
            mlflow.log_params(log_params)
    else:
        print("Starting evaluation WITHOUT MLflow tracking...")

    try:
        total_metrics = {"loss": 0.0}
        num_batches = 0

        for batch in tqdm(dataset, desc="Evaluating"):
            logits = model(batch, training=False)
            batch_size = tf.shape(logits)[0]
            labels = tf.eye(batch_size)

            loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)(labels, logits)

            metrics = compute_retrieval_metrics(logits, k=k)

            # Initialize metric keys on first batch
            if num_batches == 0:
                for key in metrics:
                    total_metrics[key] = 0.0

            for key, val in metrics.items():
                total_metrics[key] += val.numpy()

            total_metrics["loss"] += loss.numpy()
            num_batches += 1

        avg_metrics = {key: float(val / num_batches) for key, val in total_metrics.items()}
        if enable_mlflow:
            mlflow.log_metrics(avg_metrics)

        print("\n--- Final Evaluation Metrics ---")
        for key, val in avg_metrics.items():
            print(f"{key}: {val:.4f}")

        return avg_metrics
    finally:
        if enable_mlflow:
            mlflow.end_run()


def evaluate_model_numpy(  # noqa: C901
    predict_fn: Callable[[Any], np.ndarray],
    dataset: Any,
    k: int | List[int] = 5,
    enable_mlflow: bool = False,
    experiment_name: str | None = None,
    run_name: str = "Final_Evaluation_Numpy",
    log_params: Dict[str, Any] | None = None,
) -> Dict[str, float]:
    """Evaluation loop that processes the dataset and logs to MLflow using purely NumPy metrics.

    Framework-agnostic.

    Args:
        predict_fn (Callable): A function that takes a batch from the dataset and returns a NumPy array of logits (Batch Size, Batch Size).
        dataset (Any): An iterable dataset. Each iteration should yield a batch that is valid input for predict_fn.
        k (Union[int, List[int]], optional): Top-K metrics to compute. Defaults to 5.
        enable_mlflow (bool, optional): Whether to track evaluation metrics. Defaults to False.
        experiment_name (str, optional): MLflow experiment name. Defaults to None.
        run_name (str, optional): MLflow run name. Defaults to "Final_Evaluation_Numpy".
        log_params (Dict[str, Any], optional): Optional dict of params to log (e.g., {"epochs": 10}). Defaults to None.

    Returns:
        Dict[str, float]: Averaged evaluation metrics across all batches.
    """
    if enable_mlflow:
        if experiment_name:
            if not mlflow.get_experiment_by_name(experiment_name):
                mlflow.create_experiment(experiment_name)
            mlflow.set_experiment(experiment_name)

        print("Starting NumPy evaluation with MLflow tracking...")
        mlflow.start_run(run_name=run_name)
        if log_params:
            mlflow.log_params(log_params)
    else:
        print("Starting NumPy evaluation WITHOUT MLflow tracking...")

    try:
        total_metrics = {}
        num_batches = 0

        for batch in tqdm(dataset, desc="Evaluating (NumPy)"):
            logits = predict_fn(batch)
            metrics = compute_retrieval_metrics_numpy(logits, k=k)

            # Initialize metric keys on first batch
            if num_batches == 0:
                for key in metrics:
                    total_metrics[key] = 0.0

            for key, val in metrics.items():
                total_metrics[key] += val

            num_batches += 1

        avg_metrics = {key: float(val / num_batches) for key, val in total_metrics.items()}

        if enable_mlflow:
            mlflow.log_metrics(avg_metrics)

        print("\n--- Final NumPy Evaluation Metrics ---")
        for key, val in avg_metrics.items():
            print(f"{key}: {val:.4f}")

        return avg_metrics
    finally:
        if enable_mlflow:
            mlflow.end_run()
