"""Two conducting bus strips (+/-V0) with a floating conducting island
between them, embedded in uniform SiO2. Compute the electrostatic field
with and without the island (floating conductor solved exactly via
superposition), save results for plotting/reporting.

Geometry (microns):
  left strip:  x <= -1.0,        |y| <= 0.05   (Dirichlet -V0)
  right strip: x >=  1.0,        |y| <= 0.05   (Dirichlet +V0)
  island:      -0.5 <= x <= 0.5, |y| <= 0.05   (floating, zero net charge)

The strips are modeled as extending to the edge of a padded solve domain
(approximating x -> +/-infinity); the solve domain is padded well beyond
the 5x5 micron region we actually report/plot, with V=0 on the outer
domain edge as a "far field ground" approximation to open space.
"""
from __future__ import annotations

import time
from pathlib import Path

import numpy as np

from solver import EPS0, build_grid, build_laplacian, e_field, flux_around_box, solve_potential

OUT_DIR = Path(__file__).parent / "data"
OUT_DIR.mkdir(exist_ok=True)

# --- physical / geometric parameters ---
V0 = 1.0  # volts
EPS_R_SIO2 = 3.9
STRIP_INNER_X = 1.0  # strips start at |x| = 1.0 um
ISLAND_HALF_W = 0.5  # island spans x in [-0.5, 0.5] um
STRIP_HALF_H = 0.05  # strip/island half-height, y in [-0.05, 0.05] um

# --- numerical parameters ---
LX_SOLVE, LY_SOLVE = 6.0, 6.0  # padded solve domain half-width (um)
DX = DY = 0.01  # grid spacing (um) = 10 nm
DISPLAY_HALF = 2.5  # crop to central 5x5 um for reporting


def make_masks(X, Y, include_island: bool, island_value: float | None):
    left_strip = X <= -STRIP_INNER_X
    left_strip &= np.abs(Y) <= STRIP_HALF_H
    right_strip = X >= STRIP_INNER_X
    right_strip &= np.abs(Y) <= STRIP_HALF_H
    island = (np.abs(X) <= ISLAND_HALF_W) & (np.abs(Y) <= STRIP_HALF_H)

    nx, ny = X.shape
    outer = np.zeros((nx, ny), dtype=bool)
    outer[0, :] = outer[-1, :] = True
    outer[:, 0] = outer[:, -1] = True

    fixed_value = np.zeros((nx, ny))  # default outer boundary = 0 V (far field ground)
    fixed_value[left_strip] = -V0
    fixed_value[right_strip] = V0

    fixed_mask = outer | left_strip | right_strip
    if include_island:
        fixed_mask = fixed_mask | island
        fixed_value[island] = island_value

    return fixed_mask, fixed_value, left_strip, right_strip, island


def crop(field, x, y, half):
    ix = np.where(np.abs(x) <= half)[0]
    iy = np.where(np.abs(y) <= half)[0]
    return field[np.ix_(ix, iy)], x[ix], y[iy]


def main():
    t0 = time.time()
    x, y, X, Y, nx, ny = build_grid(LX_SOLVE, LY_SOLVE, DX, DY)
    print(f"grid: {nx} x {ny} = {nx*ny:,} nodes", flush=True)
    L = build_laplacian(nx, ny, DX, DY)
    print(f"Laplacian built ({time.time()-t0:.1f}s)", flush=True)

    island_box = (-ISLAND_HALF_W, ISLAND_HALF_W, -STRIP_HALF_H, STRIP_HALF_H)
    buffer = 0.10  # um, contour offset outside island bounding box

    # --- WITH island: floating via 2-subproblem superposition ---
    t1 = time.time()
    mask_A, val_A, left_m, right_m, island_m = make_masks(X, Y, include_island=True, island_value=0.0)
    V_A = solve_potential(nx, ny, mask_A, val_A, L)
    Ex_A, Ey_A = e_field(V_A, DX, DY)
    flux_A = flux_around_box(Ex_A, Ey_A, x, y, island_box, buffer)
    print(f"Problem A solved ({time.time()-t1:.1f}s), flux_A={flux_A:.6e} V (should be ~0 by symmetry)", flush=True)

    t1 = time.time()
    mask_B, val_B, _, _, _ = make_masks(X, Y, include_island=True, island_value=1.0)
    V_B = solve_potential(nx, ny, mask_B, val_B, L)
    Ex_B, Ey_B = e_field(V_B, DX, DY)
    flux_B = flux_around_box(Ex_B, Ey_B, x, y, island_box, buffer)
    print(f"Problem B solved ({time.time()-t1:.1f}s), flux_B={flux_B:.6e} V", flush=True)

    V_island = -flux_A / flux_B
    V_with = V_A + V_island * V_B
    Ex_with = Ex_A + V_island * Ex_B
    Ey_with = Ey_A + V_island * Ey_B
    flux_full_check = flux_around_box(Ex_with, Ey_with, x, y, island_box, buffer)
    Q_A = EPS0 * EPS_R_SIO2 * flux_A
    Q_B = EPS0 * EPS_R_SIO2 * flux_B
    Q_full = EPS0 * EPS_R_SIO2 * flux_full_check

    print(f"floating island potential V_island = {V_island:.6f} V (expect ~0 by symmetry)", flush=True)
    print(f"net charge on island (full solution): {Q_full:.4e} C/m (should be ~0)", flush=True)

    # --- WITHOUT island: single solve, island region left as plain SiO2 ---
    t1 = time.time()
    mask_no, val_no, _, _, _ = make_masks(X, Y, include_island=False, island_value=None)
    V_without = solve_potential(nx, ny, mask_no, val_no, L)
    Ex_without, Ey_without = e_field(V_without, DX, DY)
    print(f"No-island problem solved ({time.time()-t1:.1f}s)", flush=True)

    # --- crop to display window and save ---
    Vw, xc, yc = crop(V_with, x, y, DISPLAY_HALF)
    Exw, _, _ = crop(Ex_with, x, y, DISPLAY_HALF)
    Eyw, _, _ = crop(Ey_with, x, y, DISPLAY_HALF)
    Vwo, _, _ = crop(V_without, x, y, DISPLAY_HALF)
    Exwo, _, _ = crop(Ex_without, x, y, DISPLAY_HALF)
    Eywo, _, _ = crop(Ey_without, x, y, DISPLAY_HALF)

    np.savez(
        OUT_DIR / "results.npz",
        x=xc,
        y=yc,
        V_with=Vw,
        Ex_with=Exw,
        Ey_with=Eyw,
        V_without=Vwo,
        Ex_without=Exwo,
        Ey_without=Eywo,
        V_island=V_island,
        Q_A=Q_A,
        Q_B=Q_B,
        Q_full=Q_full,
        V0=V0,
        eps_r_sio2=EPS_R_SIO2,
        strip_inner_x=STRIP_INNER_X,
        island_half_w=ISLAND_HALF_W,
        strip_half_h=STRIP_HALF_H,
    )
    # also keep full-resolution fields (uncropped) for the edge-charge-density lineouts
    np.savez(
        OUT_DIR / "results_full.npz",
        x=x,
        y=y,
        Ex_with=Ex_with,
        Ey_with=Ey_with,
        Ex_without=Ex_without,
        Ey_without=Ey_without,
    )
    print(f"saved results to {OUT_DIR}, total time {time.time()-t0:.1f}s", flush=True)


if __name__ == "__main__":
    main()