import gdsfactory as gf
from ihp import LAYER, PDK

from gsim.palace import DrivenSim

PDK.activate()


@gf.cell
def gsg_electrode(
    length=300, s_width=20, g_width=40, gap_width=15, layer=LAYER.TopMetal2drawing
):
    c = gf.Component()
    r1 = c << gf.c.rectangle((length, g_width), centered=True, layer=layer)
    r1.move((0, (g_width + s_width) / 2 + gap_width))
    _r2 = c << gf.c.rectangle((length, s_width), centered=True, layer=layer)
    r3 = c << gf.c.rectangle((length, g_width), centered=True, layer=layer)
    r3.move((0, -(g_width + s_width) / 2 - gap_width))
    c.add_port(
        name="o1",
        center=(-length / 2, 0),
        width=s_width,
        orientation=0,
        port_type="electrical",
        layer=layer,
    )
    c.add_port(
        name="o2",
        center=(length / 2, 0),
        width=s_width,
        orientation=180,
        port_type="electrical",
        layer=layer,
    )
    return c


sim = DrivenSim()
sim.set_output_dir("./palace-sim-cpw-fields")
sim.set_geometry(gsg_electrode())
sim.set_stack(substrate_thickness=2.0, air_above=300.0)
sim.add_cpw_port("o1", layer="topmetal2", s_width=20, gap_width=15)
sim.add_cpw_port("o2", layer="topmetal2", s_width=20, gap_width=15)

# Single frequency point at 50 GHz, adaptive off so Palace does a full solve
sim.set_driven(
    fmin=50e9,
    fmax=50e9 + 1e6,  # tiny range = effectively one point
    num_points=1,
    adaptive_tol=0,
    save_step=1,
)
sim.mesh(preset="fine", planar_conductors=False)
results = sim.run()

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pyvista as pv
from matplotlib.colors import LogNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.interpolate import griddata

pv.OFF_SCREEN = True

# Get results dir from sim output (or hardcode for re-runs)
results_dir = Path(results["port-S.csv"]).parent
print(f"Results dir: {results_dir}")

# Read frequency from S-parameter CSV (first data row, first column)
s_csv = np.loadtxt(results_dir / "port-S.csv", delimiter=",", skiprows=1)
freq_ghz = s_csv[0, 0]

vol_dir = results_dir / "paraview/driven/excitation_1"
bnd_dir = results_dir / "paraview/driven_boundary/excitation_1"

# Use the last .pvtu (final solve iteration)
vol_path = sorted(vol_dir.rglob("*.pvtu"))[-1]
bnd_path = sorted(bnd_dir.rglob("*.pvtu"))[-1]

vol = pv.read(str(vol_path))
bnd = pv.read(str(bnd_path))

print(f"Frequency: {freq_ghz:.1f} GHz")
print(f"Volume: {vol.n_points:,} points, {vol.n_cells:,} cells")
print(f"Boundary: {bnd.n_points:,} points, {bnd.n_cells:,} cells")

# Slice volume at conductor top, clip boundary to conductor region
z_conductor = 16.3
vol_slice = vol.slice(normal="z", origin=(0, 0, z_conductor))
bnd_clip = bnd.clip(normal="z", origin=(0, 0, 17), invert=True)

# Fit interpolation grid to the actual data bounds
vpts_all = vol_slice.points
x_pad, y_pad = 5, 5
xi = np.linspace(vpts_all[:, 0].min() - x_pad, vpts_all[:, 0].max() + x_pad, 500)
yi = np.linspace(vpts_all[:, 1].min() - y_pad, vpts_all[:, 1].max() + y_pad, 150)
Xi, Yi = np.meshgrid(xi, yi)


def interp2d(points, values, grid_x, grid_y, method="linear"):
    return griddata(points, values, (grid_x, grid_y), method=method)


def _norm(gi, log=False):
    if log:
        vmin = np.nanpercentile(gi[gi > 0], 2) if np.any(gi > 0) else 1e-10
        return LogNorm(vmin=vmin, vmax=np.nanpercentile(gi, 98))
    return None  # use default linear


def plot_topview(pts, data, title, cmap="turbo", method="linear", log=False):
    gi = interp2d(pts[:, :2], data, Xi, Yi, method=method)
    fig, ax = plt.subplots(figsize=(14, 3.5))
    norm = _norm(gi, log)
    vmax = np.nanpercentile(gi, 98)
    im = ax.pcolormesh(
        Xi,
        Yi,
        gi,
        cmap=cmap,
        shading="auto",
        norm=norm,
        **({"vmin": 0, "vmax": vmax} if not log else {}),
    )
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02)
    ax.set_title(title)
    ax.set_aspect("equal")
    ax.set_xlabel("x (µm)")
    ax.set_ylabel("y (µm)")
    valid = ~np.isnan(gi) & (gi > 0) if method == "nearest" else ~np.isnan(gi)
    if valid.any():
        rows = np.any(valid, axis=1)
        cols = np.any(valid, axis=0)
        ax.set_xlim(xi[cols][0], xi[cols][-1])
        ax.set_ylim(yi[rows][0], yi[rows][-1])
    fig.tight_layout(pad=0.5)
    plt.show()


def plot_cross_section(yz_pts, field, title, label, zi_range=(-5, 50), log=False):
    mag = np.linalg.norm(field, axis=1)
    y_pad = 5
    yi_cs = np.linspace(yz_pts[:, 1].min() - y_pad, yz_pts[:, 1].max() + y_pad, 200)
    zi_cs = np.linspace(*zi_range, 100)
    Yi_cs, Zi_cs = np.meshgrid(yi_cs, zi_cs)

    mag_cs = interp2d(yz_pts[:, 1:3], mag, Yi_cs, Zi_cs)
    Fy_cs = interp2d(yz_pts[:, 1:3], field[:, 1], Yi_cs, Zi_cs)
    Fz_cs = interp2d(yz_pts[:, 1:3], field[:, 2], Yi_cs, Zi_cs)

    fig, ax = plt.subplots(figsize=(12, 5))
    norm = _norm(mag_cs, log)
    vmax = np.nanpercentile(mag_cs, 98)
    im = ax.pcolormesh(
        Yi_cs,
        Zi_cs,
        mag_cs,
        cmap="turbo",
        shading="auto",
        norm=norm,
        **({"vmin": 0, "vmax": vmax} if not log else {}),
    )

    skip = 8
    ax.quiver(
        Yi_cs[::skip, ::skip],
        Zi_cs[::skip, ::skip],
        Fy_cs[::skip, ::skip],
        Fz_cs[::skip, ::skip],
        color="white",
        alpha=0.7,
        scale=vmax * 15,
        width=0.002,
    )
    ax.set_title(title)
    ax.set_xlabel("y (µm)")
    ax.set_ylabel("z (µm)")
    ax.set_aspect("equal")
    valid = ~np.isnan(mag_cs)
    if valid.any():
        rows = np.any(valid, axis=1)
        cols = np.any(valid, axis=0)
        ax.set_xlim(yi_cs[cols][0], yi_cs[cols][-1])
        ax.set_ylim(zi_cs[rows][0], zi_cs[rows][-1])
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="2%", pad=0.1)
    fig.colorbar(im, cax=cax, label=label)
    fig.tight_layout(pad=0.5)
    plt.show()


vpts = vol_slice.points
bpts = bnd_clip.points

plot_topview(
    vpts,
    np.linalg.norm(vol_slice.point_data["E_real"], axis=1),
    f"Electric field |E| at {freq_ghz:.1f} GHz — field concentrated in CPW gaps (V/m)",
)

plot_topview(
    vpts,
    np.linalg.norm(vol_slice.point_data["S"], axis=1),
    f"Poynting vector |S| at {freq_ghz:.1f} GHz — power flow along the waveguide (W/m²)",
)

plot_topview(
    bpts,
    np.linalg.norm(bnd_clip.point_data["J_s_real"], axis=1),
    f"Surface current |J_s| at {freq_ghz:.1f} GHz — current crowding at conductor edges (A/m)",
    method="nearest",
)

def plot_topview_component(pts, data, title, cmap="RdBu_r", method="linear"):
    """Top-view plot for signed field components with symmetric diverging colormap."""
    gi = interp2d(pts[:, :2], data, Xi, Yi, method=method)
    fig, ax = plt.subplots(figsize=(14, 3.5))
    vlim = np.nanpercentile(np.abs(gi), 98)
    im = ax.pcolormesh(Xi, Yi, gi, cmap=cmap, shading="auto", vmin=-vlim, vmax=vlim)
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02)
    ax.set_title(title)
    ax.set_aspect("equal")
    ax.set_xlabel("x (µm)")
    ax.set_ylabel("y (µm)")
    valid = ~np.isnan(gi)
    if valid.any():
        rows = np.any(valid, axis=1)
        cols = np.any(valid, axis=0)
        ax.set_xlim(xi[cols][0], xi[cols][-1])
        ax.set_ylim(yi[rows][0], yi[rows][-1])
    fig.tight_layout(pad=0.5)
    plt.show()

plot_topview_component(
    vpts,
    vol_slice.point_data["E_real"][:, 1],
    f"E_y at {freq_ghz:.1f} GHz — transverse field across CPW gaps (V/m)",
)

yz_slice = vol.slice(normal="x", origin=(0, 0, 0))
yz_pts = yz_slice.points

plot_cross_section(
    yz_pts,
    yz_slice.point_data["E_real"],
    f"E-field cross-section at {freq_ghz:.1f} GHz — field lines from signal to ground",
    "|E| (V/m)",
)

plot_cross_section(
    yz_pts,
    yz_slice.point_data["B_real"],
    f"H-field cross-section at {freq_ghz:.1f} GHz — magnetic field circulating around conductor",
    "|B| (T)",
)

# B_z component cross-section — non-zero B_z indicates departure from pure TEM mode
B_z = yz_slice.point_data["B_real"][:, 2]

y_pad = 5
yi_cs = np.linspace(yz_pts[:, 1].min() - y_pad, yz_pts[:, 1].max() + y_pad, 200)
zi_cs = np.linspace(-5, 50, 100)
Yi_cs, Zi_cs = np.meshgrid(yi_cs, zi_cs)

Bz_cs = interp2d(yz_pts[:, 1:3], B_z, Yi_cs, Zi_cs)

fig, ax = plt.subplots(figsize=(12, 5))
vlim = np.nanpercentile(np.abs(Bz_cs), 98)
im = ax.pcolormesh(
    Yi_cs, Zi_cs, Bz_cs, cmap="RdBu_r", shading="auto", vmin=-vlim, vmax=vlim
)
ax.set_title(
    f"B_z cross-section at {freq_ghz:.1f} GHz — non-TEM indicator (should be ~0 for pure TEM)"
)
ax.set_xlabel("y (µm)")
ax.set_ylabel("z (µm)")
ax.set_aspect("equal")
valid = ~np.isnan(Bz_cs)
if valid.any():
    rows = np.any(valid, axis=1)
    cols = np.any(valid, axis=0)
    ax.set_xlim(yi_cs[cols][0], yi_cs[cols][-1])
    ax.set_ylim(zi_cs[rows][0], zi_cs[rows][-1])
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="2%", pad=0.1)
fig.colorbar(im, cax=cax, label="B_z (T)")
fig.tight_layout(pad=0.5)
plt.show()