fem

import numpy as np
import matplotlib
import sympy as sp
from scipy.constants import electron_mass
from shapely import unary_union

#matplotlib.use("Agg")
import matplotlib.pyplot as plt

from shapely.geometry import Polygon, Point, LineString, MultiPolygon
from shapely import contains_xy
from scipy.sparse import lil_matrix, csr_matrix, coo_matrix, block_diag as sp_block_diag
from scipy.sparse.linalg import spsolve
from scipy.spatial import cKDTree
from matplotlib.tri import Triangulation, LinearTriInterpolator
from matplotlib.patches import Polygon as MplPolygon, Patch
from matplotlib.collections import LineCollection

# ============================================================
# LATERAL EDGE-CONTACT FEM (graphene-metal side contact)
# ============================================================
# Two coupled FEM domains:
#   Omega_g — graphene monolayer (2D sheet conductor, Kubo conductivity)
#   Omega_m — adjacent metal contact (Ni, skin-effect surface impedance)
#
# Both domains are meshed and coupled at their shared boundary Gamma_c
# via a Robin interface condition:
#
#   -sigma_g * dV_g/dn = g_edge * (V_g - V_m)   on Gamma_c
#   -sigma_m * dV_m/dn = g_edge * (V_m - V_g)   on Gamma_c
#
# g_edge models the Ni-C bonding layer as a parallel RC+L network:
#   g_edge = 1/(R_bond + jωL_bond) + jωC_bond
#
# BCs: V_g = 1 V at graphene inlet (left edge),
#      V_m = 0 V at metal outer boundary (drain).
# ============================================================
"""
Parameter extraction: channel and overlap (μ_c, τ) from Gahoi TLM + QuantumATK ΔE_F
=====================================================================================
Replace the parameter block and compute_materials() in your FEM script with this.
"""

import numpy as np

# --- units ---
cm, um, nm, m = 0.01, 1e-6, 1e-9, 1
freq = 1e9

# --- fundamental constants ---
k     = 8.617333262145e-5       # eV/K
hbar  = 6.582119569e-16         # eV·s
T     = 300.0                   # K
q     = 1.602176634e-19         # C
vf    = 1e6                     # m/s  (graphene Fermi velocity)
MU0   = 4.0e-7 * np.pi
EPS0  = 8.854187817e-12

# ============================================================
#  EXPERIMENTAL INPUTS  (Gahoi et al., Adv. Electron. Mater. 2020)
# ============================================================
R_SH  = 1104       # Ω/□   channel sheet resistance (Ni devices, vacuum)
R_SK  = 608        # Ω/□   sheet resistance beneath Ni contact (RCE extraction)
n_ch  = 8.69e12 / cm**2   # m⁻²  channel carrier density

# Edge-contact resistance (vacuum, back-gate at Dirac)
R_contact_intrinsic = 853 * um   # Ω·m  (= 853 Ω·μm)

# ============================================================
#  AB INITIO INPUT  (QuantumATK surface band structure)
# ============================================================
# Fermi level shift of graphene beneath Ni, extracted from the
# graphene-projected surface band structure of the Ni/graphene slab.
#
# Ni is ferromagnetic: the exchange splitting pushes the majority
# and minority Dirac points to opposite sides of E_F.
#   Spin up   (majority):  E_D = -0.808 eV  → n-doped, |μ_↑| = 0.808 eV
#   Spin down (minority):  E_D = +0.209 eV  → p-doped, |μ_↓| = 0.209 eV
#
# The scalar Drude formula σ = (q²/πℏ²)μ_c τ includes a factor of 2
# for spin degeneracy.  When degeneracy is broken:
#   σ_total = (q²/2πℏ²) τ (|μ_↑| + |μ_↓|)
# Equating to the scalar form gives the effective chemical potential:
#   μ_c,eff = (|μ_↑| + |μ_↓|) / 2

mu_up   = 0.808    # eV  — majority spin, n-doped (E_D below E_F)
mu_down = 0.209    # eV  — minority spin, p-doped (E_D above E_F)
mu_c_SK_abinitio = (mu_up + mu_down) / 2   # = 0.509 eV
Delta_EF = mu_c_SK_abinitio - hbar * vf * np.sqrt(np.pi * n_ch)  # ≈ 0.165 eV

# ============================================================
#  DERIVED: chemical potential and τ in each region
# ============================================================

# --- Channel ---
mu_c_ch = hbar * vf * np.sqrt(np.pi * n_ch)   # eV  (should ≈ 0.344 eV)
tau_ch  = (np.pi * hbar**2) / (R_SH * mu_c_ch * q)   # seconds

# --- Overlap zone (Ni-doped) ---
mu_c_SK = mu_c_SK_abinitio                            # eV  (= 0.509 eV from spin avg)
n_SK    = mu_c_SK**2 / (np.pi * (hbar * vf)**2)       # m⁻²  (consistency check)
tau_SK  = (np.pi * hbar**2) / (R_SK * mu_c_SK * q)    # seconds

print("="*60)
print("CHANNEL")
print(f"  n_ch      = {n_ch*cm**2:.2e} cm⁻²")
print(f"  μ_c,ch    = {mu_c_ch:.4f} eV")
print(f"  τ_ch      = {tau_ch*1e12:.4f} ps")
print(f"  R_SH      = {R_SH} Ω/□  (input)")
print(f"  R_SH chk  = {np.pi*hbar**2/(q*mu_c_ch*tau_ch):.1f} Ω/□  (round-trip)")
print()
print("AB INITIO (QuantumATK spin-resolved)")
print(f"  |μ_↑|     = {mu_up:.3f} eV  (majority, n-doped)")
print(f"  |μ_↓|     = {mu_down:.3f} eV  (minority, p-doped)")
print(f"  μ_c,SK    = (|μ_↑| + |μ_↓|)/2 = {mu_c_SK:.4f} eV")
print(f"  ΔE_F      = μ_c,SK − μ_c,ch = {Delta_EF:.4f} eV")
print()
print("OVERLAP (Ni-doped)")
print(f"  n_SK      = {n_SK*cm**2:.2e} cm⁻²")
print(f"  τ_SK      = {tau_SK*1e12:.4f} ps")
print(f"  R_SK      = {R_SK} Ω/□  (input)")
print(f"  R_SK chk  = {np.pi*hbar**2/(q*mu_c_SK*tau_SK):.1f} Ω/□  (round-trip)")
print(f"  τ_SK/τ_ch = {tau_SK/tau_ch:.2f}")
print("="*60)


# ============================================================
#  GRAPHENE MATERIAL (thickness, permeability)
# ============================================================
d_g   = 0.34 * nm       # monolayer thickness
mu_r_g = 1.0

# ============================================================
#  NICKEL
# ============================================================
mu_r_m     = 200.0       # DC μ_r (drops to ~1 at GHz; see note in code)
eps_r_m    = 1.0
d_m        = 100 * nm
rho_nickel = 7.2e-8
R_sheet_m  = rho_nickel / d_m

# ============================================================
#  Ni–C INTERFACE BOND  (for AC model)
# ============================================================
ni_c_bondlen = 0.211 * nm
C_interface  = EPS0 * d_g / ni_c_bondlen
L_contact    = MU0 * ni_c_bondlen

# ============================================================
#  OVERLAP WIDTH
# ============================================================
L_overlap = 50 * nm    # incidental metal-on-graphene overlap

# ============================================================
#  GEOMETRY PARAMETERS
# ============================================================
W_interconnect = 100 * nm
contact_length = 1200 * nm
pad_width      = 500 * nm
taper_length   = 250 * nm
metal_gap_viz  = 35 * nm
mesh_h         = 5 * nm
eps_sample     = 2 * nm

# ============================================================
#  KUBO CONDUCTIVITY  (needs F1 integral, computed once)
# ============================================================
import sympy as sp
from scipy.integrate import quad

a = sp.symbols("a")

def _f_d(a_sym, mu):
    """Fermi-Dirac distribution (symbolic) for chemical potential mu [eV]."""
    return (sp.exp((a_sym - mu) / (k * T)) + 1)**-1


# F1 intraband integral — depends on μ_c but NOT on τ or ω
# We precompute for both regions.
def _compute_F1(mu):
    fd = _f_d(a, mu)
    integrand = a * (fd.diff(a) - _f_d(-a, mu).diff(a)) / q
    val, _ = quad(sp.lambdify(a, integrand, "mpmath"), 0, np.inf)
    return complex(val)

F1_ch = _compute_F1(mu_c_ch)
F1_SK = _compute_F1(mu_c_SK)


def _kubo_sheet_sigma(freq, mu_c, tau_val, F1_val):
    """Full Kubo sheet conductance σ_sh [S/□] at frequency freq [Hz]."""
    omega = 2.0 * np.pi * freq
    scattering = 1.0 / (2 * tau_val)
    fd_mu = _f_d(a, mu_c)
    F2, _ = quad(
        sp.lambdify(a,
                    (_f_d(-a, mu_c) - fd_mu)
                    / ((omega + 1j * 2 * scattering)**2 - 4 * (a / hbar)**2)
                    / q,
                    "mpmath"),
        0, np.inf, complex_func=True
    )
    F2 = complex(F2)
    eps_g = (EPS0
             - q**2 / (np.pi * hbar**2 * (omega + 1j*2*scattering) * omega * d_g) * F1_val
             + q**2 * (omega + 1j*2*scattering) / (np.pi * hbar**2 * omega * d_g) * F2)
    return 1j * omega * eps_g * d_g


def compute_materials(freq):
    """Return (sigma_channel, sigma_overlap, sigma_m, g_edge) at freq [Hz].

    At DC (freq→0), sigma_channel → 1/R_SH and sigma_overlap → 1/R_SK
    by construction, since τ was extracted to reproduce these values.
    """
    omega = 2.0 * np.pi * freq

    s_ch = _kubo_sheet_sigma(freq, mu_c_ch, tau_ch, F1_ch)
    s_SK = _kubo_sheet_sigma(freq, mu_c_SK, tau_SK, F1_SK)

    # Metal surface impedance with skin effect
    gamma_m = np.sqrt(1j * omega * MU0 * mu_r_m / rho_nickel)
    Z_m = rho_nickel * gamma_m / np.tanh(gamma_m * d_m)
    s_m = (1.0 / Z_m if abs(Z_m) > 0 else 0.0) + 1j * omega * EPS0 * eps_r_m * d_m

    # Edge contact: parallel RC+L bond model
    g = 1.0 / (R_contact_intrinsic + 1j * omega * L_contact) + 1j * omega * C_interface

    return s_ch, s_SK, s_m, g


# Module-level defaults computed at the default freq
#sigma_s, sigma_m, g_edge = compute_materials(freq)

# ---------- helpers ----------
def graphics_scale(a):
    return np.asarray(a) * 1e9

def fmt_ohm(val):
    if val >= 1000:
        return f"{val/1000:.2f} kΩ"
    return f"{val:.1f} Ω"

def build_structured_tri_mesh(poly, h):
    minx, miny, maxx, maxy = poly.bounds
    xs = np.arange(minx, maxx + 0.5*h, h)
    ys = np.arange(miny, maxy + 0.5*h, h)
    nx, ny = len(xs), len(ys)
    X, Y = np.meshgrid(xs, ys, indexing="xy")
    pts = np.column_stack([X.ravel(), Y.ravel()])

    # Vectorized triangle index generation
    I, J = np.meshgrid(np.arange(nx - 1), np.arange(ny - 1), indexing="xy")
    I, J = I.ravel(), J.ravel()
    n00 = J * nx + I;       n10 = J * nx + (I + 1)
    n01 = (J+1) * nx + I;  n11 = (J+1) * nx + (I + 1)
    tris = np.empty((2 * len(I), 3), dtype=int)
    tris[0::2] = np.stack([n00, n10, n11], axis=1)
    tris[1::2] = np.stack([n00, n11, n01], axis=1)

    # Vectorized containment via shapely 2.x bulk op (avoids per-centroid buffer/Point calls)
    centroids = pts[tris].mean(axis=1)
    buffered = poly.buffer(1e-9 * h)
    inside = contains_xy(buffered, centroids[:, 0], centroids[:, 1])
    tris = tris[inside]

    used = np.unique(tris)
    remap = -np.ones(len(pts), dtype=int)
    remap[used] = np.arange(len(used))
    pts = pts[used]
    tris = remap[tris]
    return pts, tris

# NOTE: triangle_area, add_robin_edges, and total_contact_current_edges
# are retained for reference but are not called by the two-domain solver.
def triangle_area(xy):
    x1, y1 = xy[0]
    x2, y2 = xy[1]
    x3, y3 = xy[2]
    return 0.5 * abs((x2 - x1)*(y3 - y1) - (x3 - x1)*(y2 - y1))


def assemble_bulk_matrix(pts, tris, sigma_per_element):
    """Assemble FEM stiffness for per-element conductivities.

    sigma_per_element : array of shape (n_tris,), complex
        Sheet conductance for each triangle.
    """
    n = len(pts)
    xy = pts[tris]
    x = xy[:, :, 0];
    y = xy[:, :, 1]
    area = 0.5 * np.abs((x[:, 1] - x[:, 0]) * (y[:, 2] - y[:, 0]) - (x[:, 2] - x[:, 0]) * (y[:, 1] - y[:, 0]))
    inv2a = 0.5 / area
    beta = np.stack([y[:, 1] - y[:, 2], y[:, 2] - y[:, 0], y[:, 0] - y[:, 1]], axis=1) * inv2a[:, None]
    gamma = np.stack([x[:, 2] - x[:, 1], x[:, 0] - x[:, 2], x[:, 1] - x[:, 0]], axis=1) * inv2a[:, None]
    # sigma_per_element is (T,) — broadcast into (T,3,3) element matrices
    Ke = sigma_per_element[:, None, None] * area[:, None, None] * (
            beta[:, :, None] * beta[:, None, :] + gamma[:, :, None] * gamma[:, None, :]
    )
    rows = np.repeat(tris, 3, axis=1).ravel()
    cols = np.tile(tris, (1, 3)).ravel()
    from scipy.sparse import coo_matrix
    A = coo_matrix((Ke.ravel(), (rows, cols)), shape=(n, n), dtype=complex).tocsr()
    return A, np.zeros(n, dtype=complex)


def boundary_edges(tris):
    # All 3T directed edges, sorted to canonical (low, high) form
    all_edges = np.sort(np.concatenate([tris[:, [0,1]], tris[:, [1,2]], tris[:, [2,0]]], axis=0), axis=1)
    n_nodes = int(tris.max()) + 1
    codes = all_edges[:, 0].astype(np.int64) * n_nodes + all_edges[:, 1]
    unique_codes, counts = np.unique(codes, return_counts=True)
    bnd = unique_codes[counts == 1]
    return list(zip((bnd // n_nodes).tolist(), (bnd % n_nodes).tolist()))

def add_robin_edges(A, pts, edges, g_edge, edge_selector):
    A = A.tolil(copy=True)
    contact_edges = []
    contact_lengths = []
    for a, b in edges:
        pa, pb = pts[a], pts[b]
        mid = 0.5 * (pa + pb)
        if edge_selector(pa, pb, mid):
            L = float(np.linalg.norm(pb - pa))
            Me = g_edge * L / 6.0 * np.array([[2.0, 1.0],[1.0, 2.0]])
            A[a, a] += Me[0,0]
            A[a, b] += Me[0,1]
            A[b, a] += Me[1,0]
            A[b, b] += Me[1,1]
            contact_edges.append((a,b))
            contact_lengths.append(L)
    return A.tocsr(), contact_edges, np.array(contact_lengths)

def apply_dirichlet_nodes(A, b, nodes, value):
    nodes = np.asarray(nodes, dtype=int)
    b = b.copy()
    # Column elimination via CSC: subtract fixed-node contributions from RHS
    A_csc = A.tocsc().astype(complex, copy=True)
    for node in nodes:
        s, e = A_csc.indptr[node], A_csc.indptr[node+1]
        b[A_csc.indices[s:e]] -= A_csc.data[s:e] * value
        A_csc.data[s:e] = 0.0
    b[nodes] = value
    # Zero fixed rows via CSR
    A_csr = A_csc.tocsr()
    for node in nodes:
        s, e = A_csr.indptr[node], A_csr.indptr[node+1]
        A_csr.data[s:e] = 0.0
    A_csr[nodes, nodes] = 1.0
    return A_csr, b

def reaction_current(A_full, V, dir_nodes):
    r = A_full @ V
    return np.sum(r[dir_nodes])

def total_contact_current_edges(pts, V, contact_edges, g_edge):
    I = 0j
    for a, b in contact_edges:
        L = float(np.linalg.norm(pts[b] - pts[a]))
        I += g_edge * L * 0.5 * (V[a] + V[b])
    return I

def inlet_nodes_for_line(pts, inlet_line, h):
    tol = 0.55 * h
    coords = np.array(inlet_line.coords)
    p0, p1 = coords[0], coords[-1]
    v = p1 - p0
    v_len = np.linalg.norm(v)
    if v_len < 1e-15:
        return np.array([], dtype=int)
    v_hat = v / v_len
    t = (pts - p0) @ v_hat                          # projection parameter along line
    closest = p0 + np.clip(t, 0, v_len)[:, None] * v_hat
    dist = np.linalg.norm(pts - closest, axis=1)
    return np.where(dist <= tol)[0].astype(int)

def sample_line_profile(tri_obj, field, line_pts):
    """Interpolate `field` (nodal values on tri_obj) along line_pts.
    tri_obj is a pre-built matplotlib.tri.Triangulation (cached in prepare_geom)."""
    interp = LinearTriInterpolator(tri_obj, field)
    line_pts = np.asarray(line_pts, dtype=float)
    vals = np.asarray(interp(line_pts[:,0], line_pts[:,1]), dtype=float)
    if np.isnan(vals).any():
        pts = np.column_stack([tri_obj.x, tri_obj.y])
        _, idx = cKDTree(pts).query(line_pts[np.isnan(vals)], k=1)
        vals[np.isnan(vals)] = field[idx]
    ds = np.linalg.norm(np.diff(line_pts, axis=0), axis=1)
    s = np.concatenate([[0.0], np.cumsum(ds)])
    return s, vals

def characteristic_length_1e(s_nm, J):
    J = np.asarray(J)
    if len(J) < 3 or J[0] <= 0:
        return None
    y = J / J[0]
    target = 1.0 / np.e
    idx = np.where(y <= target)[0]
    if len(idx) == 0 or idx[0] == 0:
        return None
    i = idx[0]
    x0, x1 = s_nm[i-1], s_nm[i]
    y0, y1 = y[i-1], y[i]
    return x0 + (target - y0) * (x1 - x0) / (y1 - y0)

def add_poly_patch(ax, shp, **kwargs):
    if isinstance(shp, MultiPolygon):
        for g in shp.geoms:
            add_poly_patch(ax, g, **kwargs)
        return
    xy = np.array(shp.exterior.coords)
    ax.add_patch(MplPolygon(graphics_scale(xy), closed=True, **kwargs))

def line_collection_from_sample(sample_nm, values_norm, cmap="plasma", lw=4):
    pts = np.asarray(sample_nm)
    segs = np.stack([pts[:-1], pts[1:]], axis=1)
    vals = 0.5 * (values_norm[:-1] + values_norm[1:])
    lc = LineCollection(segs, cmap=cmap, linewidths=lw, capstyle="round")
    lc.set_array(vals)
    lc.set_clim(0.0, 1.0)
    return lc

# ---------- geometry builders ----------
# Each builder takes `size` = total graphene–metal interface length [m].
# All geometries with the same `size` have exactly the same total interface.
#
#   Baseline    : one straight side of length d = size
#   Wrap-around : two sides + front;  d_side = (size − W) / 2,  total = 2·d_side + W
#   Square taper: fanned pad, one side of length d = size  (pad_width independent)
#   Curved pad  : semicircle, arc = π·r = size  →  r = size / π

def geom_baseline(size=None):
    """Rectangular strip contacted along one long edge. interface = size."""
    if size is None:
        size = contact_length
    W = W_interconnect
    d = size                            # contact length = total interface
    poly = Polygon([(0, -W/2), (d, -W/2), (d, W/2), (0, W/2)])
    inlet = LineString([(0, -W/2), (0, W/2)])
    metal = Polygon([(0, W/2), (d, W/2), (d, W/2 + metal_gap_viz), (0, W/2 + metal_gap_viz)])

    def edge_selector(pa, pb, mid):
        return (abs(mid[1] - W/2) < 0.6*mesh_h) and (0 - 1e-15 <= mid[0] <= d + 1e-15)

    def metal_drain_selector(pa, pb, mid):
        return abs(mid[1] - (W/2 + metal_gap_viz)) < 0.6*mesh_h

    n = max(2, int(np.ceil(d / mesh_h)))
    sample = np.column_stack([np.linspace(0, d, n), np.full(n, W/2 - eps_sample)])
    return dict(
        name="Baseline", poly=poly, inlet=inlet, metal=metal,
        edge_selector=edge_selector, metal_drain_selector=metal_drain_selector,
        sample_line=sample, sample_label="distance along side contact x (nm)"
    )

def geom_wrap(size=None):
    """Three-sided wrap contact. interface = 2·d_side + W = size.
    Requires size > W_interconnect so each side d_side = (size−W)/2 > 0."""
    if size is None:
        size = contact_length
    W = W_interconnect
    d = (size - W) / 2                 # each side length; total = 2d + W = size
    if d < mesh_h:
        raise ValueError(f"geom_wrap: size={size/nm:.0f} nm is too small; "
                         f"need size > W_interconnect={W/nm:.0f} nm")
    poly = Polygon([(0, -W/2), (d, -W/2), (d, W/2), (0, W/2)])
    inlet = LineString([(0, -W/2), (0, W/2)])

    metal_top   = Polygon([(0, W/2), (d, W/2), (d, W/2 + metal_gap_viz), (0, W/2 + metal_gap_viz)])
    metal_bot   = Polygon([(0, -W/2 - metal_gap_viz), (d, -W/2 - metal_gap_viz), (d, -W/2), (0, -W/2)])
    metal_front = Polygon([(d, -W/2 - metal_gap_viz), (d + metal_gap_viz, -W/2 - metal_gap_viz),
                           (d + metal_gap_viz, W/2 + metal_gap_viz), (d, W/2 + metal_gap_viz)])
    metal = metal_top.union(metal_bot).union(metal_front)

    def edge_selector(pa, pb, mid):
        top   = (abs(mid[1] - W/2)  < 0.6*mesh_h) and (0 <= mid[0] <= d)
        bot   = (abs(mid[1] + W/2)  < 0.6*mesh_h) and (0 <= mid[0] <= d)
        front = (abs(mid[0] - d)    < 0.6*mesh_h) and (-W/2 <= mid[1] <= W/2)
        return top or bot or front

    def metal_drain_selector(pa, pb, mid):
        top   = abs(mid[1] - (W/2 + metal_gap_viz))  < 0.6*mesh_h
        bot   = abs(mid[1] - (-W/2 - metal_gap_viz)) < 0.6*mesh_h
        right = abs(mid[0] - (d + metal_gap_viz))     < 0.6*mesh_h
        return top or bot or right

    n1 = max(2, int(np.ceil(d / mesh_h)))
    n2 = max(2, int(np.ceil(W / mesh_h)))
    seg1 = np.column_stack([np.linspace(0, d, n1), np.full(n1, -W/2 + eps_sample)])
    seg2 = np.column_stack([np.full(n2, d - eps_sample), np.linspace(-W/2 + eps_sample, W/2 - eps_sample, n2)])
    seg3 = np.column_stack([np.linspace(d, 0, n1), np.full(n1, W/2 - eps_sample)])
    sample = np.vstack([seg1, seg2[1:], seg3[1:]])
    return dict(
        name="Wrap-around", poly=poly, inlet=inlet, metal=metal,
        edge_selector=edge_selector, metal_drain_selector=metal_drain_selector,
        sample_line=sample, sample_label="position along contacted perimeter s (nm)",
        markers=[d, d + W, size], labels=["bottom side", "front", "top side"]
    )

def geom_square(size=None):
    """Straight-sided taper to pad; one long edge contacted. interface = size.
    pad_width (W2) is a fixed design parameter independent of size."""
    if size is None:
        size = contact_length
    W1 = W_interconnect
    W2 = pad_width                      # pad width (independent of interface length)
    Lt = taper_length
    d  = size                           # contact length along top of pad = total interface
    poly = Polygon([
        (0, -W1/2),
        (Lt, -W2/2),
        (Lt + d, -W2/2),
        (Lt + d, W2/2),
        (Lt, W2/2),
        (0, W1/2),
    ])
    inlet = LineString([(0, -W1/2), (0, W1/2)])
    metal = Polygon([(Lt, W2/2), (Lt + d, W2/2), (Lt + d, W2/2 + metal_gap_viz), (Lt, W2/2 + metal_gap_viz)])

    def edge_selector(pa, pb, mid):
        return (abs(mid[1] - W2/2) < 0.6*mesh_h) and (Lt - 1e-15 <= mid[0] <= Lt + d + 1e-15)

    def metal_drain_selector(pa, pb, mid):
        return abs(mid[1] - (W2/2 + metal_gap_viz)) < 0.6*mesh_h

    n = max(2, int(np.ceil(d / mesh_h)))
    sample = np.column_stack([np.linspace(Lt, Lt + d, n), np.full(n, W2/2 - eps_sample)])
    return dict(
        name="Square taper", poly=poly, inlet=inlet, metal=metal,
        edge_selector=edge_selector, metal_drain_selector=metal_drain_selector,
        sample_line=sample, sample_label="distance along side contact x (nm)"
    )

def geom_curved(size=None):
    """Semicircular pad; arc length π·r = size → r = size/π."""
    if size is None:
        size = contact_length
    r = size / np.pi                    # radius such that arc = π·r = size
    W1 = W_interconnect

    theta = np.linspace(-np.pi/2, np.pi/2, 120)
    poly_pts = np.vstack([
        [0, -r],
        np.column_stack([r*np.cos(theta), r*np.sin(theta)]),
        [0,  r],
    ])
    poly  = Polygon(poly_pts)
    inlet = LineString([(0, -W1/2), (0, W1/2)])

    outer_r = r + metal_gap_viz
    theta2 = np.linspace(-np.pi/2, np.pi/2, 140)
    xo = outer_r*np.cos(theta2);  yo = outer_r*np.sin(theta2)
    xi = r*np.cos(theta2[::-1]);   yi = r*np.sin(theta2[::-1])
    metal = Polygon(np.column_stack([np.r_[xo, xi], np.r_[yo, yi]]))

    def edge_selector(pa, pb, mid):
        rr = np.hypot(mid[0], mid[1])
        return (abs(rr - r) < 0.9*mesh_h) and (mid[0] >= -1e-15)

    def metal_drain_selector(pa, pb, mid):
        rr = np.hypot(mid[0], mid[1])
        return (abs(rr - outer_r) < 0.9*mesh_h) and (mid[0] >= -1e-15)

    n = max(2, int(np.ceil(np.pi * r / mesh_h)))
    theta_s = np.linspace(-np.pi/2, np.pi/2, n)
    sample = np.column_stack([
        (r - eps_sample)*np.cos(theta_s),
        (r - eps_sample)*np.sin(theta_s),
    ])
    return dict(
        name="Curved pad", poly=poly, inlet=inlet, metal=metal,
        edge_selector=edge_selector, metal_drain_selector=metal_drain_selector,
        sample_line=sample, sample_label="position along curved contact s (nm)"
    )

def geom_holes(size=None, n_holes=4, hole_radius=None, hole_radial_frac=0.55):
    """Semicircular pad with holes punched in the metal annulus.
    interface = π·r = size  (same rule as geom_curved)."""
    if size is None:
        size = contact_length
    W1 = W_interconnect
    r = size / np.pi                    # arc length = π·r = size (same as geom_curved)

    # ---------------- graphene ----------------
    theta = np.linspace(-np.pi/2, np.pi/2, 120)
    x_arc = r * np.cos(theta)
    y_arc = r * np.sin(theta)

    poly_pts = np.vstack([
        [0, -r],
        np.column_stack([x_arc, y_arc]),
        [0,  r],
    ])
    poly = Polygon(poly_pts)
    inlet = LineString([(0, -W1/2), (0, W1/2)])

    # ---------------- metal annulus ----------------
    outer_r = r + metal_gap_viz

    theta2 = np.linspace(-np.pi/2, np.pi/2, 160)
    xo = outer_r * np.cos(theta2)
    yo = outer_r * np.sin(theta2)
    xi = r * np.cos(theta2[::-1])
    yi = r * np.sin(theta2[::-1])

    metal_base = Polygon(np.column_stack([np.r_[xo, xi], np.r_[yo, yi]]))

    # ---------------- holes in the metal ----------------
    annulus_t = outer_r - r

    if hole_radius is None:
        hole_radius = 0.22 * annulus_t

    # Keep holes fully inside the annulus
    hole_radius = min(hole_radius, 0.45 * annulus_t)

    # Place hole centers on a mid-annulus arc
    r_hole = r + hole_radial_frac * annulus_t

    # Angular margin so end holes do not break out of the annulus ends
    ang_margin = max(
        np.deg2rad(15),
        np.arcsin(min(0.99, (hole_radius + 2 * mesh_h) / max(r_hole, 1e-15)))
    )

    if n_holes > 0:
        theta_holes = np.linspace(-np.pi/2 + ang_margin, np.pi/2 - ang_margin, n_holes)

        hole_geoms = [
            Point(r_hole * np.cos(th), r_hole * np.sin(th)).buffer(
                hole_radius, resolution=48
            )
            for th in theta_holes
        ]

        holes_union = unary_union(hole_geoms)
        metal = metal_base.difference(holes_union).buffer(0)
    else:
        metal = metal_base

    # ---------------- same contacted graphene edge as before ----------------
    def edge_selector(pa, pb, mid):
        rr = np.hypot(mid[0], mid[1])
        return (abs(rr - r) < 0.9 * mesh_h) and (mid[0] >= -1e-15)

    def metal_drain_selector(pa, pb, mid):
        rr = np.hypot(mid[0], mid[1])
        return (abs(rr - outer_r) < 0.9 * mesh_h) and (mid[0] >= -1e-15)   # outer arc only

    # Sample just inside the graphene arc
    n = max(2, int(np.ceil(np.pi * r / mesh_h)))
    theta_s = np.linspace(-np.pi/2, np.pi/2, n)
    sample = np.column_stack([
        (r - eps_sample) * np.cos(theta_s),
        (r - eps_sample) * np.sin(theta_s)
    ])

    return dict(
        name="Curved pad with metal holes",
        poly=poly, inlet=inlet, metal=metal,
        edge_selector=edge_selector, metal_drain_selector=metal_drain_selector,
        sample_line=sample, sample_label="position along curved contact s (nm)"
    )

def build_all_geometries(size=None):
    """Build all standard geometries with the same total interface length `size`.
    Defaults to the module-level contact_length."""
    if size is None:
        size = contact_length
    return [geom_baseline(size), geom_wrap(size), geom_square(size), geom_curved(size)]

# ---------- two-domain interface helpers ----------
def find_matching_metal_edges(pts_g, pts_m, graphene_iface_edges, metal_bedges, tol=None):
    """Match graphene interface edges to metal boundary edges by midpoint proximity.

    For axis-aligned interfaces (rectangular geometries) the match is exact.
    For curved interfaces both meshes are staircased so midpoints may be up to
    ~mesh_h apart; tol=1.5*mesh_h catches these while rejecting unrelated edges.

    Returns (matched_pairs, metal_iface_node_set) where matched_pairs is a list
    of ((a_g,b_g),(a_m,b_m)) tuples and each metal edge appears at most once.
    """
    if tol is None:
        tol = 1.5 * mesh_h
    if not graphene_iface_edges or not metal_bedges:
        return [], set()
    metal_mids = np.array([0.5*(pts_m[a] + pts_m[b]) for a, b in metal_bedges])
    tree = cKDTree(metal_mids)
    used = set()
    matched_pairs = []
    metal_iface_nodes = set()
    k_search = min(8, len(metal_bedges))
    for a_g, b_g in graphene_iface_edges:
        mid_g = 0.5 * (pts_g[a_g] + pts_g[b_g])
        dists, idxs = tree.query(mid_g, k=k_search)
        if k_search == 1:
            dists, idxs = [dists], [idxs]
        for dist, idx in zip(dists, idxs):
            if dist < tol and idx not in used:
                a_m, b_m = metal_bedges[idx]
                used.add(idx)
                matched_pairs.append(((int(a_g), int(b_g)), (int(a_m), int(b_m))))
                metal_iface_nodes.add(int(a_m))
                metal_iface_nodes.add(int(b_m))
                break
    return matched_pairs, metal_iface_nodes


def add_interface_coupling(A, pts_g, pts_m, n_g, matched_pairs, g_edge):
    """Add two-domain interface coupling to the global matrix A (vectorized COO assembly).

    For each matched edge pair the weak form adds ±Me blocks where
    Me = g_edge * L_edge/6 * [[2,1],[1,2]].
    """
    if not matched_pairs:
        return A.tocsr(), [], np.array([])
    pairs_g = np.array([(a, b) for (a, b), _ in matched_pairs], dtype=int)
    pairs_m = np.array([(a, b) for _, (a, b) in matched_pairs], dtype=int)
    L = np.linalg.norm(pts_g[pairs_g[:, 1]] - pts_g[pairs_g[:, 0]], axis=1)
    c = g_edge * L / 6.0                            # (N,) coupling coefficient per edge
    ag, bg = pairs_g[:, 0], pairs_g[:, 1]
    am, bm = pairs_m[:, 0] + n_g, pairs_m[:, 1] + n_g
    # 16 COO entries per edge: 4 graphene, 4 metal, 4+4 cross
    rows = np.concatenate([ag,ag,bg,bg,  am,am,bm,bm,  ag,ag,bg,bg,  am,am,bm,bm])
    cols = np.concatenate([ag,bg,ag,bg,  am,bm,am,bm,  am,bm,am,bm,  ag,bg,ag,bg])
    data = np.concatenate([2*c,c,c,2*c,  2*c,c,c,2*c,  -2*c,-c,-c,-2*c,  -2*c,-c,-c,-2*c])
    A = A.tocsr() + coo_matrix((data, (rows, cols)), shape=A.shape, dtype=complex).tocsr()
    contact_edges_g = list(zip(ag.tolist(), bg.tolist()))
    return A, contact_edges_g, L


def total_contact_current_full(pts_g, V_g, V_m, matched_pairs, g_edge):
    """Total current flowing from graphene to metal across the coupled interface."""
    if not matched_pairs:
        return 0j
    pairs_g = np.array([(a, b) for (a, b), _ in matched_pairs], dtype=int)
    pairs_m = np.array([(a, b) for _, (a, b) in matched_pairs], dtype=int)
    L = np.linalg.norm(pts_g[pairs_g[:, 1]] - pts_g[pairs_g[:, 0]], axis=1)
    dV = 0.5 * ((V_g[pairs_g[:,0]] - V_m[pairs_m[:,0]]) + (V_g[pairs_g[:,1]] - V_m[pairs_m[:,1]]))
    return g_edge * np.dot(L, dV)


# ---------- solve / plot ----------

def prepare_geom(geom):
    """Build mesh + tag graphene elements by distance to Gamma_c."""
    pts_g, tris_g = build_structured_tri_mesh(geom["poly"], mesh_h)
    bedges_g = boundary_edges(tris_g)
    n_g = len(pts_g)

    pts_m, tris_m = build_structured_tri_mesh(geom["metal"], mesh_h)
    bedges_m = boundary_edges(tris_m)
    n_m = len(pts_m)

    # Identify graphene interface edges
    graphene_iface_edges = [
        (a, b) for a, b in bedges_g
        if geom["edge_selector"](pts_g[a], pts_g[b], 0.5 * (pts_g[a] + pts_g[b]))
    ]
    matched_pairs, _ = find_matching_metal_edges(pts_g, pts_m, graphene_iface_edges, bedges_m)

    # --- NEW: Build distance field from Gamma_c for graphene elements ---
    # Collect all interface node positions
    iface_node_set = set()
    for (a, b), _ in matched_pairs:
        iface_node_set.add(a)
        iface_node_set.add(b)
    if iface_node_set:
        iface_pts = pts_g[np.array(sorted(iface_node_set))]
        iface_tree = cKDTree(iface_pts)
        # Distance from each element centroid to nearest interface node
        centroids_g = pts_g[tris_g].mean(axis=1)
        dist_to_iface, _ = iface_tree.query(centroids_g)
        is_overlap = dist_to_iface <= L_overlap
    else:
        is_overlap = np.zeros(len(tris_g), dtype=bool)

    inlet_nodes = inlet_nodes_for_line(pts_g, geom["inlet"], mesh_h)

    metal_drain_local = set()
    for a, b in bedges_m:
        pa, pb = pts_m[a], pts_m[b]
        mid = 0.5 * (pa + pb)
        if geom["metal_drain_selector"](pa, pb, mid):
            metal_drain_local.add(int(a))
            metal_drain_local.add(int(b))
    metal_drain_global = [n_g + i for i in sorted(metal_drain_local)]

    sample_line = geom.get("sample_line")
    tri_obj = Triangulation(pts_g[:, 0], pts_g[:, 1], tris_g) if sample_line is not None else None

    return dict(
        pts_g=pts_g, tris_g=tris_g, n_g=n_g,
        pts_m=pts_m, tris_m=tris_m, n_m=n_m,
        matched_pairs=matched_pairs,
        inlet_nodes=inlet_nodes,
        metal_drain_global=metal_drain_global,
        sample_line=sample_line,
        tri_obj=tri_obj,
        name=geom["name"],
        is_overlap=is_overlap,  # NEW: boolean array per graphene element
    )


def solve_prepared(prepared, sigma_channel, sigma_overlap, sigma_m, g_edge):
    """Solve with spatially varying graphene conductivity.

    sigma_channel : sheet conductance for undoped graphene [S/□]
    sigma_overlap : sheet conductance for Ni-doped overlap zone [S/□]
    sigma_m       : sheet conductance for metal [S/□] (uniform)
    g_edge        : interface conductance per unit length [S/m]
    """
    pts_g = prepared["pts_g"];
    tris_g = prepared["tris_g"];
    n_g = prepared["n_g"]
    pts_m = prepared["pts_m"];
    tris_m = prepared["tris_m"]
    n_m = len(pts_m)
    matched_pairs = prepared["matched_pairs"]
    inlet_nodes = prepared["inlet_nodes"]
    metal_drain_global = prepared["metal_drain_global"]
    is_overlap = prepared["is_overlap"]

    # Build per-element conductivity arrays
    n_tri_g = len(tris_g)
    sigma_g_elem = np.full(n_tri_g, sigma_channel, dtype=complex)
    sigma_g_elem[is_overlap] = sigma_overlap

    n_tri_m = len(tris_m)
    sigma_m_elem = np.full(n_tri_m, sigma_m, dtype=complex)

    A_g, b_g = assemble_bulk_matrix(pts_g, tris_g, sigma_g_elem)
    A_m, b_m = assemble_bulk_matrix(pts_m, tris_m, sigma_m_elem)

    A_global = sp_block_diag([A_g, A_m], format="csr", dtype=complex)
    b_global = np.concatenate([b_g, b_m])

    A_global, contact_edges, contact_lengths = add_interface_coupling(
        A_global, pts_g, pts_m, n_g, matched_pairs, g_edge)

    A_pre_bc = A_global.tocsr()
    A_bc, b_bc = apply_dirichlet_nodes(A_global, b_global, inlet_nodes, 1.0)
    A_bc, b_bc = apply_dirichlet_nodes(A_bc, b_bc, metal_drain_global, 0.0)

    V = spsolve(A_bc.tocsr(), b_bc)
    V_g = V[:n_g]
    V_m = V[n_g:]

    I_contact = total_contact_current_full(pts_g, V_g, V_m, matched_pairs, g_edge)
    I_inlet = reaction_current(A_pre_bc, V, inlet_nodes)

    result = dict(
        pts=pts_g, tris=tris_g, V=V_g,
        pts_m=pts_m, tris_m=tris_m, V_m=V_m,
        contact_edges=contact_edges, contact_lengths=contact_lengths,
        I_contact=I_contact, I_inlet=I_inlet,
    )
    if prepared["sample_line"] is not None:
        V_m_at_g = np.zeros(n_g, dtype=V_g.dtype)
        count_g = np.zeros(n_g, dtype=float)
        for (a_g, b_g), (a_m, b_m) in matched_pairs:
            V_m_at_g[a_g] += V_m[a_m];
            count_g[a_g] += 1
            V_m_at_g[b_g] += V_m[b_m];
            count_g[b_g] += 1
        nz = count_g > 0
        V_m_at_g[nz] /= count_g[nz]
        Jnodal = g_edge * (V_g - V_m_at_g)
        s, Jline = sample_line_profile(prepared["tri_obj"], np.abs(Jnodal), prepared["sample_line"])
        result.update(Jnodal=Jnodal, s=s, Jline=Jline)
    return result


def solve_geom(geom, sigma_channel, sigma_overlap, sigma_m, g_edge):
    return solve_prepared(prepare_geom(geom), sigma_channel, sigma_overlap, sigma_m, g_edge)


def make_figure(geoms, sols, out_png):
    graphene_color = "#B9D6F2"
    metal_color = "#D9A441"
    inlet_color = "#2CA25F"

    global_max = max(float(np.max(sol["Jline"])) for sol in sols)

    fig, axes = plt.subplots(2, 4, figsize=(16.6, 8.9))
    plt.subplots_adjust(top=0.84, bottom=0.09, left=0.05, right=0.92, hspace=0.46, wspace=0.25)
    fig.suptitle("Lateral graphene-metal side-contact FEM", fontsize=17, y=0.975)
    fig.text(0.5, 0.935,
             "blue = graphene, gold = adjacent metal, color = shared normalization $J_c/J_{c,\\max}^{(all)}$",
             ha="center", va="center", fontsize=13.3)
    handles = [
        Patch(facecolor=graphene_color, edgecolor="k", label="graphene"),
        Patch(facecolor=metal_color, edgecolor="none", hatch="////", label="metal"),
        plt.Line2D([0],[0], color=inlet_color, lw=4, label="1 V inlet"),
    ]
    fig.legend(handles=handles, loc="upper center", bbox_to_anchor=(0.5, 0.905), ncol=3, frameon=False, fontsize=11.4)

    last_lc = None
    for col, (geom, sol) in enumerate(zip(geoms, sols)):
        ax = axes[0, col]
        add_poly_patch(ax, geom["poly"], facecolor=graphene_color, edgecolor="k", lw=1.3, zorder=1)
        add_poly_patch(ax, geom["metal"], facecolor=metal_color, edgecolor=metal_color, hatch="////", lw=1.0, alpha=0.65, zorder=0)

        inlet_xy = graphics_scale(np.array(geom["inlet"].coords))
        ax.plot(inlet_xy[:,0], inlet_xy[:,1], color=inlet_color, lw=4, solid_capstyle="round", zorder=4)

        # color the contacted boundary path using the sampled Jline
        Jn = sol["Jline"] / global_max
        lc = line_collection_from_sample(graphics_scale(geom["sample_line"]), Jn, lw=4.2)
        ax.add_collection(lc)
        last_lc = lc

        Rc = 1.0 / sol["I_contact"]          # complex contact impedance [Ω]
        ax.set_title(f"{col+1}. {geom['name']}\n$|Z_C|$ = {fmt_ohm(abs(Rc))}, $\\angle(Z_C)$ = {np.angle(Rc, deg=True):.1f}°", fontsize=11.2)
        ax.set_aspect("equal")
        minx, miny, maxx, maxy = geom["poly"].bounds
        mminx, mminy, mmaxx, mmaxy = geom["metal"].bounds
        minx, miny = min(minx, mminx), min(miny, mminy)
        maxx, maxy = max(maxx, mmaxx), max(maxy, mmaxy)
        padx = 0.08*(maxx-minx) + metal_gap_viz
        pady = 0.15*(maxy-miny) + metal_gap_viz
        ax.set_xlim((minx - padx)*1e9, (maxx + padx)*1e9)
        ax.set_ylim((miny - pady)*1e9, (maxy + pady)*1e9)
        ax.set_xticks([]); ax.set_yticks([])
        for sp in ax.spines.values():
            sp.set_visible(False)

        ax2 = axes[1, col]
        s_nm = sol["s"] * 1e9
        ax2.plot(s_nm, Jn, lw=2.3)
        ax2.grid(True, alpha=0.3)
        ax2.set_ylim(0.0, 1.05)
        ax2.set_xlim(0.0, s_nm[-1])
        ax2.set_title("Interface current density", fontsize=11.2)
        ax2.set_xlabel(geom["sample_label"], fontsize=10.2)
        if col == 0:
            ax2.set_ylabel("$J_c/J_{c,\\max}^{(all)}$", fontsize=11)

        lam = characteristic_length_1e(s_nm, sol["Jline"])
        if lam is not None:
            y1e = (sol["Jline"][0] / global_max) / np.e
            ax2.axhline(y1e, color="0.5", ls="--", lw=1)
            ax2.axvspan(0, lam, color="0.6", alpha=0.10)
            ax2.text(0.52*lam, min(0.93, Jn[0]*0.95), f"$\\lambda_{{1/e}}\\approx {lam:.0f}$ nm",
                     ha="center", va="top", fontsize=9, color="0.25")

        if "markers" in geom:
            for m in geom["markers"][:-1]:
                ax2.axvline(m * 1e9, color="0.6", ls=":", lw=1)
            ymax = 1.02
            edges = geom["markers"]
            mids = [0.5*(0 + edges[0]), 0.5*(edges[0] + edges[1]), 0.5*(edges[1] + edges[2])]
            for mid, label in zip(mids, geom["labels"]):
                ax2.text(mid*1e9, ymax*0.95, label, ha="center", va="top", fontsize=8.6)

    cax = fig.add_axes([0.935, 0.47, 0.017, 0.31])
    cb = fig.colorbar(last_lc, cax=cax)
    cb.set_label("$J_c/J_{c,\\max}^{(all)}$", fontsize=11.5)

    fig.savefig(out_png, dpi=220, bbox_inches="tight")
    plt.show()

def write_verification(geoms, sols, path_txt):
    with open(path_txt, "w") as f:
        f.write("="*96 + "\n")
        f.write("LATERAL SIDE-CONTACT FEM CHECK\n")
        f.write("="*96 + "\n")
        f.write(f"freq             = {freq:.3e} Hz\n")
#        f.write(f"sigma_s          = {sigma_s:.4e} S  (|sigma_s|={abs(sigma_s):.4e})\n")
        f.write(f"R_sheet_m        = {R_sheet_m:.3g} ohm/sq  (mu_r={mu_r_m}, eps_r={eps_r_m}, d={d_m*1e9:.1f} nm)\n")
 #       f.write(f"g_edge           = {g_edge:.3e} S/m  (complex: R_bond + jωL_bond || jωC_bond)\n")
        f.write(f"contact_length   = {contact_length/nm:.1f} nm\n")
        f.write(f"mesh pitch       = {mesh_h/nm:.1f} nm\n\n")
        f.write(f"{'Geometry':<16} {'|ZC| (ohm)':>12} {'|ZC|W (ohm*um)':>16} {'Ierr rel':>12} {'lambda_1e (nm)':>16}\n")
        f.write("-"*96 + "\n")
        for geom, sol in zip(geoms, sols):
            Zc = 1.0 / sol["I_contact"]          # complex contact impedance [Ω]
            Zc_abs = abs(Zc)
            ZcW = Zc_abs * W_interconnect * 1e4
            rel = abs(sol["I_inlet"] - sol["I_contact"]) / abs(sol["I_contact"])
            lam = characteristic_length_1e(sol["s"] * 1e9, sol["Jline"])
            lam_str = f"{lam:.1f}" if lam is not None else "n/a"
            f.write(f"{geom['name']:<16} {Zc_abs:>12.2f} {ZcW:>16.2f} {rel:>12.3e} {lam_str:>16}\n")

if __name__ == "__main__":
    _mats = compute_materials(freq)     # returns 4-tuple now
    geoms = build_all_geometries()
    sols = [solve_geom(g, *_mats) for g in geoms]
    out_png = "contact_lateral_sidecontact_fem.png"
    out_txt = "contact_lateral_sidecontact_fem_verification.txt"
    make_figure(geoms, sols, out_png)
    write_verification(geoms, sols, out_txt)
    print(f"Saved: {out_png}")
    print(f"Saved: {out_txt}")