"""Experiment visualization utilities using matplotlib.

Generates scientific-style charts comparing evaluations within an experiment,
grouped or plotted by a user-selected metadata key.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

try:
    import matplotlib.pyplot as plt
    import matplotlib
    import numpy as np

    HAS_MATPLOTLIB = True
except ImportError:  # pragma: no cover
    HAS_MATPLOTLIB = False


def _require_matplotlib() -> None:
    if not HAS_MATPLOTLIB:
        raise ImportError(
            "matplotlib is required for plotting. Install it with: pip install matplotlib"
        )


# ---------------------------------------------------------------------------
# Metric extraction helpers
# ---------------------------------------------------------------------------

_RETRIEVAL_METRICS = [
    ("chunk_precision", "Chunk Precision"),
    ("chunk_recall", "Chunk Recall"),
    ("chunk_f1", "Chunk F1"),
    ("page_precision", "Page Precision"),
    ("page_recall", "Page Recall"),
    ("page_f1", "Page F1"),
    ("document_precision", "Doc Precision"),
    ("document_recall", "Doc Recall"),
    ("document_f1", "Doc F1"),
]

_GENERATION_METRICS = [
    ("accuracy", "Accuracy"),
    ("groundedness", "Groundedness"),
]


def _get_metric_value(evaluation: Dict[str, Any], metric_key: str) -> Optional[float]:
    """Extract a metric value from a flat or nested evaluation dict."""
    # Try flat keys first (backend response style)
    if metric_key in evaluation and evaluation[metric_key] is not None:
        return float(evaluation[metric_key])

    # Try nested style (chunk_level.precision etc.)
    mapping = {
        "chunk_precision": ("chunk_level", "precision"),
        "chunk_recall": ("chunk_level", "recall"),
        "chunk_f1": ("chunk_level", "f1_score"),
        "page_precision": ("page_level", "precision"),
        "page_recall": ("page_level", "recall"),
        "page_f1": ("page_level", "f1_score"),
        "document_precision": ("document_level", "precision"),
        "document_recall": ("document_level", "recall"),
        "document_f1": ("document_level", "f1_score"),
        "accuracy": ("generation_metrics", "accuracy"),
        "groundedness": ("generation_metrics", "groundedness"),
    }

    if metric_key in mapping:
        parent_key, child_key = mapping[metric_key]
        parent = evaluation.get(parent_key)
        if isinstance(parent, dict) and child_key in parent:
            v = parent[child_key]
            return float(v) if v is not None else None

    return None


def _pick_metrics_for_type(evaluation_type: str) -> List[Tuple[str, str]]:
    """Return the relevant (key, label) pairs for a given evaluation type."""
    if evaluation_type == "retrieval_only":
        return _RETRIEVAL_METRICS
    elif evaluation_type == "generation_only":
        return _GENERATION_METRICS
    elif evaluation_type == "retrieval_and_generation":
        return _RETRIEVAL_METRICS + _GENERATION_METRICS
    # Fallback – include everything
    return _RETRIEVAL_METRICS + _GENERATION_METRICS


def _is_numeric(values: Sequence[Any]) -> bool:
    """Return True when *all* non-None values can be interpreted as numbers."""
    for v in values:
        if v is None:
            continue
        if isinstance(v, (int, float)):
            continue
        try:
            float(v)
        except (ValueError, TypeError):
            return False
    return True


def _get_eval_metadata(ev: Dict[str, Any]) -> Dict[str, Any]:
    """Get the metadata dict from an evaluation.

    Handles both SDK-side results (``metadata`` key) and API responses
    (``run_metadata`` key).
    """
    return ev.get("run_metadata") or ev.get("metadata") or {}


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def plot_experiment(
    evaluations: List[Dict[str, Any]],
    metadata_key: str,
    metrics: Optional[List[str]] = None,
    title: Optional[str] = None,
    figsize: Tuple[int, int] = (12, 6),
    save_path: Optional[str] = None,
    show: bool = True,
) -> Any:
    """Plot evaluations from an experiment, grouped by a metadata key.

    Parameters
    ----------
    evaluations : list of dict
        Each dict must contain at minimum ``name``, ``evaluation_type``,
        ``metadata`` (dict), and the relevant metric fields (either flat or
        nested).  This is the shape returned by the ``GET /experiments/{id}``
        endpoint or by ``VectaAPIClient.get_experiment()``.
    metadata_key : str
        The metadata key to use for the x-axis / grouping.
    metrics : list of str, optional
        Specific metric keys to plot (e.g. ``["chunk_f1", "accuracy"]``).
        When *None*, all metrics applicable to the evaluation type are used.
    title : str, optional
        Chart title.  Defaults to ``"Experiment: {metadata_key}"``.
    figsize : tuple of int
        Figure size in inches.
    save_path : str, optional
        If provided the figure is saved to this path.
    show : bool
        Whether to call ``plt.show()``.

    Returns
    -------
    matplotlib.figure.Figure
    """
    _require_matplotlib()

    if not evaluations:
        raise ValueError("No evaluations to plot")

    # ----- Filter evaluations that have the requested metadata key -----
    filtered: List[Dict[str, Any]] = []
    for ev in evaluations:
        meta = _get_eval_metadata(ev)
        if metadata_key in meta:
            filtered.append(ev)

    if not filtered:
        raise ValueError(
            f"No evaluations have metadata key '{metadata_key}'. "
            f"Available keys: {_available_metadata_keys(evaluations)}"
        )

    # ----- Determine metric list -----
    eval_types = {ev.get("evaluation_type", "") for ev in filtered}
    if metrics:
        all_metric_tuples = [(k, k.replace("_", " ").title()) for k in metrics]
    else:
        # Pick metrics from the most inclusive eval type present
        if "retrieval_and_generation" in eval_types:
            all_metric_tuples = _pick_metrics_for_type("retrieval_and_generation")
        elif "retrieval_only" in eval_types and "generation_only" in eval_types:
            all_metric_tuples = _RETRIEVAL_METRICS + _GENERATION_METRICS
        elif "retrieval_only" in eval_types:
            all_metric_tuples = _pick_metrics_for_type("retrieval_only")
        elif "generation_only" in eval_types:
            all_metric_tuples = _pick_metrics_for_type("generation_only")
        else:
            all_metric_tuples = _pick_metrics_for_type("retrieval_and_generation")

    # Filter to metrics that have at least one non-None value across evals
    metric_tuples: List[Tuple[str, str]] = []
    for key, label in all_metric_tuples:
        if any(_get_metric_value(ev, key) is not None for ev in filtered):
            metric_tuples.append((key, label))

    if not metric_tuples:
        raise ValueError("No metric data found for the given evaluations")

    # ----- Extract metadata values -----
    meta_values = [_get_eval_metadata(ev)[metadata_key] for ev in filtered]

    # ----- Decide chart type -----
    if _is_numeric(meta_values):
        fig = _plot_line_chart(
            filtered, metadata_key, meta_values, metric_tuples, title, figsize
        )
    else:
        fig = _plot_bar_chart(
            filtered, metadata_key, meta_values, metric_tuples, title, figsize
        )

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")

    if show:
        plt.show()

    return fig


def get_metadata_keys(evaluations: List[Dict[str, Any]]) -> List[str]:
    """Return all unique metadata keys seen across the given evaluations."""
    keys: set[str] = set()
    for ev in evaluations:
        meta = _get_eval_metadata(ev)
        keys.update(meta.keys())
    return sorted(keys)


# ---------------------------------------------------------------------------
# Internal chart renderers
# ---------------------------------------------------------------------------


def _available_metadata_keys(evaluations: List[Dict[str, Any]]) -> List[str]:
    return get_metadata_keys(evaluations)


def _plot_line_chart(
    evaluations: List[Dict[str, Any]],
    metadata_key: str,
    meta_values: List[Any],
    metric_tuples: List[Tuple[str, str]],
    title: Optional[str],
    figsize: Tuple[int, int],
) -> Any:
    """Numeric metadata → line chart with each metric as a separate line."""
    fig, ax = plt.subplots(figsize=figsize)

    # Sort by metadata value
    sorted_pairs = sorted(
        zip(meta_values, evaluations), key=lambda p: float(p[0])
    )
    x_vals = [float(p[0]) for p in sorted_pairs]
    sorted_evals = [p[1] for p in sorted_pairs]

    for metric_key, metric_label in metric_tuples:
        y_vals = [_get_metric_value(ev, metric_key) for ev in sorted_evals]
        # Replace None with NaN so matplotlib skips them
        y_clean = [v if v is not None else float("nan") for v in y_vals]
        ax.plot(x_vals, y_clean, marker="o", label=metric_label, linewidth=2)

    ax.set_xlabel(metadata_key.replace("_", " ").title(), fontsize=12)
    ax.set_ylabel("Score", fontsize=12)
    ax.set_title(
        title or f"Experiment: {metadata_key.replace('_', ' ').title()}",
        fontsize=14,
        fontweight="bold",
    )
    ax.set_ylim(-0.05, 1.05)
    ax.legend(loc="best", fontsize=9)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    return fig


def _plot_bar_chart(
    evaluations: List[Dict[str, Any]],
    metadata_key: str,
    meta_values: List[Any],
    metric_tuples: List[Tuple[str, str]],
    title: Optional[str],
    figsize: Tuple[int, int],
) -> Any:
    """Categorical metadata → grouped bar chart."""
    fig, ax = plt.subplots(figsize=figsize)

    labels = [str(v) for v in meta_values]
    n_evals = len(evaluations)
    n_metrics = len(metric_tuples)

    x = np.arange(n_evals)
    total_width = 0.8
    bar_width = total_width / n_metrics

    for i, (metric_key, metric_label) in enumerate(metric_tuples):
        values = []
        for ev in evaluations:
            v = _get_metric_value(ev, metric_key)
            values.append(v if v is not None else 0.0)
        offset = (i - n_metrics / 2 + 0.5) * bar_width
        ax.bar(x + offset, values, bar_width, label=metric_label)

    ax.set_xlabel("Evaluation", fontsize=12)
    ax.set_ylabel("Score", fontsize=12)
    ax.set_title(
        title or f"Experiment: {metadata_key.replace('_', ' ').title()}",
        fontsize=14,
        fontweight="bold",
    )
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=9)
    ax.set_ylim(0, 1.05)
    ax.legend(loc="best", fontsize=9)
    ax.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()
    return fig
