"""WandB sweep results aggregation and visualization for scIB benchmarking."""

from dataclasses import fields
from pathlib import Path
from tempfile import TemporaryDirectory

import numpy as np
import pandas as pd
import scanpy as sc
import wandb
from scib_metrics.benchmark import BatchCorrection, Benchmarker, BioConservation
from scib_metrics.benchmark._core import metric_name_cleaner
from tqdm import tqdm

from scembed.logging import logger
from scembed.utils import _download_artifact_by_run_id


class scIBAggregator:
    """Aggregator for WandB sweep results with scIB metrics visualization.

    Retrieves runs from a WandB project, organizes them by integration method,
    and provides visualization capabilities using scIB metrics formatting
    :cite:`luecken2022benchmarking`.

    Parameters
    ----------
    entity
        WandB entity name.
    project
        WandB project name.
    output_dir
        Directory for saving downloaded models and embeddings. If None, creates temporary directory.
    """

    def __init__(self, entity: str, project: str, output_dir: str | Path | None = None):
        self.entity = entity
        self.project = project

        # Setup output directories
        self._temp_dir = None
        if output_dir is None:
            self._temp_dir = TemporaryDirectory()
            output_dir = Path(self._temp_dir.name)
        elif isinstance(output_dir, str):
            output_dir = Path(output_dir)

        self.output_dir = output_dir
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Raw data storage
        self.raw_df: pd.DataFrame | None = None

        # Processed data organized by method
        self.method_data: dict[str, dict[str, pd.DataFrame | Benchmarker]] = {}

        # Data quality tracking
        self.n_runs_fetched = 0
        self.n_runs_filtered = 0
        self.missing_config_runs: list[str] = []
        self.missing_metrics_runs: list[str] = []

        # Available metrics in the data
        self.available_scib_metrics: set[str] = set()

        # Mapping of metric display names to their types
        self.metric_to_type: dict[str, str] = self._build_metric_type_mapping()

        # Aggregated results (best run per method)
        self.results: dict[str, pd.DataFrame | Benchmarker] | None = None

        logger.info("Initialized scIBAggregator for %s/%s", entity, project)

    def fetch_runs(self) -> None:
        """Fetch runs from WandB and process into internal storage."""
        logger.info("Fetching runs from %s/%s...", self.entity, self.project)

        # Reset state for fresh fetch
        self.raw_df = None
        self.method_data = {}
        self.n_runs_fetched = 0
        self.n_runs_filtered = 0
        self.missing_config_runs = []
        self.missing_metrics_runs = []
        self.available_scib_metrics = set()

        # Initialize WandB API
        api = wandb.Api()

        # Fetch all runs for the project
        runs = api.runs(f"{self.entity}/{self.project}")

        # Extract parameter configurations and final results
        data = []
        for run in tqdm(runs, desc="Processing runs"):
            run_data = {
                "run_id": run.id,
                "name": run.name,
                **dict(run.summary),
            }
            if "config" not in run.config.keys():
                run_data["config"] = run.config
            else:
                run_data.update(**run.config)
            data.append(run_data)

        # Convert to DataFrame
        self.raw_df = pd.DataFrame(data)
        self.n_runs_fetched = len(self.raw_df)

        logger.info("Fetched %d runs", self.n_runs_fetched)

        # Process the data
        self._process_runs()

        logger.info(
            "Processing complete: %d methods, %d runs filtered out", len(self.method_data), self.n_runs_filtered
        )

    def _process_runs(self) -> None:
        """Process raw runs into method-organized data structures."""
        if self.raw_df is None:
            raise ValueError("No runs fetched. Call fetch_runs() first.")

        # Filter runs with missing configs
        valid_runs = []
        for idx, row in self.raw_df.iterrows():
            config = row.get("config")
            if config is None or not isinstance(config, dict) or "method" not in config:
                self.missing_config_runs.append(row["run_id"])
                continue
            valid_runs.append(idx)

        if self.missing_config_runs:
            logger.warning(
                "Filtered out %d runs with missing/invalid configs: %s",
                len(self.missing_config_runs),
                self.missing_config_runs[:5],
            )

        valid_df = self.raw_df.loc[valid_runs].copy()

        # Extract scIB metrics and validate
        scib_metrics_data = []
        other_logs_data = []

        # Define aggregate metrics to exclude from individual metrics
        aggregate_metrics = {"Batch correction", "Bio conservation", "Total"}

        for _, row in valid_df.iterrows():
            config = row["config"]
            method = config["method"]

            # Extract scIB metrics from the 'scib' column
            scib_data = row.get("scib", {})

            # Handle various types of scib data (dict, SummarySubDict, etc.)
            if scib_data is None:
                self.missing_metrics_runs.append(row["run_id"])
                continue

            # Convert to regular dict if it's a wandb SummarySubDict or similar
            if hasattr(scib_data, "keys") and hasattr(scib_data, "items"):
                scib_data = dict(scib_data)

            # Now check if it's empty or not a dict
            if not isinstance(scib_data, dict) or len(scib_data) == 0:
                self.missing_metrics_runs.append(row["run_id"])
                continue

            # Separate actual scIB metrics from aggregates
            actual_scib_metrics = {k: v for k, v in scib_data.items() if k not in aggregate_metrics}

            # Build other logs entry (everything else from the run)
            other_logs_entry = {
                "run_id": row["run_id"],
                "method": method,
                **{k: v for k, v in row.items() if k not in ["config", "scib", "run_id"]},
            }
            other_logs_data.append(other_logs_entry)

            # Build scIB metrics entry (only actual metrics, no aggregates)
            metrics_entry = {"run_id": row["run_id"], "method": method, **actual_scib_metrics}
            scib_metrics_data.append(metrics_entry)

            # Track available metrics (only actual metrics)
            self.available_scib_metrics.update(actual_scib_metrics.keys())

        if self.missing_metrics_runs:
            logger.warning(
                "Filtered out %d runs with missing scIB metrics: %s",
                len(self.missing_metrics_runs),
                self.missing_metrics_runs[:5],
            )

        # Update filter count
        self.n_runs_filtered = len(self.missing_config_runs) + len(self.missing_metrics_runs)

        # Create dataframes
        if not scib_metrics_data:
            logger.warning("No valid runs found with both config and scIB metrics")
            return

        full_metrics_df = pd.DataFrame(scib_metrics_data).set_index("run_id")
        full_other_logs_df = pd.DataFrame(other_logs_data).set_index("run_id")

        # Organize by method - create method-specific config DataFrames without NaN columns
        # Use runs that passed BOTH config and scIB metrics validation
        valid_run_ids = set(full_metrics_df.index)  # These runs have both valid configs AND scIB metrics

        methods = {row["config"]["method"] for _, row in valid_df.iterrows() if row["run_id"] in valid_run_ids}
        for method in methods:
            # Get runs for this method that have BOTH valid configs AND scIB metrics
            method_rows = [
                row
                for _, row in valid_df.iterrows()
                if row["config"]["method"] == method and row["run_id"] in valid_run_ids
            ]
            method_run_ids = [row["run_id"] for row in method_rows]

            # Create method-specific config DataFrame preserving original structure
            method_configs = []
            for row in method_rows:
                config_entry = {"run_id": row["run_id"], **row["config"]}
                method_configs.append(config_entry)

            method_config_df = pd.DataFrame(method_configs).set_index("run_id")
            method_metrics_df = full_metrics_df.loc[method_run_ids]
            method_other_logs_df = full_other_logs_df.loc[method_run_ids]

            # Create Benchmarker object for this method's scIB metrics
            try:
                benchmarker = self._create_benchmarker_for_method(method_metrics_df)
            except (ValueError, KeyError, ImportError) as e:
                logger.warning("Failed to create Benchmarker for method %s: %s", method, e)
                # Fallback to storing as DataFrame if Benchmarker creation fails
                benchmarker = method_metrics_df

            self.method_data[method] = {
                "configs": method_config_df,
                "scib_benchmarker": benchmarker,
                "other_logs": method_other_logs_df,
            }

        logger.info("Available scIB metrics: %s", sorted(self.available_scib_metrics))

    def _build_metric_type_mapping(self) -> dict[str, str]:
        """Build mapping from metric names to metric types using scIB dataclasses."""
        metric_to_type = {}

        for name, display_name in metric_name_cleaner.items():
            if any(name.startswith(field.name) for field in fields(BioConservation)):
                metric_to_type[display_name] = "Bio conservation"
            if any(name.startswith(field.name) for field in fields(BatchCorrection)):
                metric_to_type[display_name] = "Batch correction"

        return metric_to_type

    def _create_benchmarker_for_method(self, scib_df: pd.DataFrame) -> Benchmarker:
        """Create a Benchmarker object for scIB metrics results.

        Uses the scIB benchmarking framework :cite:`luecken2022benchmarking` to
        organize and visualize integration method evaluation results.

        Parameters
        ----------
        scib_df
            DataFrame with scIB metrics (runs as rows, metrics as columns)

        Returns
        -------
        Benchmarker
            Benchmarker object with results loaded
        """
        # Transpose so metrics are rows and runs are columns (as expected by scIB)
        scib_plot = scib_df.drop(columns=["method"]).T.copy()

        # Add metric type information
        scib_plot["Metric Type"] = scib_plot.index.map(self.metric_to_type)

        # Filter to only metrics that have type mappings
        valid_metrics = [metric for metric in scib_plot.index if metric in self.metric_to_type]
        scib_plot = scib_plot.loc[valid_metrics]

        # Reorder rows based on the order in self.metric_to_type
        ordered_metrics = [metric for metric in self.metric_to_type.keys() if metric in valid_metrics]
        scib_plot = scib_plot.loc[ordered_metrics]

        # Create minimal dummy AnnData for Benchmarker initialization
        n_obs = 10
        n_vars = 5
        dummy_adata = sc.AnnData(np.random.random((n_obs, n_vars)))
        dummy_adata.obs["batch"] = ["batch1", "batch2"] * (n_obs // 2)
        dummy_adata.obs["celltype"] = ["type1", "type2"] * (n_obs // 2)

        # Create Benchmarker instance

        bm = Benchmarker(
            adata=dummy_adata,
            batch_key="batch",
            label_key="celltype",
            embedding_obsm_keys=list(scib_plot.columns[:-1]),  # Exclude 'Metric Type' column
            bio_conservation_metrics=BioConservation(),
            batch_correction_metrics=BatchCorrection(),
            n_jobs=-1,
        )

        # Inject the results directly (skip actual benchmarking)
        bm._results = scib_plot
        bm._benchmarked = True  # Mark as benchmarked to avoid re-running

        return bm

    def aggregate(self, sort_by: str = "Total", min_max_scale_for_selection: bool = True) -> None:
        """
        Aggregate best run per method into unified results structure.

        Creates self.results with the same structure as individual method results:
        - configs: DataFrame with best configurations per method
        - scib_benchmarker: Benchmarker object with all best runs
        - other_logs: DataFrame with other logs from best runs per method

        Parameters
        ----------
        sort_by
            Metric to sort by for selecting best run per method.
            Default is "Total" (overall scIB score).
        min_max_scale_for_selection
            Whether to use min-max scaled metrics for selecting the best run per method.
            If True, uses scaled metrics for fair comparison across different metric ranges.
            If False, uses raw metrics for selection.
            Default is True. Note: reported metrics are always raw (non-scaled) values.
        """
        if not self.method_data:
            raise ValueError("No data available. Call fetch_runs() first.")

        best_configs = []
        best_metrics = []
        best_other_logs = []

        for method, data in self.method_data.items():
            benchmarker = data["scib_benchmarker"]
            configs_df = data["configs"]
            other_logs_df = data["other_logs"]

            # Get metrics DataFrame from Benchmarker for SELECTION (user choice of scaling)
            selection_results_df = benchmarker.get_results(min_max_scale=min_max_scale_for_selection)
            # Remove the "Metric Type" row if present
            if "Metric Type" in selection_results_df.index:
                selection_metrics_df = selection_results_df.drop(index="Metric Type")
            else:
                selection_metrics_df = selection_results_df

            if sort_by not in selection_metrics_df.columns:
                logger.warning("Metric '%s' not found for method '%s', skipping", sort_by, method)
                continue

            # Find best run using the selection metrics (scaled or unscaled as requested)
            best_idx = selection_metrics_df[sort_by].idxmax()

            # Get metrics DataFrame for REPORTING (always unscaled)
            reporting_results_df = benchmarker.get_results(min_max_scale=False)
            # Remove the "Metric Type" row if present
            if "Metric Type" in reporting_results_df.index:
                reporting_metrics_df = reporting_results_df.drop(index="Metric Type")
            else:
                reporting_metrics_df = reporting_results_df

            # Collect best run data
            best_config = configs_df.loc[best_idx].copy()
            best_config["run_id"] = best_idx  # Add run_id from the index
            best_config.name = method  # Use method name as index
            best_configs.append(best_config)

            # Extract metrics for best run using REPORTING metrics (always unscaled)
            best_metric = reporting_metrics_df.loc[best_idx].copy()
            best_metric.name = method
            best_metrics.append(best_metric)

            # Collect other logs for best run
            best_other = other_logs_df.loc[best_idx].copy()
            best_other.name = method
            best_other_logs.append(best_other)

        if not best_configs:
            logger.warning("No valid runs found with metric '%s'", sort_by)
            self.results = None
            return

        # Create aggregated DataFrames
        configs_df = pd.DataFrame(best_configs)

        # Remove redundant 'method' column (since DataFrame is indexed by method)
        if "method" in configs_df.columns:
            configs_df = configs_df.drop(columns=["method"])

        # Ensure 'run_id' is the first column
        if "run_id" in configs_df.columns:
            cols = ["run_id"] + [col for col in configs_df.columns if col != "run_id"]
            configs_df = configs_df[cols]

        other_logs_df = pd.DataFrame(best_other_logs)

        # Create metrics DataFrame and Benchmarker
        metrics_df = pd.DataFrame(best_metrics)
        metrics_df["method"] = metrics_df.index

        # Create Benchmarker object for aggregated results
        try:
            benchmarker = self._create_benchmarker_for_method(metrics_df)
        except (ValueError, KeyError, ImportError) as e:
            logger.warning("Failed to create aggregated Benchmarker: %s", e)
            # Fallback to storing as DataFrame
            benchmarker = metrics_df

        # Store results in the same format as individual methods
        self.results = {
            "configs": configs_df,
            "scib_benchmarker": benchmarker,
            "other_logs": other_logs_df,
        }

        logger.info("Aggregated best runs for %d methods using metric '%s'", len(best_configs), sort_by)

    def get_method_runs(
        self, method: str, sort_by: str = "Total", min_max_scale: bool = True
    ) -> dict[str, pd.DataFrame | Benchmarker]:
        """
        Get all data for a specific method, optionally sorted by a metric.

        Parameters
        ----------
        method
            Method name to retrieve data for.
        sort_by
            Metric name to sort runs by. Default is "Total".
            Must be a valid metric name (e.g., 'Total', 'Silhouette label', etc.).
        min_max_scale
            Whether to use min-max scaled metrics for sorting. Default is True.
            If True, uses scaled metrics for fair comparison across different metric ranges.
            If False, uses raw metrics for sorting.

        Returns
        -------
        dict
            Dictionary with 'configs' (DataFrame), 'scib_benchmarker' (Benchmarker),
            and 'other_logs' (DataFrame). The 'configs' and 'other_logs' DataFrames
            are sorted by the specified metric in descending order (best runs first).
        """
        if method not in self.method_data:
            raise ValueError(f"Method '{method}' not found. Available methods: {list(self.method_data.keys())}")

        method_data = self.method_data[method].copy()

        # Get benchmarker and extract results for sorting
        benchmarker = method_data["scib_benchmarker"]

        # Get results from Benchmarker
        results_df = benchmarker.get_results(min_max_scale=min_max_scale)

        # Remove 'Metric Type' row if present
        if "Metric Type" in results_df.index:
            results_df = results_df.drop(index="Metric Type")

        if sort_by not in results_df.columns:
            available_metrics = list(results_df.columns)
            raise ValueError(f"Metric '{sort_by}' not found. Available metrics: {available_metrics}")

        # Sort by the specified metric (descending - best first)
        sorted_order = results_df[sort_by].sort_values(ascending=False).index

        # Apply sorting to configs and other_logs DataFrames
        method_data["configs"] = method_data["configs"].loc[sorted_order]
        method_data["other_logs"] = method_data["other_logs"].loc[sorted_order]

        return method_data

    def get_models_and_embeddings(
        self,
        model_artifact_name: str = "trained_model",
        embedding_artifact_name: str = "embedding",
    ) -> None:
        """
        Download models and embeddings for the best-performing run per method.

        Creates a folder structure in output_dir with models and embeddings organized by method.
        Models are only downloaded for methods that have them (GPU-based methods).

        Parameters
        ----------
        model_artifact_name
            Name of the model artifact in wandb.
        embedding_artifact_name
            Name of the embedding artifact in wandb.
        """
        if self.results is None:
            raise ValueError("Must call aggregate() first to determine best runs per method")

        configs_df = self.results["configs"]
        assert isinstance(configs_df, pd.DataFrame), "configs should always be a DataFrame"

        for method in configs_df.index:
            row = configs_df.loc[method]
            run_id = row["run_id"]

            logger.debug("Processing method '%s', run_id: %s", method, run_id)

            # Create method directory
            method_dir = self.output_dir / method
            method_dir.mkdir(parents=True, exist_ok=True)

            # Create subdirectories
            models_dir = method_dir / "models"
            embeddings_dir = method_dir / "embeddings"
            models_dir.mkdir(exist_ok=True)
            embeddings_dir.mkdir(exist_ok=True)

            # Track what was downloaded for summary
            downloaded_items = []
            missing_items = []

            # Try to download model (may not exist for CPU methods)
            model_path = _download_artifact_by_run_id(
                run_id=run_id,
                entity=self.entity,
                project=self.project,
                artifact_name=model_artifact_name,
                download_dir=models_dir,
            )

            if model_path is None:
                logger.debug("No model artifact found for method '%s'", method)
                missing_items.append("model")
            else:
                logger.debug("Downloaded model for method '%s'", method)
                downloaded_items.append("models")

            # Try to download embedding
            embedding_path = _download_artifact_by_run_id(
                run_id=run_id,
                entity=self.entity,
                project=self.project,
                artifact_name=embedding_artifact_name,
                download_dir=embeddings_dir,
            )

            if embedding_path is None:
                missing_items.append("embedding")
            else:
                downloaded_items.append("embeddings")

            # Summary logging per method
            if downloaded_items and missing_items:
                logger.info(
                    "Downloaded %s for method '%s' to %s/%s/",
                    " and ".join(downloaded_items),
                    method,
                    self.output_dir,
                    method,
                )
                logger.warning("Missing %s for method '%s' (run_id: %s)", " and ".join(missing_items), method, run_id)
            elif downloaded_items:
                logger.info(
                    "Downloaded %s for method '%s' to %s/%s/",
                    " and ".join(downloaded_items),
                    method,
                    self.output_dir,
                    method,
                )
            else:
                logger.warning("No artifacts found for method '%s' (run_id: %s)", method, run_id)

    @property
    def available_methods(self) -> list[str]:
        """List of available methods in the data."""
        return list(self.method_data.keys())

    def __repr__(self) -> str:
        """String representation with data summary."""
        if self.raw_df is None:
            return f"scIBAggregator({self.entity}/{self.project}, no data fetched)"

        n_methods = len(self.method_data)
        n_valid_runs = self.n_runs_fetched - self.n_runs_filtered
        n_metrics = len(self.available_scib_metrics)

        status_parts = []
        if self.missing_config_runs:
            status_parts.append(f"{len(self.missing_config_runs)} missing configs")
        if self.missing_metrics_runs:
            status_parts.append(f"{len(self.missing_metrics_runs)} missing metrics")
        if self.results is not None:
            status_parts.append("aggregated")

        status = f" ({', '.join(status_parts)})" if status_parts else ""

        return (
            f"scIBAggregator("
            f"{self.entity}/{self.project}, "
            f"{n_valid_runs}/{self.n_runs_fetched} runs, "
            f"{n_methods} methods, "
            f"{n_metrics} metrics, "
            f"output_dir='{self.output_dir}'"
            f"{status})"
        )
