"""Visualization for TQCE components.

Provides publication-quality plots for:
    - Convergence history (fixed-point + PPGD)
    - 3D paradox pressure landscapes
    - CTN graph rendering with edge types
    - TEC probability distributions
    - Bootstrap amplification traces
    - Animated convergence sequences
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Sequence

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
from matplotlib.figure import Figure
from numpy.typing import NDArray

if TYPE_CHECKING:
    from tqce.network.ctn import CausalTensorNetwork


# Color Palette
CAUSAL_COLOR = "#3b82f6"
RETRO_COLOR = "#ef4444"
PRESSURE_CMAP = "inferno"
CONVERGE_COLOR = "#10b981"
LOSS_COLOR = "#f59e0b"
TEC_COLOR = "#8b5cf6"
AMP_COLOR = "#06b6d4"
BG_COLOR = "#0f172a"
TEXT_COLOR = "#e2e8f0"
GRID_COLOR = "#1e293b"


def _apply_dark_theme(fig: Figure, ax: plt.Axes | None = None) -> None:
    """Apply a consistent dark theme to figures."""
    fig.patch.set_facecolor(BG_COLOR)
    if ax is not None:
        ax.set_facecolor(BG_COLOR)
        ax.tick_params(colors=TEXT_COLOR)
        ax.xaxis.label.set_color(TEXT_COLOR)
        ax.yaxis.label.set_color(TEXT_COLOR)
        ax.title.set_color(TEXT_COLOR)
        for spine in ax.spines.values():
            spine.set_color(GRID_COLOR)
        ax.grid(True, color=GRID_COLOR, alpha=0.5, linestyle="--")


def plot_convergence(
    iteration_history: Sequence[float],
    loss_history: Sequence[float] | None = None,
    title: str = "TQCE Convergence",
) -> Figure:
    """Plot fixed-point inconsistency (and optionally PPGD loss) over iterations.

    Args:
        iteration_history: Total state delta at each fixed-point iteration.
        loss_history: Optional PPGD loss values to overlay.
        title: Plot title.

    Returns:
        Matplotlib Figure.
    """
    fig, ax1 = plt.subplots(figsize=(10, 6))
    _apply_dark_theme(fig, ax1)

    iters = range(len(iteration_history))
    ax1.semilogy(iters, iteration_history, color=CONVERGE_COLOR, linewidth=2,
                 label="State Inconsistency (dS)")
    ax1.set_xlabel("Iteration")
    ax1.set_ylabel("Inconsistency (log scale)", color=CONVERGE_COLOR)

    if loss_history is not None:
        ax2 = ax1.twinx()
        ax2.set_facecolor("none")
        loss_iters = range(len(loss_history))
        ax2.semilogy(loss_iters, loss_history, color=LOSS_COLOR, linewidth=2,
                     linestyle="--", label="PPGD Loss")
        ax2.set_ylabel("PPGD Loss (log scale)", color=LOSS_COLOR)
        ax2.tick_params(axis="y", colors=LOSS_COLOR)
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc="upper right",
                   facecolor=BG_COLOR, edgecolor=GRID_COLOR, labelcolor=TEXT_COLOR)
    else:
        ax1.legend(facecolor=BG_COLOR, edgecolor=GRID_COLOR, labelcolor=TEXT_COLOR)

    ax1.set_title(title, fontsize=14, fontweight="bold")
    fig.tight_layout()
    return fig


def plot_paradox_pressure_landscape(
    pressures: dict[str, float],
    title: str = "Paradox Pressure Map",
) -> Figure:
    """Bar chart of paradox pressure per node.

    Args:
        pressures: Mapping of node_id to paradox pressure value.
        title: Plot title.

    Returns:
        Matplotlib Figure.
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    _apply_dark_theme(fig, ax)

    nodes = list(pressures.keys())
    values = list(pressures.values())

    colors = plt.colormaps[PRESSURE_CMAP](
        np.array(values) / (max(values) + 1e-10)
    )

    bars = ax.bar(nodes, values, color=colors, edgecolor="none", width=0.6)
    ax.set_xlabel("Node")
    ax.set_ylabel("Paradox Pressure")
    ax.set_title(title, fontsize=14, fontweight="bold")

    for bar, val in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                f"{val:.4f}", ha="center", va="bottom", color=TEXT_COLOR,
                fontsize=9)

    fig.tight_layout()
    return fig


def plot_ctn_graph(
    ctn: CausalTensorNetwork,
    title: str = "Causal Tensor Network",
) -> Figure:
    """Visualize the CTN graph with distinct edge types.

    Forward-causal edges are blue with forward arrows.
    Retrocausal edges are red with backward arrows and dashed lines.
    Node size is proportional to paradox pressure.

    Args:
        ctn: The Causal Tensor Network to visualize.
        title: Plot title.

    Returns:
        Matplotlib Figure.
    """
    import networkx as nx

    fig, ax = plt.subplots(figsize=(12, 8))
    _apply_dark_theme(fig, ax)
    ax.grid(False)

    pos = {}
    time_groups: dict[float, list[str]] = {}
    for nid, node in ctn.nodes.items():
        t = node.time_coordinate
        time_groups.setdefault(t, []).append(nid)

    for t, group in time_groups.items():
        for i, nid in enumerate(group):
            y_offset = (i - (len(group) - 1) / 2) * 1.0
            pos[nid] = (t, y_offset)

    base_size = 800
    sizes = [
        base_size + ctn.nodes[nid].paradox_pressure * 2000
        for nid in ctn.graph.nodes()
    ]

    nx.draw_networkx_nodes(
        ctn.graph, pos, ax=ax, node_size=sizes, node_color=CAUSAL_COLOR,
        alpha=0.9, edgecolors=TEXT_COLOR, linewidths=1.5,
    )
    nx.draw_networkx_labels(
        ctn.graph, pos, ax=ax, font_color=TEXT_COLOR, font_weight="bold",
        font_size=11,
    )

    causal_edges = [(u, v) for u, v, d in ctn.graph.edges(data=True)
                    if d.get("type") == "causal"]
    retro_edges = [(u, v) for u, v, d in ctn.graph.edges(data=True)
                   if d.get("type") == "retrocausal"]

    nx.draw_networkx_edges(
        ctn.graph, pos, edgelist=causal_edges, ax=ax,
        edge_color=CAUSAL_COLOR, width=2, arrows=True,
        arrowsize=20, arrowstyle="-|>", connectionstyle="arc3,rad=0.1",
    )
    nx.draw_networkx_edges(
        ctn.graph, pos, edgelist=retro_edges, ax=ax,
        edge_color=RETRO_COLOR, width=2, style="dashed", arrows=True,
        arrowsize=20, arrowstyle="-|>", connectionstyle="arc3,rad=-0.3",
    )

    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color=CAUSAL_COLOR, linewidth=2, label="Causal (past -> future)"),
        Line2D([0], [0], color=RETRO_COLOR, linewidth=2, linestyle="--",
               label="Retrocausal (future -> past)"),
    ]
    ax.legend(handles=legend_elements, loc="upper left", facecolor=BG_COLOR,
              edgecolor=GRID_COLOR, labelcolor=TEXT_COLOR)

    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.set_xlabel("Time Coordinate", color=TEXT_COLOR)
    fig.tight_layout()
    return fig


def plot_tec_probabilities(
    probabilities: NDArray[np.floating],
    title: str = "Timeline Probabilities (TEC)",
) -> Figure:
    """Bar chart of timeline probabilities after causal interference.

    Args:
        probabilities: Array of probabilities for each timeline.
        title: Plot title.

    Returns:
        Matplotlib Figure.
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    _apply_dark_theme(fig, ax)

    n = len(probabilities)
    indices = range(n)
    colors = [TEC_COLOR if p == max(probabilities) else "#475569" for p in probabilities]

    bars = ax.bar(indices, probabilities, color=colors, edgecolor="none", width=0.7)
    ax.set_xlabel("Timeline Index")
    ax.set_ylabel("Probability |alpha|^2")
    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.set_xticks(range(n))

    best = int(np.argmax(probabilities))
    ax.annotate(
        f"COLLAPSED\nP={probabilities[best]:.3f}",
        xy=(best, probabilities[best]),
        xytext=(best + 0.5, probabilities[best] + 0.05),
        arrowprops=dict(arrowstyle="->", color=TEC_COLOR),
        fontsize=10, color=TEC_COLOR, fontweight="bold",
    )

    fig.tight_layout()
    return fig


def plot_amplification_trace(
    amplification_history: Sequence[float],
    target: float | None = None,
    title: str = "Bootstrap Amplification Trace",
) -> Figure:
    """Plot signal amplification factor over loop iterations.

    Args:
        amplification_history: Amplification factor at each iteration.
        target: Optional target amplification line.
        title: Plot title.

    Returns:
        Matplotlib Figure.
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    _apply_dark_theme(fig, ax)

    iters = range(len(amplification_history))
    ax.plot(iters, amplification_history, color=AMP_COLOR, linewidth=2,
            label="Amplification Factor")

    if target is not None:
        ax.axhline(y=target, color=RETRO_COLOR, linestyle="--", linewidth=1.5,
                   label=f"Target ({target}x)")

    ax.axhline(y=1.0, color=GRID_COLOR, linestyle=":", linewidth=1,
               label="No amplification (1x)")

    ax.set_xlabel("Loop Iteration")
    ax.set_ylabel("Amplification Factor")
    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.legend(facecolor=BG_COLOR, edgecolor=GRID_COLOR, labelcolor=TEXT_COLOR)
    fig.tight_layout()
    return fig


def animate_convergence(
    state_histories: Sequence[Sequence[float]],
    interval: int = 100,
    title: str = "Convergence Animation",
) -> animation.FuncAnimation:
    """Create an animated plot of convergence over time.

    Each entry in state_histories is one node's inconsistency trace.

    Args:
        state_histories: List of per-node inconsistency sequences.
        interval: Milliseconds between animation frames.
        title: Plot title.

    Returns:
        Matplotlib FuncAnimation object.
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    _apply_dark_theme(fig, ax)

    max_len = max(len(h) for h in state_histories)
    lines = []
    colors = plt.colormaps["viridis"](
        np.linspace(0, 1, len(state_histories))
    )

    for i, history in enumerate(state_histories):
        (line,) = ax.plot([], [], color=colors[i], linewidth=1.5)
        lines.append(line)

    ax.set_xlim(0, max_len)
    all_vals = [v for h in state_histories for v in h]
    ax.set_ylim(min(all_vals) * 0.9, max(all_vals) * 1.1)
    ax.set_xlabel("Iteration")
    ax.set_ylabel("Inconsistency")
    ax.set_title(title, fontsize=14, fontweight="bold")

    def init():
        for line in lines:
            line.set_data([], [])
        return lines

    def update(frame):
        for i, (line, history) in enumerate(zip(lines, state_histories)):
            end = min(frame + 1, len(history))
            line.set_data(range(end), history[:end])
        return lines

    anim = animation.FuncAnimation(
        fig, update, init_func=init, frames=max_len,
        interval=interval, blit=True,
    )
    return anim
