from __future__ import annotations

from typing import Dict, Optional, Tuple

import matplotlib

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

from .cubic_curves import (
    EPS33_ROW,
    LATERAL_ROW,
    MODE_TO_ROW,
    default_x_for_mode,
    evaluate_cubic_no_intercept,
    fit_cubic_no_intercept,
)


def _set_axis_limits(
    ax: plt.Axes,
    y: np.ndarray,
    *,
    force_min: Optional[float] = None,
    force_max: Optional[float] = None,
    min_floor: Optional[float] = None,
    max_ceiling: Optional[float] = None,
    pad_frac: float = 0.2,
) -> None:
    y = np.asarray(y, dtype=np.float32).reshape(-1)
    if y.size == 0:
        return
    y_min = float(np.min(y))
    y_max = float(np.max(y))
    if force_min is not None:
        y_min = min(y_min, float(force_min))
    if force_max is not None:
        y_max = max(y_max, float(force_max))
    if min_floor is not None:
        y_min = min(y_min, float(min_floor))
    if max_ceiling is not None:
        y_max = max(y_max, float(max_ceiling))
    span = y_max - y_min
    pad = span * pad_frac if span > 0 else max(abs(y_max), abs(y_min), 1.0) * 0.1
    bottom = y_min - pad
    top = y_max + pad
    if force_min is not None:
        bottom = float(force_min)
    if force_max is not None:
        top = float(force_max)
    ax.set_ylim(bottom, top)


def plot_single_curve(coeffs: np.ndarray, mode: str, title: str, ylabel: str) -> plt.Figure:
    x = default_x_for_mode(mode)
    y = evaluate_cubic_no_intercept(coeffs, x)
    fig, ax = plt.subplots(figsize=(3.2, 2.5))
    ax.plot(x, y, color="k", linewidth=2)
    ax.set_title(title)
    ax.set_xlabel("strain")
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.25)
    fig.tight_layout()
    return fig


def plot_required_curves(cond_coeffs_7x3: np.ndarray) -> plt.Figure:
    cond = np.asarray(cond_coeffs_7x3, dtype=np.float32).reshape(7, 3)
    fig, axes = plt.subplots(3, 3, figsize=(14, 12))
    modes = ["11", "22", "12"]
    for col, mode in enumerate(modes):
        x = default_x_for_mode(mode)
        stress_row = MODE_TO_ROW[mode]
        axes[0, col].plot(x, evaluate_cubic_no_intercept(cond[stress_row], x), color="k", linewidth=2)
        axes[0, col].set_title(f"Mode {mode} stress")
        axes[0, col].set_xlabel("strain")
        axes[0, col].set_ylabel("stress (MPa)")
        axes[0, col].grid(True, alpha=0.25)

        if mode in ("11", "22"):
            lat_row = int(LATERAL_ROW[mode])
            e33_row = int(EPS33_ROW[mode])
            axes[1, col].plot(x, evaluate_cubic_no_intercept(cond[lat_row], x), color="k", linewidth=2)
            axes[1, col].set_title(f"Mode {mode} lateral")
            axes[1, col].set_xlabel("strain")
            axes[1, col].set_ylabel("lateral strain")
            axes[1, col].grid(True, alpha=0.25)

            axes[2, col].plot(x, evaluate_cubic_no_intercept(cond[e33_row], x), color="k", linewidth=2)
            axes[2, col].set_title(f"Mode {mode} eps33")
            axes[2, col].set_xlabel("strain")
            axes[2, col].set_ylabel("eps33")
            axes[2, col].grid(True, alpha=0.25)
        else:
            axes[1, col].axis("off")
            axes[2, col].axis("off")

    fig.tight_layout()
    return fig


def plot_required_curves_from_series(curves: Dict[str, np.ndarray]) -> plt.Figure:
    fig, axes = plt.subplots(3, 3, figsize=(14, 12))

    axes[0, 0].plot(curves["eps11"], curves["sig11_mpa"], color="k", linewidth=2)
    axes[0, 0].set_title("Mode 11 stress")
    axes[0, 0].set_xlabel("strain")
    axes[0, 0].set_ylabel("stress (MPa)")
    axes[0, 0].grid(True, alpha=0.25)
    _set_axis_limits(axes[0, 0], curves["sig11_mpa"])

    axes[1, 0].plot(curves["eps11"], curves["eps22_from_eps11"], color="k", linewidth=2)
    axes[1, 0].set_title("Mode 11 lateral")
    axes[1, 0].set_xlabel("strain")
    axes[1, 0].set_ylabel("lateral strain")
    axes[1, 0].grid(True, alpha=0.25)
    _set_axis_limits(axes[1, 0], curves["eps22_from_eps11"], min_floor=-0.01, max_ceiling=0.0, pad_frac=0.0)

    axes[2, 0].plot(curves["eps11"], curves["eps33_from_eps11"], color="k", linewidth=2)
    axes[2, 0].set_title("Mode 11 eps33")
    axes[2, 0].set_xlabel("strain")
    axes[2, 0].set_ylabel("eps33")
    axes[2, 0].grid(True, alpha=0.25)
    _set_axis_limits(axes[2, 0], curves["eps33_from_eps11"])

    axes[0, 1].plot(curves["eps22"], curves["sig22_mpa"], color="k", linewidth=2)
    axes[0, 1].set_title("Mode 22 stress")
    axes[0, 1].set_xlabel("strain")
    axes[0, 1].set_ylabel("stress (MPa)")
    axes[0, 1].grid(True, alpha=0.25)
    _set_axis_limits(axes[0, 1], curves["sig22_mpa"])

    axes[1, 1].plot(curves["eps22"], curves["eps11_from_eps22"], color="k", linewidth=2)
    axes[1, 1].set_title("Mode 22 lateral")
    axes[1, 1].set_xlabel("strain")
    axes[1, 1].set_ylabel("lateral strain")
    axes[1, 1].grid(True, alpha=0.25)
    _set_axis_limits(axes[1, 1], curves["eps11_from_eps22"], min_floor=-0.01, max_ceiling=0.0, pad_frac=0.0)

    axes[2, 1].plot(curves["eps22"], curves["eps33_from_eps22"], color="k", linewidth=2)
    axes[2, 1].set_title("Mode 22 eps33")
    axes[2, 1].set_xlabel("strain")
    axes[2, 1].set_ylabel("eps33")
    axes[2, 1].grid(True, alpha=0.25)
    _set_axis_limits(axes[2, 1], curves["eps33_from_eps22"])

    axes[0, 2].plot(curves["eps12"], curves["sig12_mpa"], color="k", linewidth=2)
    axes[0, 2].set_title("Mode 12 stress")
    axes[0, 2].set_xlabel("strain")
    axes[0, 2].set_ylabel("stress (MPa)")
    axes[0, 2].grid(True, alpha=0.25)
    _set_axis_limits(axes[0, 2], curves["sig12_mpa"])
    axes[1, 2].axis("off")
    axes[2, 2].axis("off")

    fig.tight_layout()
    return fig


def plot_condition_and_simulations_v5(
    cond_coeffs_7x3: np.ndarray,
    simulations_by_design: Dict[int, Dict[str, Dict[str, np.ndarray]]],
) -> plt.Figure:
    cond = np.asarray(cond_coeffs_7x3, dtype=np.float32).reshape(7, 3)
    fig, axes = plt.subplots(3, 3, figsize=(14, 12))
    modes = ["11", "22", "12"]
    colors = ["b", "c", "m", "g", "r"]

    for col, mode in enumerate(modes):
        x_cond = default_x_for_mode(mode)
        stress_ys = [evaluate_cubic_no_intercept(cond[MODE_TO_ROW[mode]], x_cond)]
        lateral_ys = []
        eps33_ys = []
        x_min, x_max = None, None
        for _, sim in sorted(simulations_by_design.items()):
            if mode in sim:
                x = sim[mode]["strain"]
                x_min = float(np.min(x))
                x_max = float(np.max(x))
                break
        stress_row = MODE_TO_ROW[mode]
        axes[0, col].plot(
            x_cond,
            evaluate_cubic_no_intercept(cond[stress_row], x_cond),
            color="k",
            linewidth=2,
            label="Condition",
        )

        if mode in ("11", "22"):
            lat_row = int(LATERAL_ROW[mode])
            e33_row = int(EPS33_ROW[mode])
            lateral_ys.append(evaluate_cubic_no_intercept(cond[lat_row], x_cond))
            eps33_ys.append(evaluate_cubic_no_intercept(cond[e33_row], x_cond))
            axes[1, col].plot(
                x_cond,
                lateral_ys[-1],
                color="k",
                linewidth=2,
                label="Condition",
            )
            axes[2, col].plot(
                x_cond,
                eps33_ys[-1],
                color="k",
                linewidth=2,
                label="Condition",
            )
        else:
            axes[1, col].axis("off")
            axes[2, col].axis("off")

        for idx, (design_idx, sim) in enumerate(sorted(simulations_by_design.items(), key=lambda kv: kv[0])):
            if mode not in sim:
                continue
            color = colors[idx % len(colors)]
            data = sim[mode]
            strain = np.asarray(data["strain"], dtype=np.float32)
            stress = np.asarray(data["stress"], dtype=np.float32)
            stress_ys.append(stress)
            axes[0, col].plot([], [], color=color, linewidth=2, label=f"Gen {design_idx}")
            axes[0, col].scatter(strain, stress, color=color, s=22, alpha=0.85, label="_nolegend_")
            stress_coeffs = fit_cubic_no_intercept(strain, stress, degree=3)
            if stress_coeffs is not None:
                x_sim = np.linspace(float(np.min(strain)), float(np.max(strain)), 200, dtype=np.float32)
                axes[0, col].plot(x_sim, evaluate_cubic_no_intercept(stress_coeffs, x_sim), color=color, linewidth=2, alpha=0.9, label="_nolegend_")

            if mode in ("11", "22") and data.get("lateral") is not None:
                lateral = np.asarray(data["lateral"], dtype=np.float32)
                lateral_ys.append(lateral)
                axes[1, col].plot([], [], color=color, linewidth=2, label=f"Gen {design_idx}")
                axes[1, col].scatter(strain, lateral, color=color, s=22, alpha=0.85, label="_nolegend_")
                lat_coeffs = fit_cubic_no_intercept(strain, lateral, degree=3)
                if lat_coeffs is not None:
                    axes[1, col].plot(x_sim, evaluate_cubic_no_intercept(lat_coeffs, x_sim), color=color, linewidth=2, alpha=0.9, label="_nolegend_")

            if mode in ("11", "22") and data.get("eps33") is not None:
                eps33 = np.asarray(data["eps33"], dtype=np.float32)
                eps33_ys.append(eps33)
                axes[2, col].plot([], [], color=color, linewidth=2, label=f"Gen {design_idx}")
                axes[2, col].scatter(strain, eps33, color=color, s=22, alpha=0.85, label="_nolegend_")
                e33_coeffs = fit_cubic_no_intercept(strain, eps33, degree=3)
                if e33_coeffs is not None:
                    axes[2, col].plot(x_sim, evaluate_cubic_no_intercept(e33_coeffs, x_sim), color=color, linewidth=2, alpha=0.9, label="_nolegend_")

        if x_min is None or x_max is None or x_max <= x_min:
            x_min = float(np.min(x_cond))
            x_max = float(np.max(x_cond))
        else:
            x_min = min(float(np.min(x_cond)), x_min)
            x_max = max(float(np.max(x_cond)), x_max)
        axes[0, col].set_xlim(x_min, x_max)
        if mode in ("11", "22"):
            axes[1, col].set_xlim(x_min, x_max)
            axes[2, col].set_xlim(x_min, x_max)

        axes[0, col].set_title(f"Mode {mode} stress")
        axes[0, col].set_xlabel("strain")
        axes[0, col].set_ylabel("stress (MPa)")
        axes[0, col].grid(True, alpha=0.25)
        axes[0, col].legend(fontsize=8)
        _set_axis_limits(axes[0, col], np.concatenate(stress_ys))

        if mode in ("11", "22"):
            axes[1, col].set_title(f"Mode {mode} lateral")
            axes[1, col].set_xlabel("strain")
            axes[1, col].set_ylabel("lateral strain")
            axes[1, col].grid(True, alpha=0.25)
            axes[1, col].legend(fontsize=8)
            _set_axis_limits(axes[1, col], np.concatenate(lateral_ys), min_floor=-0.01, max_ceiling=0.0, pad_frac=0.0)

            axes[2, col].set_title(f"Mode {mode} eps33")
            axes[2, col].set_xlabel("strain")
            axes[2, col].set_ylabel("eps33")
            axes[2, col].grid(True, alpha=0.25)
            axes[2, col].legend(fontsize=8)
            _set_axis_limits(axes[2, col], np.concatenate(eps33_ys))

    fig.tight_layout()
    return fig