"""Matplotlib visualization functions for the Gradio demo.

Four figure types:
1. Prediction comparison bar chart (neural vs analytical)
2. Deformed shape diagram with exaggerated deflection
3. Safety factor gauge
4. Error bar plot with uncertainty bands
"""

import io

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np


def create_comparison_chart(
    neural_stress: float,
    analytical_stress: float,
    neural_deflection: float,
    analytical_deflection: float,
    stress_ci: tuple[float, float] | None = None,
    deflection_ci: tuple[float, float] | None = None,
) -> plt.Figure:
    """Bar chart comparing neural prediction vs analytical solution."""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

    # Stress comparison
    bars1 = ax1.bar(
        ["Neural", "Analytical"],
        [neural_stress / 1e6, analytical_stress / 1e6],
        color=["#4A90D9", "#7CB342"],
        width=0.5,
    )
    if stress_ci:
        ax1.errorbar(
            0, neural_stress / 1e6,
            yerr=[[neural_stress / 1e6 - stress_ci[0] / 1e6],
                  [stress_ci[1] / 1e6 - neural_stress / 1e6]],
            fmt="none", color="black", capsize=5,
        )
    error_pct = abs(neural_stress - analytical_stress) / analytical_stress * 100
    ax1.set_title(f"Max Stress (Error: {error_pct:.2f}%)")
    ax1.set_ylabel("Stress [MPa]")
    ax1.bar_label(bars1, fmt="%.1f")

    # Deflection comparison
    bars2 = ax2.bar(
        ["Neural", "Analytical"],
        [neural_deflection * 1e3, analytical_deflection * 1e3],
        color=["#4A90D9", "#7CB342"],
        width=0.5,
    )
    if deflection_ci:
        ax2.errorbar(
            0, neural_deflection * 1e3,
            yerr=[[neural_deflection * 1e3 - deflection_ci[0] * 1e3],
                  [deflection_ci[1] * 1e3 - neural_deflection * 1e3]],
            fmt="none", color="black", capsize=5,
        )
    error_pct_d = abs(neural_deflection - analytical_deflection) / analytical_deflection * 100
    ax2.set_title(f"Max Deflection (Error: {error_pct_d:.2f}%)")
    ax2.set_ylabel("Deflection [mm]")
    ax2.bar_label(bars2, fmt="%.3f")

    fig.tight_layout()
    return fig


def create_beam_deformation(
    length: float,
    height: float,
    max_deflection: float,
    config_id: str,
    n_points: int = 100,
) -> plt.Figure:
    """Beam deformation diagram with exaggerated deflection."""
    fig, ax = plt.subplots(figsize=(10, 4))

    x = np.linspace(0, length, n_points)

    # Deflection shape depends on config
    if "ss" in config_id and "point" in config_id:
        # Simply supported, central point: cubic segments
        mid = length / 2
        y = np.where(
            x <= mid,
            max_deflection * (3 * length * x**2 - 4 * x**3) / length**3,
            max_deflection * (3 * length * (length - x)**2 - 4 * (length - x)**3) / length**3,
        )
    elif "ss" in config_id and "udl" in config_id:
        # Simply supported, UDL: quartic
        y = max_deflection * 16 * x * (length - x) * (length**2 + x * (length - x) - x * (length - x)) / (5 * length**4)
        # Simplified shape
        y = max_deflection * np.sin(np.pi * x / length)
    elif "cantilever" in config_id and "point" in config_id:
        # Cantilever, tip load: cubic
        y = max_deflection * (3 * length * x**2 - x**3) / (2 * length**3)
    elif "cantilever" in config_id and "udl" in config_id:
        # Cantilever, UDL: quartic
        y = max_deflection * (6 * length**2 * x**2 - 4 * length * x**3 + x**4) / (3 * length**4)
    elif "fixed" in config_id and "point" in config_id:
        # Fixed-fixed, central point
        mid = length / 2
        y_half = max_deflection * 16 * (x[:n_points//2])**2 * (3 * mid - 2 * x[:n_points//2]) / (3 * length**3) * (3 * length)
        y = np.concatenate([y_half, y_half[::-1]])
        y = y[:n_points]
    else:
        # Default: sinusoidal
        y = max_deflection * np.sin(np.pi * x / length)

    # Scale factor for visibility
    scale = max(length / (20 * max_deflection), 1.0) if max_deflection > 0 else 1.0
    scale = min(scale, 100.0)

    # Undeformed beam
    ax.fill_between(x, -height/2, height/2, alpha=0.15, color="gray", label="Undeformed")
    ax.plot(x, np.full_like(x, height/2), "k--", linewidth=0.5, alpha=0.3)
    ax.plot(x, np.full_like(x, -height/2), "k--", linewidth=0.5, alpha=0.3)

    # Deformed beam (exaggerated)
    y_scaled = -y * scale
    ax.fill_between(x, y_scaled - height/2, y_scaled + height/2, alpha=0.6, color="#4A90D9", label="Deformed")
    ax.plot(x, y_scaled + height/2, "b-", linewidth=1.5)
    ax.plot(x, y_scaled - height/2, "b-", linewidth=1.5)

    ax.set_xlabel("Position along beam [m]")
    ax.set_ylabel("Displacement [m]")
    ax.set_title(f"Deformed Shape (scale: {scale:.0f}x)")
    ax.legend(loc="upper right")
    ax.set_aspect("auto")
    ax.grid(True, alpha=0.3)

    fig.tight_layout()
    return fig


def create_safety_gauge(safety_factor: float) -> plt.Figure:
    """Semicircular gauge showing safety factor value."""
    fig, ax = plt.subplots(figsize=(5, 3), subplot_kw={"projection": "polar"})

    # Gauge range: 0 to 4 (mapped to 0 to pi)
    max_sf = 4.0
    sf_clamped = min(safety_factor, max_sf)
    angle = np.pi * (1 - sf_clamped / max_sf)

    # Color zones
    theta_fail = np.linspace(np.pi, np.pi * (1 - 1.0/max_sf), 50)
    theta_marg = np.linspace(np.pi * (1 - 1.0/max_sf), np.pi * (1 - 2.0/max_sf), 50)
    theta_safe = np.linspace(np.pi * (1 - 2.0/max_sf), 0, 50)

    for theta, color in [(theta_fail, "#EF5350"), (theta_marg, "#FFA726"), (theta_safe, "#66BB6A")]:
        ax.barh(1, np.diff(theta).mean(), left=theta[:-1], height=0.3, color=color, alpha=0.5)

    # Needle
    ax.plot([angle, angle], [0, 0.9], color="black", linewidth=2)
    ax.plot(angle, 0.9, "ko", markersize=8)

    ax.set_ylim(0, 1.3)
    ax.set_thetamin(0)
    ax.set_thetamax(180)
    ax.set_rticks([])
    ax.set_thetagrids([0, 45, 90, 135, 180], ["4.0", "3.0", "2.0", "1.0", "0"])
    ax.set_title(f"Safety Factor: {safety_factor:.2f}", pad=20, fontsize=14, fontweight="bold")

    fig.tight_layout()
    return fig