import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

try:
    from scpn_fusion.core.stability_mhd import (
        QProfile,
        MercierResult,
        BallooningResult,
        KruskalShafranovResult,
        TroyonResult,
        NTMResult,
        StabilitySummary,
        compute_q_profile,
        mercier_stability,
        ballooning_stability,
        kruskal_shafranov_stability,
        troyon_beta_limit,
        ntm_stability,
        run_full_stability_check,
    )
    print("Imported scpn_fusion.core.stability_mhd successfully.")
except ImportError as exc:
    raise ImportError(
        "Could not import stability_mhd.  Make sure scpn-fusion-core is installed:\n"
        "    pip install -e '.[dev]'\n"
        f"Original error: {exc}"
    ) from exc

# Plotting defaults
plt.rcParams.update({
    "figure.dpi": 120,
    "axes.grid": True,
    "grid.alpha": 0.3,
    "font.size": 11,
})
print("Matplotlib backend:", plt.get_backend())

# ------------------------------------------------------------------
# ITER-like equilibrium parameters
# ------------------------------------------------------------------
N = 200
rho = np.linspace(0.0, 1.0, N)

# Plasma profiles (peaked parabolic shapes)
ne = 10.0 * (1.0 - 0.8 * rho**2)          # [10^19 m^-3]
Ti = 20.0 * (1.0 - rho**2)**1.5            # [keV]
Te = 18.0 * (1.0 - rho**2)**1.5            # [keV]

# Machine parameters
R0    = 6.2    # major radius [m]
a     = 2.0    # minor radius [m]
B0    = 5.3    # toroidal field on axis [T]
Ip_MA = 15.0   # plasma current [MA]

# Compute q-profile
qp = compute_q_profile(rho, ne, Ti, Te, R0, a, B0, Ip_MA)

print(f"q on axis    : {qp.q[0]:.3f}")
print(f"q_min        : {qp.q_min:.3f}  at rho = {qp.q_min_rho:.3f}")
print(f"q_edge       : {qp.q_edge:.3f}")
print(f"Shear at edge: {qp.shear[-1]:.3f}")

# ------------------------------------------------------------------
# Plot: q(rho) and shear(rho)
# ------------------------------------------------------------------
fig, ax1 = plt.subplots(figsize=(8, 4.5))

color_q = "#2563eb"
ax1.plot(qp.rho, qp.q, color=color_q, lw=2, label="$q(\\rho)$")
ax1.axhline(1.0, color="grey", ls="--", lw=0.8, alpha=0.6)
ax1.axhline(2.0, color="grey", ls=":", lw=0.8, alpha=0.4)
ax1.set_xlabel(r"$\rho$ (normalised radius)")
ax1.set_ylabel(r"$q(\rho)$", color=color_q)
ax1.tick_params(axis="y", labelcolor=color_q)
ax1.set_ylim(0, max(qp.q) * 1.15)

# Overlay shear on twin axis
ax2 = ax1.twinx()
color_s = "#dc2626"
ax2.plot(qp.rho, qp.shear, color=color_s, lw=1.5, ls="--", label="$s(\\rho)$")
ax2.set_ylabel(r"Magnetic shear $s$", color=color_s)
ax2.tick_params(axis="y", labelcolor=color_s)

# Combined legend
lines1, labels1 = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax1.legend(lines1 + lines2, labels1 + labels2, loc="upper left")

ax1.set_title("Safety Factor and Magnetic Shear (ITER-like)")
plt.tight_layout()
plt.show()

# Run Mercier stability check
mr = mercier_stability(qp)

if mr.first_unstable_rho is not None:
    print(f"Mercier UNSTABLE: first violation at rho = {mr.first_unstable_rho:.3f}")
else:
    print("Mercier STABLE across entire profile.")

# ------------------------------------------------------------------
# Plot: D_M vs rho with unstable regions shaded
# ------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(8, 4.5))

ax.plot(mr.rho, mr.D_M, color="#2563eb", lw=2, label=r"$D_M(\rho)$")
ax.axhline(0.0, color="black", lw=1, ls="-")

# Shade unstable regions (D_M < 0)
unstable_mask = mr.D_M < 0
if np.any(unstable_mask):
    ax.fill_between(
        mr.rho, mr.D_M, 0,
        where=unstable_mask,
        color="#dc2626", alpha=0.25,
        label="Unstable ($D_M < 0$)",
    )

# Mark first unstable point
if mr.first_unstable_rho is not None:
    idx = np.argmin(np.abs(mr.rho - mr.first_unstable_rho))
    ax.plot(mr.first_unstable_rho, mr.D_M[idx], "rv", ms=10, zorder=5,
            label=f"First unstable ($\\rho={mr.first_unstable_rho:.2f}$)")

ax.set_xlabel(r"$\rho$ (normalised radius)")
ax.set_ylabel(r"$D_M$")
ax.set_title("Mercier Interchange Stability")
ax.legend(loc="best")
plt.tight_layout()
plt.show()

# Run ballooning stability check
br = ballooning_stability(qp)

n_unstable = int(np.sum(~br.stable))
print(f"Ballooning: {n_unstable}/{len(br.rho)} radial points unstable")

# ------------------------------------------------------------------
# Plot: s-alpha diagram with stability boundary
# ------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(7, 6))

# Draw the stability boundary
s_bnd = np.linspace(0, max(br.s) * 1.1, 300)
alpha_bnd = np.where(s_bnd < 1.0, s_bnd * (1.0 - s_bnd / 2.0), 0.6 * s_bnd)
ax.plot(s_bnd, alpha_bnd, "k-", lw=2, label=r"$\alpha_{\mathrm{crit}}(s)$")
ax.fill_between(s_bnd, alpha_bnd, max(alpha_bnd) * 1.3,
                color="#dc2626", alpha=0.08, label="Unstable region")

# Scatter: operating points coloured by stability
stable_pts = br.stable
ax.scatter(br.s[stable_pts], br.alpha[stable_pts],
           c="#16a34a", s=12, alpha=0.7, label="Stable", zorder=3)
if n_unstable > 0:
    ax.scatter(br.s[~stable_pts], br.alpha[~stable_pts],
               c="#dc2626", s=18, marker="x", alpha=0.9,
               label="Unstable", zorder=4)

# Annotate selected radii
for rho_label in [0.2, 0.4, 0.6, 0.8]:
    idx = np.argmin(np.abs(br.rho - rho_label))
    ax.annotate(f"$\\rho$={rho_label:.1f}",
                (br.s[idx], br.alpha[idx]),
                textcoords="offset points", xytext=(6, 6),
                fontsize=8, color="#475569")

ax.set_xlabel(r"Magnetic shear $s$")
ax.set_ylabel(r"Normalised pressure gradient $\alpha$")
ax.set_title(r"$s$-$\alpha$ Ballooning Stability Diagram")
ax.set_xlim(left=0)
ax.set_ylim(bottom=0)
ax.legend(loc="upper left")
plt.tight_layout()
plt.show()

# Run Kruskal-Shafranov check
ks = kruskal_shafranov_stability(qp)

status = "STABLE" if ks.stable else "UNSTABLE"
print(f"Kruskal-Shafranov: {status}")
print(f"  q_edge = {ks.q_edge:.3f}")
print(f"  margin = q_edge - 1 = {ks.margin:.3f}")

# ------------------------------------------------------------------
# Visual: bar showing q_edge relative to threshold
# ------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(6, 2))

bar_color = "#16a34a" if ks.stable else "#dc2626"
ax.barh(["$q_{\\mathrm{edge}}$"], [ks.q_edge], color=bar_color, height=0.4, alpha=0.8)
ax.axvline(1.0, color="#dc2626", lw=2, ls="--", label="KS limit ($q=1$)")
ax.set_xlim(0, max(ks.q_edge * 1.3, 2.0))
ax.set_xlabel(r"$q_{\mathrm{edge}}$")
ax.set_title(f"Kruskal-Shafranov: {status} (margin = {ks.margin:.2f})")
ax.legend(loc="lower right", fontsize=9)
plt.tight_layout()
plt.show()

# Sweep beta_t and evaluate Troyon limit
beta_t_values = np.linspace(0.005, 0.06, 50)
beta_N_arr = np.empty_like(beta_t_values)
margin_nw_arr = np.empty_like(beta_t_values)
margin_w_arr = np.empty_like(beta_t_values)
stable_nw_arr = np.empty(len(beta_t_values), dtype=bool)

for i, bt in enumerate(beta_t_values):
    tr = troyon_beta_limit(bt, Ip_MA, a, B0)
    beta_N_arr[i] = tr.beta_N
    margin_nw_arr[i] = tr.margin_nowall
    margin_w_arr[i] = tr.beta_N_crit_wall - tr.beta_N
    stable_nw_arr[i] = tr.stable_nowall

# Print the reference point at beta_t = 0.025 (ITER design)
tr_ref = troyon_beta_limit(0.025, Ip_MA, a, B0)
print(f"ITER reference point (beta_t = 2.5%):")
print(f"  beta_N          = {tr_ref.beta_N:.2f} [% m T / MA]")
print(f"  No-wall limit   = {tr_ref.beta_N_crit_nowall:.1f}  -> {'STABLE' if tr_ref.stable_nowall else 'UNSTABLE'}")
print(f"  Ideal-wall limit= {tr_ref.beta_N_crit_wall:.1f}  -> {'STABLE' if tr_ref.stable_wall else 'UNSTABLE'}")
print(f"  No-wall margin  = {tr_ref.margin_nowall:.2f}")

# ------------------------------------------------------------------
# Plot: stability margins vs beta_N
# ------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(8, 4.5))

ax.plot(beta_N_arr, margin_nw_arr, color="#2563eb", lw=2,
        label=r"No-wall margin ($g=2.8$)")
ax.plot(beta_N_arr, margin_w_arr, color="#7c3aed", lw=2, ls="--",
        label=r"Ideal-wall margin ($g=3.5$)")
ax.axhline(0.0, color="black", lw=1)
ax.axvline(2.8, color="#2563eb", lw=1, ls=":", alpha=0.5)
ax.axvline(3.5, color="#7c3aed", lw=1, ls=":", alpha=0.5)

# Shade unstable region (no-wall)
ax.fill_between(beta_N_arr, margin_nw_arr, 0,
                where=(margin_nw_arr < 0),
                color="#dc2626", alpha=0.15, label="No-wall unstable")

# Mark ITER reference
ax.plot(tr_ref.beta_N, tr_ref.margin_nowall, "ko", ms=8, zorder=5,
        label=f"ITER design ($\\beta_N$={tr_ref.beta_N:.2f})")

ax.set_xlabel(r"$\beta_N$ [% m T / MA]")
ax.set_ylabel("Stability margin")
ax.set_title("Troyon Beta Limit: Stability Margin vs $\\beta_N$")
ax.legend(loc="upper right", fontsize=9)
plt.tight_layout()
plt.show()

# ------------------------------------------------------------------
# Construct realistic current density profiles for NTM analysis
# ------------------------------------------------------------------

# Total current density: parabolic profile (consistent with q-profile)
mu0 = 4.0 * np.pi * 1e-7
j_total = (2.0 * Ip_MA * 1e6 / (np.pi * a**2)) * (1.0 - rho**2)   # [A/m^2]
j_total = np.maximum(j_total, 1e3)  # floor to avoid division issues at edge

# Bootstrap current: peaked off-axis (typical of H-mode)
# j_bs ~ f_bs * j_total * rho * exp(-rho^2)
f_bs = 0.35  # bootstrap fraction at peak
j_bs = f_bs * j_total[0] * rho * np.exp(-2.0 * rho**2)

# Run NTM analysis
ntm = ntm_stability(qp, j_bs, j_total, a, r_s_delta_prime=-2.0)

n_ntm_unstable = int(np.sum(ntm.ntm_unstable))
print(f"NTM: {n_ntm_unstable}/{len(ntm.rho)} radial points with positive marginal width")
if ntm.most_unstable_rho is not None:
    print(f"Most unstable location: rho = {ntm.most_unstable_rho:.3f}")
    idx_peak = np.argmin(np.abs(ntm.rho - ntm.most_unstable_rho))
    print(f"  w_marginal = {ntm.w_marginal[idx_peak]*100:.2f} cm")
    print(f"  j_bs/j_total = {ntm.j_bs_drive[idx_peak]:.3f}")
else:
    print("No NTM-unstable regions found.")

# ------------------------------------------------------------------
# Plot: marginal island width and bootstrap drive vs rho
# ------------------------------------------------------------------
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 7), sharex=True)

# Panel 1: Marginal island width
ax1.plot(ntm.rho, ntm.w_marginal * 100, color="#2563eb", lw=2,
         label=r"$w_{\mathrm{marg}}$")
if np.any(ntm.ntm_unstable):
    ax1.fill_between(
        ntm.rho, 0, ntm.w_marginal * 100,
        where=ntm.ntm_unstable,
        color="#dc2626", alpha=0.2, label="NTM-unstable",
    )
if ntm.most_unstable_rho is not None:
    ax1.axvline(ntm.most_unstable_rho, color="#dc2626", ls=":", lw=1, alpha=0.6)
ax1.set_ylabel(r"$w_{\mathrm{marg}}$ [cm]")
ax1.set_title("NTM Marginal Island Width")
ax1.legend(loc="upper right", fontsize=9)

# Panel 2: Bootstrap current fraction
ax2.plot(ntm.rho, ntm.j_bs_drive, color="#7c3aed", lw=2,
         label=r"$j_{\mathrm{bs}} / j_\phi$")
ax2.set_xlabel(r"$\rho$ (normalised radius)")
ax2.set_ylabel(r"Bootstrap fraction $j_{\mathrm{bs}}/j_\phi$")
ax2.set_title("Bootstrap Current Drive")
ax2.legend(loc="upper right", fontsize=9)

plt.tight_layout()
plt.show()

# Run the full five-criterion suite
summary = run_full_stability_check(
    qp,
    beta_t=0.025,
    Ip_MA=Ip_MA,
    a=a,
    B0=B0,
    j_bs=j_bs,
    j_total=j_total,
)

# ------------------------------------------------------------------
# Formatted summary table
# ------------------------------------------------------------------
def _status(ok: bool) -> str:
    return "STABLE" if ok else "UNSTABLE"

# Per-criterion status
mercier_ok = summary.mercier.first_unstable_rho is None
ballooning_ok = bool(np.all(summary.ballooning.stable))
ks_ok = summary.kruskal_shafranov.stable
troyon_ok = summary.troyon.stable_nowall if summary.troyon else None
ntm_ok = not bool(np.any(summary.ntm.ntm_unstable)) if summary.ntm else None

print("=" * 64)
print("  MHD STABILITY SUMMARY  (ITER-like, beta_t = 2.5%)")
print("=" * 64)
print(f"{'Criterion':<30} {'Status':<12} {'Detail'}")
print("-" * 64)

# 1. Mercier
detail_m = (
    f"first violation at rho={summary.mercier.first_unstable_rho:.3f}"
    if summary.mercier.first_unstable_rho is not None
    else "D_M >= 0 everywhere"
)
print(f"{'1. Mercier':<30} {_status(mercier_ok):<12} {detail_m}")

# 2. Ballooning
n_unst_b = int(np.sum(~summary.ballooning.stable))
detail_b = f"{n_unst_b} unstable points" if n_unst_b > 0 else "all points below boundary"
print(f"{'2. Ballooning':<30} {_status(ballooning_ok):<12} {detail_b}")

# 3. Kruskal-Shafranov
detail_ks = f"q_edge={summary.kruskal_shafranov.q_edge:.3f}, margin={summary.kruskal_shafranov.margin:.3f}"
print(f"{'3. Kruskal-Shafranov':<30} {_status(ks_ok):<12} {detail_ks}")

# 4. Troyon
if summary.troyon is not None:
    detail_t = f"beta_N={summary.troyon.beta_N:.2f}, limit={summary.troyon.beta_N_crit_nowall:.1f}, margin={summary.troyon.margin_nowall:.2f}"
    print(f"{'4. Troyon (no-wall)':<30} {_status(troyon_ok):<12} {detail_t}")
else:
    print(f"{'4. Troyon':<30} {'SKIPPED':<12} (beta_t not provided)")

# 5. NTM
if summary.ntm is not None:
    n_unst_n = int(np.sum(summary.ntm.ntm_unstable))
    detail_n = (
        f"most unstable at rho={summary.ntm.most_unstable_rho:.3f}"
        if summary.ntm.most_unstable_rho is not None
        else "no NTM-unstable points"
    )
    print(f"{'5. NTM':<30} {_status(ntm_ok):<12} {detail_n}")
else:
    print(f"{'5. NTM':<30} {'SKIPPED':<12} (j_bs not provided)")

print("-" * 64)
overall_str = _status(summary.overall_stable)
print(f"{'OVERALL':<30} {overall_str:<12} {summary.n_criteria_stable}/{summary.n_criteria_checked} criteria satisfied")
print("=" * 64)

%%timeit -n1 -r3
run_full_stability_check(
    qp, beta_t=0.025, Ip_MA=Ip_MA, a=a, B0=B0,
    j_bs=j_bs, j_total=j_total,
)

#	Criterion	Key Condition	Reference
1	Mercier (interchange)	$D_M \ge 0$	Freidberg, Ideal MHD (2014), Ch. 12
2	Ballooning (pressure-driven)	$\alpha < \alpha_{\mathrm{crit}}(s)$	Connor, Hastie & Taylor, PRL 40:396 (1978)
3	Kruskal-Shafranov (external kink)	$q_{\mathrm{edge}} > 1$	Kruskal & Schwarzschild, Proc. R. Soc. A 223:348 (1954)
4	Troyon (beta limit)	$\beta_N < g$	Troyon et al., PPCF 26:209 (1984)
5	NTM (neoclassical tearing)	$w_{\mathrm{marg}} \to 0$	La Haye, Phys. Plasmas 13:055501 (2006)

Boundary condition	Troyon coefficient $g$
No conducting wall	$g \approx 2.8$
Ideal conducting wall	$g \approx 3.5$

08 — MHD Stability Analysis: Five-Criterion Suite¶

Safety-Factor Profile ($q$-profile)¶

Criterion 1: Mercier Interchange Stability¶

Criterion 2: First-Stability Ballooning¶

Criterion 3: Kruskal-Shafranov External Kink¶

Criterion 4: Troyon Beta Limit¶

Criterion 5: NTM Seeding¶

Combined Stability Summary¶

Benchmark¶

Summary and Controller Integration¶

Controller integration¶

References¶