"""
Visualization module — Plotting functionality for BetaSpectrumAnalyzer.
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import matplotlib.pyplot as plt

from beta_spectrum.logging_utils import get_git_short_hash
from beta_spectrum.utils import T_to_W

if TYPE_CHECKING:
    from matplotlib.figure import Figure

    from beta_spectrum.output.core import BetaSpectrumAnalyzer


class Visualizer:
    """Visualization functionality for BetaSpectrumAnalyzer."""

    @staticmethod
    def plot_analysis(
        analyzer: BetaSpectrumAnalyzer,
        save_path: Optional[str] = None,
        show_components: bool = True,
    ) -> None:
        """
        Create visualization of the spectrum and all correction factors.

        Parameters
        ----------
        analyzer : BetaSpectrumAnalyzer
            The analyzer instance to visualize.
        save_path : str, optional
            Path to save the figure.
        show_components : bool
            If True, show full debug view.
            If False, show only the spectrum plot.
        """
        total = analyzer.total_spectrum(normalize=True)
        components = analyzer.components
        analyzer._timestamp = datetime.now(timezone.utc).strftime(
            "%Y-%m-%d %H:%M:%S UTC"
        )

        if show_components:
            Visualizer._plot_debug_view(analyzer, total, components, save_path)
        else:
            Visualizer._plot_spectra_view(analyzer, total, components, save_path)

    @staticmethod
    def _build_figure_header(analyzer: BetaSpectrumAnalyzer, fig: "Figure") -> None:
        """Add a compact header with nuclide info, e_step, commit and timestamp outside axes."""
        parent = analyzer._element_symbol(analyzer.config.Z_parent)
        daughter = analyzer._element_symbol(analyzer.config.Z_daughter)
        e_step_keV = analyzer.config.e_step_MeV * 1000
        header = (
            f"{parent}{analyzer.config.A_number} -> {daughter}{analyzer.config.A_number}  "
            f"|  E₀={analyzer.config.endpoint_MeV * 1000:.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",
        )

        # Add neutrino mass info to header if non-zero
        if analyzer.config.neutrino_mass_MeV > 0:
            nu_mass_text = f"|  m_ν={analyzer.config.neutrino_mass_MeV:.6f} MeV"
            fig.text(
                0.01,
                0.96,
                nu_mass_text,
                transform=fig.transFigure,
                fontsize=8,
                verticalalignment="top",
                horizontalalignment="left",
                fontfamily="monospace",
                color="red",
            )

        id_text = f"{get_git_short_hash(6)}  |  {analyzer._timestamp}"
        fig.text(
            0.01,
            0.01,
            id_text,
            transform=fig.transFigure,
            fontsize=8,
            verticalalignment="bottom",
            horizontalalignment="left",
            fontfamily="monospace",
            color="gray",
        )

    @staticmethod
    def _build_figure_title(analyzer: BetaSpectrumAnalyzer, fig: "Figure") -> str:
        """Build a unified figure title."""
        parent = analyzer._element_symbol(analyzer.config.Z_parent)
        daughter = analyzer._element_symbol(analyzer.config.Z_daughter)
        n = analyzer.config.A_number
        title = f"Beta-decay: {n}{parent} -> {n}{daughter}"
        if analyzer._n_branches > 1:
            # Count active branches (not excluded by neutrino mass or intensity cutoff)
            active_branches = 0
            for i, branch in enumerate(analyzer._unified_branches):
                branch_endpoint_keV = branch.endpoint_MeV * 1000
                neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
                is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
                is_excluded_by_intensity = (
                    branch.intensity < analyzer.config.intensity_cutoff
                )
                if not is_excluded_by_mass and not is_excluded_by_intensity:
                    active_branches += 1

            if active_branches < analyzer._n_branches:
                title += f" ({active_branches} of {analyzer._n_branches} branches)"
            else:
                title += f" ({analyzer._n_branches} branches)"
        return title

    @staticmethod
    def _format_branch_label(analyzer: BetaSpectrumAnalyzer, i: int) -> str:
        """Abbreviated branch label for legend."""
        branch = analyzer._unified_branches[i]
        e0 = branch.endpoint_MeV * 1000
        intensity_pct = branch.intensity * 100

        # Check if branch is excluded
        branch_endpoint_keV = branch.endpoint_MeV * 1000
        neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
        is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
        is_excluded_by_intensity = branch.intensity < analyzer.config.intensity_cutoff

        if is_excluded_by_mass:
            return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (endpoint < m_nu)"
        elif is_excluded_by_intensity:
            return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%) - excluded (intensity < cutoff)"
        else:
            return f"Br. {i + 1}: {e0:.1f} keV ({intensity_pct:.1f}%)"

    @staticmethod
    def _plot_spectra_view(
        analyzer: BetaSpectrumAnalyzer,
        total: np.ndarray,
        components: Dict[str, np.ndarray],
        save_path: Optional[str],
    ) -> None:
        """Plot total spectrum with individual branch decomposition."""
        fig, ax = plt.subplots(figsize=(12, 7))

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

        e_step_keV = analyzer.config.e_step_MeV * 1000

        # Compute branch contributions (universal * intensity * branch_spectrum)
        branch_contribs = []
        # Universal components are shared (no per-branch key prefix)
        universal_names = [
            "Fermi",
            "Screening",
            "FiniteSizeL0",
            "ChargeDistributionU",
            "Exchange",
        ]
        for i in range(analyzer._n_branches):
            universal_product = np.ones_like(analyzer.W)
            for name in universal_names:
                key = analyzer._get_component_key(i, name)
                if key in components:
                    universal_product *= components[key]
                elif name in components:  # fallback: universal components
                    universal_product *= components[name]
            bs = analyzer._get_branch_spectrum(i)
            contrib = universal_product * analyzer._unified_branches[i].intensity * bs
            branch_contribs.append(contrib)

        total_raw = sum(branch_contribs)

        ax.plot(
            analyzer.energies_MeV,
            total_raw,
            "k-",
            lw=2.5,
            label="Total spectrum",
            zorder=10,
        )

        # Only plot active branches (not excluded by neutrino mass or intensity cutoff)
        active_branch_indices = []
        for i in range(analyzer._n_branches):
            branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000
            neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
            is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
            is_excluded_by_intensity = (
                analyzer._unified_branches[i].intensity
                < analyzer.config.intensity_cutoff
            )
            if not is_excluded_by_mass and not is_excluded_by_intensity:
                active_branch_indices.append(i)

        for i in active_branch_indices:
            contrib = branch_contribs[i]
            color = branch_colors[i % len(branch_colors)]
            ax.plot(
                analyzer.energies_MeV,
                contrib,
                color=color,
                lw=1.5,
                alpha=0.7,
                label=Visualizer._format_branch_label(analyzer, i),
            )

        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(
            Visualizer._build_figure_title(analyzer, fig),
            fontsize=14,
            fontweight="bold",
        )
        ax.set_yscale("log")
        ax.grid(True, alpha=0.3)
        ax.legend(
            loc="best",
            fontsize=9,
            ncol=min(2, len(active_branch_indices) + 1),
        )

        Visualizer._build_figure_header(analyzer, fig)

        plt.tight_layout()

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

    @staticmethod
    def _plot_debug_view(
        analyzer: BetaSpectrumAnalyzer,
        total: np.ndarray,
        components: Dict[str, np.ndarray],
        save_path: Optional[str],
    ) -> None:
        """Debug view: vertical layout with 4 panels."""
        n_branches = analyzer._n_branches
        e_step_keV = analyzer.config.e_step_MeV * 1000

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

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

        # Panel 1: Total spectrum + per-branch spectra
        branch_contribs = []
        universal_names = [
            "Fermi",
            "Screening",
            "FiniteSizeL0",
            "ChargeDistributionU",
            "Exchange",
        ]
        for i in range(n_branches):
            universal_product = np.ones_like(analyzer.W)
            for name in universal_names:
                key = analyzer._get_component_key(i, name)
                if key in components:
                    universal_product *= components[key]
                elif name in components:
                    universal_product *= components[name]
            bs = analyzer._get_branch_spectrum(i)
            contrib = universal_product * analyzer._unified_branches[i].intensity * bs
            branch_contribs.append(contrib)

        total_raw = sum(branch_contribs)

        ax = axes[0]
        ax.plot(
            analyzer.energies_MeV,
            total_raw,
            "k-",
            lw=2,
            label="Total spectrum",
            zorder=10,
        )

        # Only plot active branches (not excluded by neutrino mass or intensity cutoff)
        active_branch_indices = []
        for i in range(n_branches):
            branch_endpoint_keV = analyzer._unified_branches[i].endpoint_MeV * 1000
            neutrino_mass_keV = analyzer.config.neutrino_mass_MeV * 1000
            is_excluded_by_mass = branch_endpoint_keV <= neutrino_mass_keV
            is_excluded_by_intensity = (
                analyzer._unified_branches[i].intensity
                < analyzer.config.intensity_cutoff
            )
            if not is_excluded_by_mass and not is_excluded_by_intensity:
                active_branch_indices.append(i)

        for i in active_branch_indices:
            contrib = branch_contribs[i]
            color = branch_colors[i % len(branch_colors)]
            ax.plot(
                analyzer.energies_MeV,
                contrib,
                color=color,
                lw=1.5,
                alpha=0.7,
                label=Visualizer._format_branch_label(analyzer, i),
            )
        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=min(2, len(active_branch_indices) + 1))

        # Panel 2: Fermi function + all phase spaces
        ax = axes[1]
        if "Fermi" in components:
            fermi_vals = components["Fermi"]
            fermi_scaled = fermi_vals / np.max(fermi_vals)
            ax.plot(
                analyzer.energies_MeV,
                fermi_scaled,
                "r-",
                lw=1.5,
                label="Fermi",
                alpha=0.8,
            )
        for i in active_branch_indices:
            color = branch_colors[i % len(branch_colors)]
            branch = analyzer._unified_branches[i]
            key = analyzer._get_component_key(i, "PhaseSpace")
            if key in components:
                ps_vals = components[key]
                W0_branch = T_to_W(branch.endpoint_MeV)
                mask = analyzer.W <= W0_branch
                ps_masked = np.where(mask, ps_vals, 0.0)
                ps_max = np.max(ps_masked)
                if ps_max > 0:
                    ps_scaled = ps_masked / ps_max
                else:
                    ps_scaled = ps_masked
                ax.plot(
                    analyzer.energies_MeV,
                    ps_scaled,
                    color=color,
                    lw=1.5,
                    alpha=0.7,
                    label=f"PS Br. {i + 1}",
                )
        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 W0-independent corrections
        ax = axes[2]
        universal_components = [
            "Screening",
            "FiniteSizeL0",
            "ChargeDistributionU",
            "Exchange",
        ]
        colors = ["green", "orange", "purple", "brown"]
        for name, color in zip(universal_components, colors):
            if name in components:
                values = components[name]
                values_scaled = values / np.max(values)
                ax.plot(
                    analyzer.energies_MeV,
                    values_scaled,
                    color=color,
                    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 for each branch)
        ax = axes[3]
        for i in active_branch_indices:
            color = branch_colors[i % len(branch_colors)]
            key = analyzer._get_component_key(i, "Radiative")
            if key in components:
                rad_vals = components[key]
                ax.plot(
                    analyzer.energies_MeV,
                    rad_vals,
                    color=color,
                    lw=1.5,
                    alpha=0.7,
                    label=f"Br. {i + 1}",
                )
        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=min(2, len(active_branch_indices) + 1))

        Visualizer._build_figure_header(analyzer, fig)

        plt.tight_layout()

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