"""Data utilities for Entropice dashboard."""

import json
import pickle
import subprocess
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Literal

import antimeridian
import geopandas as gpd
import pandas as pd
import streamlit as st
import xarray as xr
from shapely.geometry import shape

import entropice.spatial.grids
import entropice.utils.paths
from entropice.dashboard.utils.formatters import TrainingResultDisplayInfo
from entropice.ml.dataset import DatasetEnsemble, TrainingSet
from entropice.utils.training import Training
from entropice.utils.types import GridConfig, TargetDataset, Task, all_target_datasets, all_tasks


def _fix_hex_geometry(geom):
    """Fix hexagon geometry crossing the antimeridian."""
    try:
        return shape(antimeridian.fix_shape(geom))
    except ValueError as e:
        st.error(f"Error fixing geometry: {e}")
        return geom


@dataclass
class TrainingResult:
    """Wrapper for training result data and metadata."""

    path: Path
    experiment: str
    run: Training
    created_at: float
    available_metrics: list[str]
    files: list[Path]

    @classmethod
    def from_path(cls, result_path: Path, experiment_name: str | None = None) -> "TrainingResult":
        """Load a TrainingResult from a given result directory path."""
        all_files = list(result_path.iterdir())
        run = Training.load(result_path)

        created_at = result_path.stat().st_ctime
        available_metrics = [str(v) for v in run.metrics["metric"].unique()]

        return cls(
            path=result_path,
            experiment=experiment_name or "N/A",
            run=run,
            created_at=created_at,
            available_metrics=available_metrics,
            files=all_files,
        )

    @property
    def display_info(self) -> TrainingResultDisplayInfo:
        """Get display information for the training result."""
        return TrainingResultDisplayInfo(
            task=self.run.task,
            target=self.run.target,
            model=self.run.model_type,
            grid=self.run.dataset.grid,
            level=self.run.dataset.level,
            timestamp=datetime.fromtimestamp(self.created_at),
        )

    def load_best_model(self) -> object | None:
        """Load the best model from a training result."""
        model_file = self.path / "best_estimator_model.pkl"
        if not model_file.exists():
            return None

        try:
            with open(model_file, "rb") as f:
                model = pickle.load(f)
            return model
        except Exception as e:
            st.error(f"Error loading model: {e}")
            return None

    def load_model_state(self) -> xr.Dataset | None:
        """Load the model state from a training result."""
        state_file = self.path / "best_estimator_state.nc"
        if not state_file.exists():
            return None

        try:
            state = xr.open_dataset(state_file, engine="h5netcdf")
            return state
        except Exception as e:
            st.error(f"Error loading model state: {e}")
            return None

    def load_predictions(self) -> pd.DataFrame | None:
        """Load predictions from a training result."""
        preds_file = self.path / "predicted_probabilities.parquet"
        if not preds_file.exists():
            return None

        try:
            preds = pd.read_parquet(preds_file)
            return preds
        except Exception as e:
            st.error(f"Error loading predictions: {e}")
            return None

    def get_cv_results_dataframe(self) -> pd.DataFrame | None:
        """Get a DataFrame of available metrics for this training result."""
        results = self.run.cv_results
        if results is None:
            return None
        metric_cols = [col for col in results.columns if col.startswith("mean_test_")]
        if not metric_cols:
            return None
        metric_data = []
        for col in metric_cols:
            metric_name = col.replace("mean_test_", "").replace("neg_", "").title()
            metrics = results[col]
            # Check if the metric is negative
            if col.startswith("mean_test_neg_"):
                task_multiplier = 1 if self.run.task != "density" else 100
                task_multiplier *= -1
                metrics = metrics * task_multiplier

            metric_data.append(
                {
                    "Metric": metric_name,
                    "Best": f"{metrics.max():.4f}",
                    "Mean": f"{metrics.mean():.4f}",
                    "Std": f"{metrics.std():.4f}",
                    "Worst": f"{metrics.min():.4f}",
                }
            )

        return pd.DataFrame(metric_data)

    def _get_best_metric_name(self) -> str:
        """Get the primary metric name for a given task."""
        match self.run.task:
            case "binary":
                return "f1"
            case "count_regimes" | "density_regimes":
                return "f1_weighted"
            case _:  # regression tasks
                return "r2"

    def _get_best_score(self, split: Literal["train", "test", "complete"]) -> float:
        """Get the best score for the primary metric on a given split."""
        metric = self._get_best_metric_name()
        scores = self.run.metrics[(self.run.metrics["metric"] == metric) & (self.run.metrics["split"] == split)]
        # Should leave a single row
        assert len(scores) == 1, (
            f"Expected exactly one score for metric {metric} and split {split}, but found {len(scores)}: {scores}"
        )
        return float(scores["score"].iloc[0])

    @staticmethod
    def to_dataframe(training_results: list["TrainingResult"]) -> pd.DataFrame:
        """Convert a list of TrainingResult objects to a DataFrame for display."""
        records = []
        for tr in training_results:
            info = tr.display_info
            best_metric_name = tr._get_best_metric_name()
            best_train_score = tr._get_best_score("train")
            best_test_score = tr._get_best_score("test")

            record = {
                "Experiment": tr.experiment if tr.experiment else "N/A",
                "Task": info.task,
                "Target": info.target,
                "Model": info.model,
                "Grid": GridConfig.from_grid_level((info.grid, info.level)).display_name,
                "Created At": info.timestamp.strftime("%Y-%m-%d %H:%M"),
                "Score-Metric": best_metric_name.title(),
                "Best Models Score (Train-Set)": best_train_score,
                "Best Models Score (Test-Set)": best_test_score,
                "Trials": tr.run.n_trials or "N/A",
                "Path": str(tr.path.name),
            }
            records.append(record)
        return pd.DataFrame.from_records(records)

    @staticmethod
    def calculate_inference_maps(training_results: list["TrainingResult"]) -> gpd.GeoDataFrame:
        """Calculate the mean and standard deviation of inference maps across multiple training results."""
        assert len({tr.run.dataset.grid for tr in training_results}) == 1, (
            "All training results must have the same grid"
        )
        assert len({tr.run.dataset.level for tr in training_results}) == 1, (
            "All training results must have the same level"
        )

        grid = training_results[0].run.dataset.grid
        level = training_results[0].run.dataset.level
        gridfile = entropice.utils.paths.get_grid_file(grid, level)
        cells = gpd.read_parquet(gridfile, columns=["cell_id", "geometry"])
        if grid == "hex":
            cells["cell_id"] = cells["cell_id"].apply(lambda x: int(x, 16))
        cells = cells.set_index("cell_id")

        vals = []
        for tr in training_results:
            preds = tr.run.predictions.set_index("cell_id")[["predicted"]]
            if preds["predicted"].dtype == "category":
                # We can do this because the categories are ordered
                preds["predicted"] = preds["predicted"].cat.codes
            vals.append(preds)
        all_preds = pd.concat(vals, axis=1)
        mean_preds = all_preds.mean(axis=1)
        std_preds = all_preds.std(axis=1)
        cells["mean_prediction"] = mean_preds
        cells["std_prediction"] = std_preds

        return cells.reset_index()


@st.cache_data(ttl=300)  # Cache for 5 minutes
def load_all_training_results() -> list[TrainingResult]:
    """Load all training results from the results directory."""
    results_dir = entropice.utils.paths.RESULTS_DIR
    training_results: list[TrainingResult] = []
    incomplete_results: list[tuple[Path, Exception]] = []
    for result_path in results_dir.iterdir():
        if not result_path.is_dir():
            continue
        try:
            training_result = TrainingResult.from_path(result_path)
            training_results.append(training_result)
        except FileNotFoundError as e:
            is_experiment_dir = False
            for experiment_path in result_path.iterdir():
                if not experiment_path.is_dir():
                    continue
                try:
                    experiment_name = experiment_path.parent.name
                    training_result = TrainingResult.from_path(experiment_path, experiment_name)
                    training_results.append(training_result)
                    is_experiment_dir = True
                except FileNotFoundError as e2:
                    incomplete_results.append((experiment_path, e2))
            if not is_experiment_dir:
                incomplete_results.append((result_path, e))

    if len(incomplete_results) > 0:
        st.warning(
            f"Found {len(incomplete_results)} incomplete training results that were skipped:\n - "
            + "\n - ".join(f"{p}: {e}" for p, e in incomplete_results)
        )
    # Sort by creation time (most recent first)
    training_results.sort(key=lambda tr: tr.created_at, reverse=True)
    return training_results


def load_training_sets(ensemble: DatasetEnsemble) -> dict[TargetDataset, dict[Task, TrainingSet]]:
    """Load training sets for all target-task combinations in the ensemble."""
    train_data_dict: dict[TargetDataset, dict[Task, TrainingSet]] = defaultdict(dict)
    for target in all_target_datasets:
        for task in all_tasks:
            train_data_dict[target][task] = ensemble.create_training_set(target=target, task=task)
    return train_data_dict


def get_available_experiments() -> list[str]:
    """Get list of available experiment names from the results directory.

    Returns:
        List of experiment directory names

    """
    results_dir = entropice.utils.paths.RESULTS_DIR
    experiments = []

    for item in results_dir.iterdir():
        if not item.is_dir():
            continue

        # Check if this directory contains training results (has subdirs with training results files)
        has_training_results = False
        for subitem in item.iterdir():
            if not subitem.is_dir():
                continue
            # Check for either RandomSearchCV or AutoGluon result files
            if (subitem / "search_results.parquet").exists() or (subitem / "leaderboard.parquet").exists():
                has_training_results = True
                break

        if has_training_results:
            experiments.append(item.name)

    return sorted(experiments)


def load_experiment_training_results(experiment_name: str) -> list[TrainingResult]:
    """Load all training results for a specific experiment.

    Args:
        experiment_name: Name of the experiment directory

    Returns:
        List of TrainingResult objects for the experiment

    """
    experiment_dir = entropice.utils.paths.RESULTS_DIR / experiment_name
    if not experiment_dir.exists():
        return []

    training_results: list[TrainingResult] = []
    for result_path in experiment_dir.iterdir():
        if not result_path.is_dir():
            continue
        try:
            training_result = TrainingResult.from_path(result_path, experiment_name)
            training_results.append(training_result)
        except FileNotFoundError:
            pass  # Skip incomplete results

    # Sort by creation time (most recent first)
    training_results.sort(key=lambda tr: tr.created_at, reverse=True)
    return training_results


def create_experiment_summary_df(training_results: list[TrainingResult]) -> pd.DataFrame:
    """Create a summary DataFrame for all results in an experiment.

    Args:
        training_results: List of TrainingResult objects

    Returns:
        DataFrame with summary statistics for the experiment

    """
    records = []

    # Add RandomSearchCV results
    for tr in training_results:
        info = tr.display_info
        best_metric_name = tr._get_best_metric_name()

        record = {
            "method": "RandomSearchCV",
            "task": info.task,
            "target": info.target,
            "model": info.model,
            "grid": info.grid,
            "level": info.level,
            "grid_level": f"{info.grid}_{info.level}",
            "train_score": tr._get_best_score("train"),
            "test_score": tr._get_best_score("test"),
            "complete_score": tr._get_best_score("complete"),
            "best_metric": best_metric_name,
            "n_trials": tr.run.n_trials or "N/A",
            "created_at": tr.created_at,
            "path": tr.path,
        }

        # Add all train metrics
        for metric, value in tr.run.get_metrics_from_split("train").items():
            record[f"train_{metric}"] = value

        # Add all test metrics
        for metric, value in tr.run.get_metrics_from_split("test").items():
            record[f"test_{metric}"] = value

        # Add all complete metrics
        for metric, value in tr.run.get_metrics_from_split("complete").items():
            record[f"complete_{metric}"] = value

        records.append(record)

    return pd.DataFrame(records)


@dataclass
class StorageInfo:
    """Storage information for a directory."""

    name: str
    size_bytes: int
    file_count: int
    display_size: str


def _parse_size_to_bytes(size_str: str) -> int:
    """Convert dust's human-readable size string to bytes.

    Examples: "92K" -> 92*1024, "1.5M" -> 1.5*1024*1024, "928B" -> 928
    """
    size_str = size_str.strip().upper()
    if not size_str:
        return 0

    # Extract numeric part and unit
    numeric_part = ""
    unit = ""
    for char in size_str:
        if char.isdigit() or char == ".":
            numeric_part += char
        else:
            unit += char

    try:
        value = float(numeric_part) if numeric_part else 0
    except ValueError:
        return 0

    # Convert based on unit
    unit = unit.strip()
    multipliers = {
        "B": 1,
        "K": 1024,
        "M": 1024**2,
        "G": 1024**3,
        "T": 1024**4,
        "P": 1024**5,
    }

    return int(value * multipliers.get(unit, 1))


def _run_dust_command(data_dir: Path, for_files: bool = False) -> dict | None:
    """Run dust command and return parsed JSON output.

    Args:
        data_dir: Directory to analyze
        for_files: If True, count files (-f flag); if False, count disk space

    Returns:
        Parsed JSON dict or None if command failed

    """
    cmd = ["dust", "-j", "-d", "1"]
    if for_files:
        cmd.append("-f")
    cmd.append(str(data_dir))

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode != 0:
            return None
        return json.loads(result.stdout)
    except (subprocess.TimeoutExpired, json.JSONDecodeError):
        return None


def _build_file_counts_lookup(files_data: dict | None) -> dict[str, int]:
    """Build lookup dict for file counts from dust JSON output."""
    file_counts = {}
    if files_data and "children" in files_data:
        for child in files_data["children"]:
            name = Path(child["name"]).name
            count_str = child.get("size", "0")
            file_counts[name] = _parse_size_to_bytes(count_str)
    return file_counts


@st.cache_data(ttl=300)  # Cache for 5 minutes
def load_storage_statistics() -> tuple[list[StorageInfo], int, int]:
    """Load storage statistics for DATA_DIR subdirectories using dust.

    Returns:
        Tuple of (subdirectory stats list, total size in bytes, total file count)

    """
    data_dir = entropice.utils.paths.DATA_DIR

    if not data_dir.exists():
        return [], 0, 0

    try:
        # Run dust for disk space and file counts
        space_data = _run_dust_command(data_dir, for_files=False)
        files_data = _run_dust_command(data_dir, for_files=True)

        if not space_data:
            st.warning("Failed to get storage statistics from dust")
            return [], 0, 0

        # Build lookup dict for file counts
        file_counts = _build_file_counts_lookup(files_data)

        # Extract subdirectory information from space data
        storage_infos = []
        total_size = 0
        total_files = 0

        if "children" in space_data:
            for child in space_data["children"]:
                full_path = child.get("name", "")
                dir_name = Path(full_path).name
                size_str = child.get("size", "0")
                size_bytes = _parse_size_to_bytes(size_str)
                file_count = file_counts.get(dir_name, 0)

                storage_infos.append(
                    StorageInfo(
                        name=dir_name,
                        size_bytes=size_bytes,
                        file_count=file_count,
                        display_size=size_str,
                    )
                )
                total_size += size_bytes
                total_files += file_count

        # Sort by size descending
        storage_infos.sort(key=lambda x: x.size_bytes, reverse=True)

        return storage_infos, total_size, total_files

    except FileNotFoundError:
        st.error("dust command not found. Please install dust: https://github.com/bootandy/dust")
        return [], 0, 0
    except Exception as e:
        st.error(f"Error getting storage statistics: {e}")
        return [], 0, 0
