# Cell 2: Imports
import sys
import json
import tempfile
import os

import numpy as np

try:
    import matplotlib
    import matplotlib.pyplot as plt
    from matplotlib.patches import Rectangle
    %matplotlib inline
    plt.rcParams.update({
        'figure.dpi': 100,
        'font.size': 11,
        'axes.titlesize': 13,
    })
    HAS_MPL = True
except ImportError:
    HAS_MPL = False
    print("WARNING: matplotlib not found. Plots will be skipped.")

try:
    from scpn_fusion.core.fusion_kernel import FusionKernel, CoilSet
    HAS_FUSION = True
except ImportError:
    HAS_FUSION = False
    print("WARNING: scpn_fusion not importable. ")
    print("Install with:  pip install -e '.[dev]'  from the SCPN-Fusion-Core root.")

print(f"numpy {np.__version__}")
if HAS_MPL:
    print(f"matplotlib {matplotlib.__version__}")
if HAS_FUSION:
    print("scpn_fusion.core.fusion_kernel loaded OK")

# Cell 4: Define CoilSet with 6 PF coils

# -- Reactor configuration (ITER-like) --
R0 = 6.2       # Major radius [m]
a  = 2.0        # Minor radius [m]
Ip = 15.0       # Plasma current [MA-turns normalised units]

# Coil positions (R, Z) in metres and initial currents [A]
coil_positions = [
    (3.5,  4.0),    # PF1 - inner upper
    (3.5, -4.0),    # PF2 - inner lower
    (9.0,  4.0),    # PF3 - outer upper
    (9.0, -4.0),    # PF4 - outer lower
    (6.2,  5.5),    # PF5 - vertical upper
    (6.2, -5.5),    # PF6 - vertical lower
]

initial_currents = np.array([5.0, 5.0, -3.0, -3.0, -1.5, -1.5])
coil_turns = [200, 200, 150, 150, 100, 100]

# Engineering current limits (absolute) [A]
current_limits = np.array([8.0, 8.0, 6.0, 6.0, 4.0, 4.0])

# Build the CoilSet dataclass
coils = CoilSet(
    positions=coil_positions,
    currents=initial_currents.copy(),
    turns=coil_turns,
    current_limits=current_limits,
    target_flux_points=None,   # set later for optimization
)

print("CoilSet defined:")
for i, (pos, I, n, lim) in enumerate(
    zip(coils.positions, coils.currents, coils.turns, coils.current_limits)
):
    print(f"  PF{i+1}: R={pos[0]:.1f} m, Z={pos[1]:+.1f} m, "
          f"I={I:+.2f} A, n={n}, |I|_max={lim:.1f} A")

# Cell 6: Fixed-boundary equilibrium solve

# Build the FusionKernel config (matches the config JSON format)
config_fixed = {
    "reactor_name": "ITER-like (fixed-boundary)",
    "grid_resolution": [65, 65],
    "dimensions": {
        "R_min": 1.0, "R_max": 9.0,
        "Z_min": -5.0, "Z_max": 5.0,
    },
    "physics": {
        "plasma_current_target": Ip,
        "vacuum_permeability": 1.2566370614e-6,
    },
    "coils": [
        {"r": pos[0], "z": pos[1], "current": float(I)}
        for pos, I in zip(coil_positions, initial_currents)
    ],
    "solver": {
        "max_iterations": 200,
        "convergence_threshold": 1e-6,
        "relaxation_factor": 0.1,
        "solver_method": "sor",
    },
}

# Write to a temporary file (FusionKernel expects a JSON path)
_fd, config_fixed_path = tempfile.mkstemp(suffix=".json")
os.close(_fd)
with open(config_fixed_path, 'w') as f:
    json.dump(config_fixed, f)

# Solve
kernel_fixed = FusionKernel(config_fixed_path)
result_fixed = kernel_fixed.solve_equilibrium()

print(f"Fixed-boundary solve:")
print(f"  Converged: {result_fixed['converged']}")
print(f"  Iterations: {result_fixed['iterations']}")
print(f"  Final residual: {result_fixed['residual']:.4e}")
print(f"  Wall time: {result_fixed['wall_time_s']:.2f} s")

# Store grid and solution for plotting
R_grid = kernel_fixed.R
Z_grid = kernel_fixed.Z
RR, ZZ = np.meshgrid(R_grid, Z_grid)
psi_fixed = kernel_fixed.Psi.copy()

# -- Plot fixed-boundary psi contours --
if HAS_MPL:
    fig, ax = plt.subplots(figsize=(6, 8))
    levels = np.linspace(psi_fixed.min(), psi_fixed.max(), 30)
    cs = ax.contour(RR, ZZ, psi_fixed, levels=levels, cmap='RdBu_r')
    ax.set_xlabel('R [m]')
    ax.set_ylabel('Z [m]')
    ax.set_title('Fixed-Boundary Equilibrium $\\Psi$ Contours')
    ax.set_aspect('equal')

    # Mark coil positions
    for i, (pos, I) in enumerate(zip(coil_positions, initial_currents)):
        color = 'red' if I > 0 else 'blue'
        ax.plot(pos[0], pos[1], 's', color=color, markersize=10)
        ax.annotate(f'PF{i+1}', (pos[0]+0.15, pos[1]+0.15), fontsize=8)

    plt.colorbar(cs, ax=ax, label='$\\Psi$ [Wb/rad]')
    plt.tight_layout()
    plt.show()

# Cell 8: Free-boundary solve with initial (fixed) coil currents

# Re-create kernel for the free-boundary solve
config_free = config_fixed.copy()
config_free["reactor_name"] = "ITER-like (free-boundary)"

_fd2, config_free_path = tempfile.mkstemp(suffix=".json")
os.close(_fd2)
with open(config_free_path, 'w') as f:
    json.dump(config_free, f)

kernel_free = FusionKernel(config_free_path)

# Reset coil currents to initial values
coils_fb = CoilSet(
    positions=coil_positions,
    currents=initial_currents.copy(),
    turns=coil_turns,
    current_limits=current_limits,
    target_flux_points=None,
)

fb_result = kernel_free.solve_free_boundary(
    coils_fb,
    max_outer_iter=15,
    tol=1e-4,
    optimize_shape=False,
)

psi_free = kernel_free.Psi.copy()

print(f"Free-boundary solve (fixed coil currents):")
print(f"  Outer iterations: {fb_result['outer_iterations']}")
print(f"  Final max |delta Psi|: {fb_result['final_diff']:.4e}")
print(f"  Coil currents: {fb_result['coil_currents']}")

# Find the X-point and separatrix
(R_x, Z_x), psi_x = kernel_free.find_x_point(psi_free)
print(f"  X-point: R={R_x:.2f} m, Z={Z_x:.2f} m, Psi_x={psi_x:.6f}")

# -- Plot free-boundary psi contours with separatrix --
if HAS_MPL:
    fig, axes = plt.subplots(1, 2, figsize=(13, 8))

    # Left: fixed-boundary reference
    levels_f = np.linspace(psi_fixed.min(), psi_fixed.max(), 30)
    axes[0].contour(RR, ZZ, psi_fixed, levels=levels_f, cmap='RdBu_r', linewidths=0.8)
    axes[0].set_title('Fixed-Boundary')
    axes[0].set_xlabel('R [m]')
    axes[0].set_ylabel('Z [m]')
    axes[0].set_aspect('equal')

    # Right: free-boundary with separatrix
    levels_fb = np.linspace(psi_free.min(), psi_free.max(), 30)
    axes[1].contour(RR, ZZ, psi_free, levels=levels_fb, cmap='RdBu_r', linewidths=0.8)
    # Highlight separatrix (psi = psi at X-point)
    axes[1].contour(RR, ZZ, psi_free, levels=[psi_x], colors='lime',
                    linewidths=2.5, linestyles='--')
    axes[1].plot(R_x, Z_x, 'x', color='lime', markersize=14, markeredgewidth=3)
    axes[1].set_title('Free-Boundary (separatrix in green)')
    axes[1].set_xlabel('R [m]')
    axes[1].set_ylabel('Z [m]')
    axes[1].set_aspect('equal')

    # Mark coils on both panels
    for ax_i in axes:
        for i, (pos, I) in enumerate(zip(coil_positions, initial_currents)):
            color = 'red' if I > 0 else 'blue'
            ax_i.plot(pos[0], pos[1], 's', color=color, markersize=9)
            ax_i.annotate(f'PF{i+1}', (pos[0]+0.15, pos[1]+0.15), fontsize=7)

    plt.suptitle('Fixed-Boundary vs Free-Boundary Equilibrium', fontsize=14)
    plt.tight_layout()
    plt.show()

# Cell 10: Define target separatrix points and run coil optimization

# Define target flux points on a D-shaped ellipse
# (R0 +/- a*cos(theta), kappa*a*sin(theta)) with elongation kappa
n_target = 32
theta_target = np.linspace(0, 2 * np.pi, n_target, endpoint=False)
kappa = 1.7       # elongation
delta_tri = 0.33   # triangularity

# Parameterised D-shape (Miller geometry)
R_target = R0 + a * np.cos(theta_target + delta_tri * np.sin(theta_target))
Z_target = kappa * a * np.sin(theta_target)

target_flux_points = np.column_stack([R_target, Z_target])

# The target flux value at each point = psi at the X-point from the
# initial free-boundary solve (all separatrix points should have the same psi)
target_psi_value = psi_x
target_psi = np.full(n_target, target_psi_value)

# Set target flux points on the coil set
coils_opt = CoilSet(
    positions=coil_positions,
    currents=initial_currents.copy(),
    turns=coil_turns,
    current_limits=current_limits,
    target_flux_points=target_flux_points,
)

print(f"Target separatrix: {n_target} control points")
print(f"  R range: [{R_target.min():.2f}, {R_target.max():.2f}] m")
print(f"  Z range: [{Z_target.min():.2f}, {Z_target.max():.2f}] m")
print(f"  Target psi: {target_psi_value:.6f} Wb/rad")
print(f"  Tikhonov lambda: 1e-3")
print()

# Run the optimization
optimized_currents = kernel_free.optimize_coil_currents(
    coils_opt,
    target_flux=target_psi,
    tikhonov_alpha=1e-3,
)

print("Optimization results:")
print(f"{'Coil':<6} {'Initial [A]':>12} {'Optimized [A]':>14} {'Limit [A]':>10} {'At limit?':>10}")
print("-" * 56)
for i in range(len(coil_positions)):
    at_lim = abs(abs(optimized_currents[i]) - current_limits[i]) < 0.01
    print(f"PF{i+1:<4} {initial_currents[i]:>+12.4f} {optimized_currents[i]:>+14.4f} "
          f"{current_limits[i]:>10.1f} {'  YES' if at_lim else '  no':>10}")

%%timeit -n1 -r3
_coils = CoilSet(
    positions=coil_positions,
    currents=initial_currents.copy(),
    turns=coil_turns,
    current_limits=current_limits,
    target_flux_points=target_flux_points,
)
kernel_free.optimize_coil_currents(_coils, target_flux=target_psi, tikhonov_alpha=1e-3)

# Cell 12: Manual outer iteration with convergence tracking

# Fresh kernel for the self-consistent loop
config_loop = config_fixed.copy()
config_loop["reactor_name"] = "ITER-like (self-consistent)"

_fd3, config_loop_path = tempfile.mkstemp(suffix=".json")
os.close(_fd3)
with open(config_loop_path, 'w') as f:
    json.dump(config_loop, f)

kernel_loop = FusionKernel(config_loop_path)

# Coil set with target points
coils_loop = CoilSet(
    positions=coil_positions,
    currents=initial_currents.copy(),
    turns=coil_turns,
    current_limits=current_limits,
    target_flux_points=target_flux_points,
)

n_outer = 10
psi_residuals = []
current_history = [coils_loop.currents.copy()]
tikhonov_lambda = 1e-3

# Initial external flux from coils
psi_ext = kernel_loop._compute_external_flux(coils_loop)

for outer in range(n_outer):
    # Step 1: Apply external flux as boundary condition
    kernel_loop.Psi[0, :]  = psi_ext[0, :]
    kernel_loop.Psi[-1, :] = psi_ext[-1, :]
    kernel_loop.Psi[:, 0]  = psi_ext[:, 0]
    kernel_loop.Psi[:, -1] = psi_ext[:, -1]

    # Step 2: Inner GS solve
    psi_old = kernel_loop.Psi.copy()
    kernel_loop.solve_equilibrium(
        preserve_initial_state=True,
        boundary_flux=psi_ext,
    )

    # Step 3: Extract current flux at target points
    actual_psi_at_targets = np.array([
        kernel_loop._interp_psi(R_t, Z_t)
        for R_t, Z_t in target_flux_points
    ])

    # Step 4: Optimize coil currents
    new_currents = kernel_loop.optimize_coil_currents(
        coils_loop,
        target_psi,
        tikhonov_alpha=tikhonov_lambda,
    )
    coils_loop.currents = new_currents
    current_history.append(new_currents.copy())

    # Step 5: Recompute external flux
    psi_ext = kernel_loop._compute_external_flux(coils_loop)

    # Track convergence
    diff = float(np.max(np.abs(kernel_loop.Psi - psi_old)))
    psi_residuals.append(diff)

    print(f"  Outer iter {outer+1:2d}: max|dPsi|={diff:.4e}, "
          f"I=[{', '.join(f'{c:+.3f}' for c in new_currents)}]")

    if diff < 1e-5:
        print(f"  Converged at outer iteration {outer+1}.")
        break

psi_optimized = kernel_loop.Psi.copy()

# -- Plot convergence --
if HAS_MPL:
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.semilogy(range(1, len(psi_residuals) + 1), psi_residuals, 'o-',
                color='navy', linewidth=2, markersize=7)
    ax.set_xlabel('Outer Iteration')
    ax.set_ylabel('max $|\\Delta\\Psi|$')
    ax.set_title('GS-Coil Self-Consistency: Convergence of $\\Psi$ Residual')
    ax.grid(True, alpha=0.3)
    ax.set_xticks(range(1, len(psi_residuals) + 1))
    plt.tight_layout()
    plt.show()

# Cell 14: Bar chart - initial vs optimized coil currents with limits

final_currents = coils_loop.currents.copy()
coil_labels = [f'PF{i+1}' for i in range(len(coil_positions))]
x = np.arange(len(coil_labels))
bar_width = 0.35

if HAS_MPL:
    fig, ax = plt.subplots(figsize=(10, 5))

    bars_init = ax.bar(x - bar_width/2, initial_currents, bar_width,
                       label='Initial', color='steelblue', alpha=0.8,
                       edgecolor='navy')
    bars_opt  = ax.bar(x + bar_width/2, final_currents, bar_width,
                       label='Optimized', color='coral', alpha=0.8,
                       edgecolor='darkred')

    # Draw current limits as horizontal lines spanning each coil
    for i in range(len(coil_labels)):
        ax.plot([x[i] - 0.45, x[i] + 0.45], [+current_limits[i]]*2,
                'k--', linewidth=1.5, alpha=0.6)
        ax.plot([x[i] - 0.45, x[i] + 0.45], [-current_limits[i]]*2,
                'k--', linewidth=1.5, alpha=0.6)

    # Add a single legend entry for the limits
    ax.plot([], [], 'k--', linewidth=1.5, alpha=0.6, label='Current limit')

    ax.set_xlabel('Coil')
    ax.set_ylabel('Current [A]')
    ax.set_title('Initial vs Optimized Coil Currents')
    ax.set_xticks(x)
    ax.set_xticklabels(coil_labels)
    ax.legend(loc='upper right')
    ax.axhline(y=0, color='gray', linewidth=0.5)
    ax.grid(True, axis='y', alpha=0.3)

    # Annotate values on bars
    for bar in bars_init:
        h = bar.get_height()
        ax.annotate(f'{h:+.2f}',
                    xy=(bar.get_x() + bar.get_width()/2, h),
                    xytext=(0, 3 if h >= 0 else -12),
                    textcoords='offset points',
                    ha='center', fontsize=8)
    for bar in bars_opt:
        h = bar.get_height()
        ax.annotate(f'{h:+.2f}',
                    xy=(bar.get_x() + bar.get_width()/2, h),
                    xytext=(0, 3 if h >= 0 else -12),
                    textcoords='offset points',
                    ha='center', fontsize=8)

    plt.tight_layout()
    plt.show()

# Print summary table
print()
print(f"{'Coil':<6} {'Initial [A]':>12} {'Optimized [A]':>14} {'Delta [A]':>10} "
      f"{'Limit [A]':>10} {'Usage %':>8}")
print("-" * 64)
for i in range(len(coil_positions)):
    delta = final_currents[i] - initial_currents[i]
    usage = abs(final_currents[i]) / current_limits[i] * 100
    print(f"PF{i+1:<4} {initial_currents[i]:>+12.4f} {final_currents[i]:>+14.4f} "
          f"{delta:>+10.4f} {current_limits[i]:>10.1f} {usage:>7.1f}%")

# Cleanup temporary config files
for path in [config_fixed_path, config_free_path, config_loop_path]:
    try:
        os.unlink(path)
    except OSError:
        pass

print("Temporary files cleaned up.")

	Fixed-Boundary	Free-Boundary
Boundary condition	$\Psi = 0$ on a prescribed contour (e.g., last closed flux surface)	$\Psi_{\text{boundary}} = \Psi_{\text{ext}}$ from coils
Plasma shape	Known a priori	Emerges self-consistently
Use case	Profile studies, initial design	Shape control, coil optimization, experiment matching
Iteration	GS PDE only	GS PDE + boundary update + (optional) coil optimization

Class/Method	Purpose
`CoilSet`	Dataclass holding coil geometry, currents, limits, and target points
`FusionKernel.solve_equilibrium()`	Fixed-boundary Picard-iteration GS solver
`FusionKernel.solve_free_boundary()`	Free-boundary solver with optional shape optimization
`FusionKernel.optimize_coil_currents()`	Bounded least-squares with Tikhonov regularization

09 — Coil Current Optimization for Free-Boundary Equilibria¶

Physics Overview¶

Coil Current Optimization¶

Setting Up the Coil Set¶

Fixed-Boundary vs Free-Boundary¶

Free-Boundary Solve with Fixed Coil Currents¶

Coil Current Optimization¶

Benchmark¶

Outer Iteration: GS-Coil Self-Consistency¶

Optimized Coil Currents Summary¶

Summary¶

Key Classes and Methods¶

References¶

Limitations¶

Connection to Real-Time Shape Control¶

Coil	R [m]	Z [m]	Role	Turns
PF1	3.5	+4.0	Inner upper shaping	200
PF2	3.5	-4.0	Inner lower shaping	200
PF3	9.0	+4.0	Outer upper shaping	150
PF4	9.0	-4.0	Outer lower shaping	150
PF5	6.2	+5.5	Vertical field	100
PF6	6.2	-5.5	Vertical field	100