"""
Visualization module — plotting built ON TOP OF the CSV export.

All graphical representations are produced from the exported CSV file
(data columns + metadata header) and never recomputed from the analyzer in
parallel. The CSV file is the single source of truth: a plot is produced
exactly from the bytes a researcher would download, which turns every plot
into an additional quality check of the export routine.
"""

from __future__ import annotations

import re
import shutil
import tempfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
from matplotlib.figure import Figure

from .csv_export import CSVEmitter

if TYPE_CHECKING:
    from beta_spectrum.output.core import BetaSpectrumAnalyzer

_BRANCH_LINE_RE = re.compile(
    r"^\s*branch_(\d+):\s*E₀=([0-9.]+)\s*keV,\s*transition=(\w+),\s*"
    r"intensity=([0-9.]+)(.*)$"
)
_EXCLUDED_SUFFIX_RE = re.compile(r"-\s*excluded:\s*(.*)$")

_BRANCH_COLORS = [
    "#1f77b4",
    "#ff7f0e",
    "#2ca02c",
    "#d62728",
    "#9467bd",
    "#8c564b",
    "#e377c2",
    "#7f7f7f",
    "#bcbd22",
    "#17becf",
]


def parse_spectrum_csv(
    csv_path: str,
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    """Read an exported spectrum CSV and parse its metadata header.

    Parameters
    ----------
    csv_path : str
        Path to a CSV file produced by CSVEmitter.export_to_csv.

    Returns
    -------
    tuple[pandas.DataFrame, dict]
        The data frame (comment lines skipped, separator auto-detected)
        and a metadata dict with keys: version, build, timestamp, source,
        nuclide, endpoint_keV, neutrino_mass_MeV, corrections, e_step_MeV,
        n_points, normalization, branches.

    Raises
    ------
    ValueError
        If the file cannot be read or parsed.
    """
    try:
        df = pd.read_csv(csv_path, sep=None, engine="python", comment="#")
        raw_lines = []
        with open(csv_path, "r") as f:
            for line in f:
                if line.startswith("#"):
                    raw_lines.append(line.rstrip("\n"))
                else:
                    break  # header block ended
    except OSError as exc:
        raise ValueError(f"Cannot read spectrum CSV {csv_path}: {exc}") from exc

    meta: Dict[str, Any] = {
        "version": None,
        "build": None,
        "timestamp": None,
        "source": None,
        "nuclide": None,
        "endpoint_keV": None,
        "neutrino_mass_MeV": None,
        "corrections": [],
        "e_step_MeV": None,
        "n_points": None,
        "normalization": None,
        "branches": [],
    }

    for line in raw_lines:
        body = line.lstrip("#").strip()
        m = re.match(r"beta-spectrum v([\d.]+)\(build\.(\w+)\)", body)
        if m:
            meta["version"] = m.group(1)
            meta["build"] = m.group(2)
            continue
        m = re.match(r"timestamp:\s*(.+)", body)
        if m:
            meta["timestamp"] = m.group(1).strip()
            continue
        m = re.match(r"source:\s*(\S+)", body)
        if m:
            meta["source"] = m.group(1)
            continue
        m = re.match(r"nuclide:\s*(.+?)\s+\(Z=", body)
        if m:
            meta["nuclide"] = m.group(1).strip()
            continue
        m = re.match(r"basic_endpoint:\s*([0-9.]+)\s*keV", body)
        if m:
            meta["endpoint_keV"] = float(m.group(1))
            continue
        m = re.match(r"neutrino_mass_MeV:\s*([0-9.]+)", body)
        if m:
            meta["neutrino_mass_MeV"] = float(m.group(1))
            continue
        m = re.match(r"corrections:\s*(.*)", body)
        if m:
            meta["corrections"] = [
                c.strip() for c in m.group(1).split(",") if c.strip()
            ]
            continue
        m = re.match(r"e_step:\s*([0-9.]+)\s*MeV", body)
        if m:
            meta["e_step_MeV"] = float(m.group(1))
            continue
        m = re.match(r"n_points:\s*(\d+)", body)
        if m:
            meta["n_points"] = int(m.group(1))
            continue
        m = re.match(r"branch normalization:\s*(\S+)", body)
        if m:
            meta["normalization"] = m.group(1)
            continue
        m = _BRANCH_LINE_RE.match(body)
        if m:
            excluded = _EXCLUDED_SUFFIX_RE.search(m.group(5))
            meta["branches"].append(
                {
                    "index": int(m.group(1)),
                    "endpoint_keV": float(m.group(2)),
                    "transition": m.group(3),
                    "intensity": float(m.group(4)),
                    "excluded": None if excluded is None else excluded.group(1).strip(),
                }
            )
            continue

    if meta["n_points"] is None:
        meta["n_points"] = len(df)

    return df, meta


def _branch_column(df: pd.DataFrame, index: int) -> Optional[str]:
    """Find the branch spectrum column name for a 1-based branch index."""
    for candidate in (f"branch_{index}_spectrum", f"branch_{index}"):
        if candidate in df.columns:
            return candidate
    return None


def _active_branches(meta: Dict[str, Any]) -> List[Dict[str, Any]]:
    return [b for b in meta["branches"] if b["excluded"] is None]


def _spectrum_title(meta: Dict[str, Any]) -> str:
    n_branches = len(meta["branches"])
    active = len(_active_branches(meta))
    title = f"Beta-decay: {meta['nuclide']}"
    if n_branches > 1:
        if active < n_branches:
            title += f" ({active} of {n_branches} branches)"
        else:
            title += f" ({n_branches} branches)"
    return title


def _apply_figure_header(fig: Figure, meta: Dict[str, Any]) -> None:
    """Write the figure header/footer lines from the CSV metadata."""
    e_step_keV = (meta["e_step_MeV"] or 0.0) * 1000.0
    endpoint = meta["endpoint_keV"]
    header = (
        f"{meta['nuclide']}  "
        f"|  E₀={endpoint:.1f} keV  "
        f"|  ΔE={e_step_keV:.1f} keV"
    )
    fig.text(
        0.01,
        0.98,
        header,
        transform=fig.transFigure,
        fontsize=9,
        verticalalignment="top",
        horizontalalignment="left",
        fontfamily="monospace",
    )

    if meta["neutrino_mass_MeV"]:
        fig.text(
            0.01,
            0.96,
            f"|  m_ν={meta['neutrino_mass_MeV']:.6f} MeV",
            transform=fig.transFigure,
            fontsize=8,
            verticalalignment="top",
            horizontalalignment="left",
            fontfamily="monospace",
            color="red",
        )

    if meta["build"]:
        fig.text(
            0.01,
            0.01,
            f"{meta['build']}  |  {meta['timestamp']}",
            transform=fig.transFigure,
            fontsize=8,
            verticalalignment="bottom",
            horizontalalignment="left",
            fontfamily="monospace",
            color="gray",
        )


def _plot_spectra_view_from_df(
    df: pd.DataFrame,
    meta: Dict[str, Any],
    save_path: Optional[str],
) -> Figure:
    """Plot total spectrum and per-branch spectra from CSV columns."""
    energies = df["energy_MeV"].to_numpy()
    e_step_keV = (meta["e_step_MeV"] or 0.0) * 1000.0

    fig, ax = plt.subplots(figsize=(12, 7))

    if "total_spectrum" in df.columns:
        ax.plot(
            energies,
            df["total_spectrum"].to_numpy(),
            "k-",
            lw=2.5,
            label="Total spectrum",
            zorder=10,
        )

    for color_i, b in enumerate(_active_branches(meta)):
        col = _branch_column(df, int(b["index"]))
        if col is None:
            continue
        ax.plot(
            energies,
            df[col].to_numpy(),
            color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
            lw=1.5,
            alpha=0.7,
            label=(
                f"Br. {b['index']}: {b['endpoint_keV']:.1f} keV "
                f"({b['intensity'] * 100:.1f}%)"
            ),
        )

    ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=12)
    ax.set_ylabel(f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=12)
    fig.suptitle(_spectrum_title(meta), fontsize=14, fontweight="bold")
    ax.set_yscale("log")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best", fontsize=9, ncol=2)

    _apply_figure_header(fig, meta)
    fig.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Figure saved to {save_path}")
    return fig


def _plot_debug_view_from_df(
    df: pd.DataFrame,
    meta: Dict[str, Any],
    save_path: Optional[str],
) -> Figure:
    """Debug view (4 panels) built entirely from CSV columns."""
    energies = df["energy_MeV"].to_numpy()
    e_step_keV = (meta["e_step_MeV"] or 0.0) * 1000.0
    active = _active_branches(meta)

    fig, axes = plt.subplots(4, 1, figsize=(12, 16), sharex=True)
    fig.suptitle(_spectrum_title(meta), fontsize=14, fontweight="bold")

    # Panel 1: total spectrum + per-branch spectra
    ax = axes[0]
    if "total_spectrum" in df.columns:
        ax.plot(
            energies,
            df["total_spectrum"].to_numpy(),
            "k-",
            lw=2,
            label="Total spectrum",
            zorder=10,
        )
    for color_i, b in enumerate(active):
        col = _branch_column(df, int(b["index"]))
        if col is None:
            continue
        ax.plot(
            energies,
            df[col].to_numpy(),
            color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
            lw=1.5,
            alpha=0.7,
            label=f"Br. {b['index']} ({b['intensity'] * 100:.1f}%)",
        )
    ax.set_yscale("log")
    ax.set_ylabel(f"Normalized Counts / {e_step_keV:.1f} keV", fontsize=10)
    ax.set_title("Spectrum", fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8, loc="best", ncol=2)

    # Panel 2: Fermi function + all phase spaces
    ax = axes[1]
    if "Fermi" in df.columns:
        fermi = df["Fermi"].to_numpy()
        fermi_scaled = fermi / np.max(fermi)
        ax.plot(energies, fermi_scaled, "r-", lw=1.5, label="Fermi", alpha=0.8)
    for color_i, b in enumerate(active):
        col = f"branch_{b['index']}_PhaseSpace"
        if col not in df.columns:
            continue
        ps_vals = df[col].to_numpy()
        ps_max = np.max(ps_vals) if np.any(ps_vals > 0) else 1.0
        ax.plot(
            energies,
            ps_vals / ps_max if ps_max > 0 else ps_vals,
            color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
            lw=1.5,
            alpha=0.7,
            label=f"PS Br. {b['index']}",
        )
    ax.set_ylabel("Normalized Factor", fontsize=10)
    ax.set_title("Fermi & Phase Space", fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8, loc="best")

    # Panel 3: universal higher-order corrections
    ax = axes[2]
    correction_colors = ["green", "orange", "purple", "brown"]
    for corr_i, name in enumerate(
        ["Screening", "FiniteSizeL0", "ChargeDistributionU", "Exchange"]
    ):
        if name in df.columns:
            values = df[name].to_numpy()
            values_scaled = values / np.max(values)
            ax.plot(
                energies,
                values_scaled,
                color=correction_colors[corr_i],
                lw=1.5,
                label=name,
                alpha=0.8,
            )
    ax.set_ylabel("Normalized Factor", fontsize=10)
    ax.set_title("Universal Corrections", fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8, loc="best")

    # Panel 4: W0-dependent corrections (Radiative per branch)
    ax = axes[3]
    for color_i, b in enumerate(active):
        col = f"branch_{b['index']}_Radiative"
        if col not in df.columns:
            continue
        ax.plot(
            energies,
            df[col].to_numpy(),
            color=_BRANCH_COLORS[color_i % len(_BRANCH_COLORS)],
            lw=1.5,
            alpha=0.7,
            label=f"Br. {b['index']}",
        )
    ax.set_ylabel("Radiative Correction", fontsize=10)
    ax.set_xlabel(r"Electron kinetic energy $E$ [MeV]", fontsize=10)
    ax.set_title("Radiative Corrections", fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8, loc="best", ncol=2)

    _apply_figure_header(fig, meta)
    fig.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Figure saved to {save_path}")
    return fig


class Visualizer:
    """Visualization built from the exported CSV data (single source of truth)."""

    @staticmethod
    def plot_from_csv(
        csv_path: str,
        save_path: Optional[str] = None,
        show_components: bool = True,
    ) -> Figure:
        """Create a visualization from an exported CSV spectrum file.

        Parameters
        ----------
        csv_path : str
            Path to a CSV exported by CSVEmitter.export_to_csv. The plots
            are built exclusively from the file contents (data columns and
            metadata header) — nothing is recomputed from the model.
        save_path : str, optional
            Path to save the figure.
        show_components : bool
            If True, produce the full 4-panel debug view; requires the
            'full' CSV verbosity. If the file lacks component columns, a
            warning is printed and the simple spectrum view is produced.
        """
        df, meta = parse_spectrum_csv(csv_path)
        if show_components and "Fermi" in df.columns:
            return _plot_debug_view_from_df(df, meta, save_path)
        if show_components:
            print(
                "Warning: requested component panels but CSV lacks component "
                "columns (export with --csv-verbosity full); falling back to "
                "simple spectrum view."
            )
        return _plot_spectra_view_from_df(df, meta, save_path)

    @staticmethod
    def plot_analysis(
        analyzer: BetaSpectrumAnalyzer,
        save_path: Optional[str] = None,
        show_components: bool = True,
    ) -> Figure:
        """Create a visualization by exporting CSV and plotting from it.

        The export goes to a temporary file using the 'full' verbosity so
        that the plots exercise the same data pipeline (header + columns)
        a researcher receives.
        """
        tmp_dir = tempfile.mkdtemp(prefix="bs_plot_")
        try:
            csv_path = f"{tmp_dir}/spectrum.csv"
            CSVEmitter.export_to_csv(
                analyzer,
                csv_path,
                source_type="plot",
                csv_verbosity="full",
            )
            return Visualizer.plot_from_csv(
                csv_path,
                save_path=save_path,
                show_components=show_components,
            )
        finally:
            shutil.rmtree(tmp_dir, ignore_errors=True)