"""Coupled multi-layer finite-difference solver.

Each included copper layer is a 2D 5-point sheet with per-layer face
conductance sigma_s = t/rho [S] (square cells: independent of h); via and
plated-through-pad barrels add vertical conductances between vertically
aligned cells of the layers they span AND reach copper on. A barrel
passing an antipad still bridges the layers above/below it with the full
barrel length. At freq > 0 the per-layer sheet conductances and the
barrel walls get the 1D skin-effect correction (see skin.py; AC results
are a rigorous lower bound - lateral redistribution is not modeled).

Two contact models for the terminals (each terminal = merged parts):

- "uniform" (default): a conductor pressed onto the contact area injects
  the current orthogonally with UNIFORM surface density: every contact
  cell sources (sinks) I/N. The in-plane current density ramps across
  the contact instead of being zero. The pure-Neumann system is grounded
  at one V- cell (that cell's sink share is exactly the flux that exits
  through the ground reference, so the solution equals the singular
  system's). R = (<V over V+ cells> - <V over V- cells>) / I; because
  the injection and averaging weights coincide, sum(edge powers) = I^2 R
  holds exactly and remains the consistency check.

- "equipotential": ideal bonded lug; contact cells are Dirichlet
  (V+ = 1 V, V- = 0). R from the exact discrete electrode flux. Touching
  terminals are rejected (a direct face would short the Dirichlet
  regions); with "uniform" contacts touching is physically fine.

The two models bracket a real contact: R_equipotential <= R_real <=
R_uniform. Missing neighbors give no matrix term = insulated boundary.
Current density per layer comes from face currents (np.gradient across
the NaN staircase boundary would pollute the field). Power density per
layer distributes each in-plane edge's dissipation half to each endpoint
cell. All reported fields are rescaled to the test current I_test.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field

import numpy as np
from scipy import sparse
from scipy.sparse import csgraph
from scipy.sparse import linalg as sla

from . import config, skin
from .errors import ConnectivityError, ElectrodeError
from .geometry import Problem
from .raster import RasterStack, electrodes_touch


@dataclass
class SolveInfo:
    method: str                   # "spsolve" | "cg+jacobi"
    n_unknowns: int
    iterations: int | None = None
    residual: float | None = None


@dataclass
class Edges:
    a: np.ndarray                 # int64 flat cell ids
    b: np.ndarray
    w: np.ndarray                 # conductance [S]
    via_index: np.ndarray         # int32; -1 = in-plane edge


@dataclass
class ViaReport:
    x_mm: float
    y_mm: float
    kind: str
    drill_mm: float
    current_a: float              # max barrel-segment current @ I_test
    power_w: float                # total barrel dissipation @ I_test


@dataclass
class Result:
    R_ohm: float
    i_test: float
    V: np.ndarray                 # (L, ny, nx) volts @ I_test, NaN off-copper
    Jmag: np.ndarray              # (L, ny, nx) A/m^2 @ I_test
    Parea: np.ndarray             # (L, ny, nx) W/m^2 @ I_test
    layer_names: list[str]
    P_total: float                # I_test^2 * R
    P_layers: list[float]         # in-plane dissipation per layer @ I_test
    P_vias: float                 # total barrel dissipation @ I_test
    power_balance_rel: float      # |sum(edge powers) - I^2 R| / I^2 R
    via_reports: list[ViaReport]  # sorted by current, descending
    I1_a: float                   # electrode currents (unit drive)
    I2_a: float
    mismatch_rel: float
    n_free: int
    solve_info: SolveInfo
    # per-part terminal currents @ I_test: [(label, amps), ...];
    # computed flux for "equipotential", prescribed area share for "uniform"
    part_currents1: list = field(default_factory=list)
    part_currents2: list = field(default_factory=list)
    contact_model: str = "uniform"
    freq_hz: float = 0.0
    skin_depth_um: float | None = None
    rs_ratios: list[float] = field(default_factory=list)  # R_AC/R_DC per layer
    timings: dict = field(default_factory=dict)


def _shifts2d():
    return [
        ((slice(None), slice(None, -1)), (slice(None), slice(1, None))),
        ((slice(None, -1), slice(None)), (slice(1, None), slice(None))),
    ]


def _sigma_2d(stack: RasterStack, li: int, sigma_layer: float,
              sigma_buildup: float) -> np.ndarray | None:
    """Per-cell sheet conductance for one layer, or None if uniform."""
    if stack.buildup is None or sigma_buildup <= 0 \
            or not stack.buildup[li].any():
        return None
    s = np.full(stack.shape2d, sigma_layer)
    s[stack.buildup[li]] += sigma_buildup
    return s


def build_edges(stack: RasterStack, problem: Problem, sigmas: list[float],
                via_factor: float = 1.0,
                sigma_buildup: float = 0.0) -> Edges:
    """All copper-copper conductances: in-plane faces + via barrels.
    sigmas: effective (possibly AC) sheet conductance per layer;
    via_factor: R_AC/R_DC of the barrel wall; sigma_buildup: extra sheet
    conductance on solder-buildup cells. Faces between cells of unequal
    conductance use the harmonic mean (series half-cells), which reduces
    exactly to sigma for uniform regions."""
    L, ny, nx = stack.masks.shape
    plane = ny * nx
    aa, bb, ww, vv = [], [], [], []

    for li in range(L):
        m = stack.masks[li]
        sig = sigmas[li]
        scell = _sigma_2d(stack, li, sig, sigma_buildup)
        base = li * plane
        for src, dst in _shifts2d():
            pair = m[src] & m[dst]
            ii, jj = np.nonzero(pair)
            if src[0] == slice(None):            # horizontal: j, j+1
                a = base + ii * nx + jj
                b = a + 1
            else:                                 # vertical: i, i+1
                a = base + ii * nx + jj
                b = a + nx
            aa.append(a.astype(np.int64))
            bb.append(b.astype(np.int64))
            if scell is None:
                ww.append(np.full(len(a), sig))
            else:
                s_a = scell[src][pair]
                s_b = scell[dst][pair]
                ww.append(2.0 * s_a * s_b / (s_a + s_b))
            vv.append(np.full(len(a), -1, dtype=np.int32))

    for vi, via in enumerate(problem.vias):
        cell = stack.cell_of(via.x, via.y)
        if cell is None:
            continue
        i, j = cell
        present = [li for li, layer in enumerate(problem.layers)
                   if via.spans(layer.z_nm) and stack.masks[li, i, j]]
        for la, lb in zip(present[:-1], present[1:]):
            length = problem.layers[lb].z_nm - problem.layers[la].z_nm
            if length <= 0:
                continue
            r = via.barrel_resistance(length, problem.rho_ohm_m,
                                      problem.plating_nm) * via_factor
            aa.append(np.array([la * plane + i * nx + j], dtype=np.int64))
            bb.append(np.array([lb * plane + i * nx + j], dtype=np.int64))
            ww.append(np.array([1.0 / r]))
            vv.append(np.array([vi], dtype=np.int32))

    if not aa:
        raise ConnectivityError("No copper found on the selected layers.")
    return Edges(a=np.concatenate(aa), b=np.concatenate(bb),
                 w=np.concatenate(ww), via_index=np.concatenate(vv))


def connected_restrict(stack: RasterStack, e1: np.ndarray, e2: np.ndarray,
                       edges: Edges) -> bool:
    """Keep only components (through-plane AND through-via) touching both
    terminals. Mutates stack.masks / e1 / e2. Returns True if anything
    was dropped (caller must rebuild edges)."""
    n = stack.masks.size
    graph = sparse.coo_matrix(
        (np.ones(len(edges.a)), (edges.a, edges.b)), shape=(n, n))
    _, labels = csgraph.connected_components(graph, directed=False)
    labels3 = labels.reshape(stack.masks.shape)
    common = np.intersect1d(np.unique(labels3[e1]), np.unique(labels3[e2]))
    if len(common) == 0:
        raise ConnectivityError(
            "The two terminals are not connected by the selected fill "
            "layers (not even through vias). Check the layer selection and "
            "that the fills are up to date."
        )
    keep = np.isin(labels3, common) & stack.masks
    changed = bool((stack.masks & ~keep).any())
    stack.masks &= keep
    e1 &= keep
    e2 &= keep
    return changed


def _assemble(state: np.ndarray, edges: Edges, rhs_extra: np.ndarray | None):
    """Weighted-Laplacian assembly with Dirichlet elimination.
    state: 0 off, 1 free, 2 Dirichlet@1V, 3 Dirichlet@0V.
    rhs_extra: per-flat-cell current injection [A] added for free cells."""
    n = state.size
    sa, sb = state[edges.a], state[edges.b]
    short = ((sa == 2) & (sb == 3)) | ((sa == 3) & (sb == 2))
    if short.any():
        n_via = int((edges.via_index[short] >= 0).sum())
        raise ElectrodeError(
            f"The terminals are directly connected by {int(short.sum())} "
            f"conductance(s) ({n_via} via barrel(s)) without any free copper "
            f"in between - move the contacts apart."
        )

    free = state == 1
    n_free = int(free.sum())
    if n_free == 0:
        raise ElectrodeError(
            "No free copper cells remain between the terminals - the "
            "contacts cover the whole fill at this grid resolution."
        )
    idx = np.full(n, -1, dtype=np.int64)
    idx[free] = np.arange(n_free)

    diag = np.zeros(n_free)
    rhs = np.zeros(n_free)
    fa, fb = sa == 1, sb == 1
    np.add.at(diag, idx[edges.a[fa]], edges.w[fa])
    np.add.at(diag, idx[edges.b[fb]], edges.w[fb])
    r1a = fa & (sb == 2)
    r1b = fb & (sa == 2)
    np.add.at(rhs, idx[edges.a[r1a]], edges.w[r1a])
    np.add.at(rhs, idx[edges.b[r1b]], edges.w[r1b])
    if rhs_extra is not None:
        rhs += rhs_extra[free]

    ff = fa & fb
    rows = np.concatenate([idx[edges.a[ff]], idx[edges.b[ff]],
                           np.arange(n_free)])
    cols = np.concatenate([idx[edges.b[ff]], idx[edges.a[ff]],
                           np.arange(n_free)])
    vals = np.concatenate([-edges.w[ff], -edges.w[ff], diag])
    A = sparse.coo_matrix((vals, (rows, cols)),
                          shape=(n_free, n_free)).tocsr()
    return A, rhs, idx


def solve_system(A: sparse.csr_matrix, b: np.ndarray) -> tuple[np.ndarray, SolveInfo]:
    n = A.shape[0]
    if n <= config.SPSOLVE_MAX_UNKNOWNS:
        x = sla.spsolve(A.tocsc(), b)
        return x, SolveInfo(method="spsolve", n_unknowns=n)

    # The matrix is SPD, so CG is guaranteed to converge. Jacobi is the
    # only preconditioner in scipy that keeps the preconditioned operator
    # SPD without a factorization that can break down at this scale.
    d = A.diagonal()
    M = sla.LinearOperator((n, n), lambda v: v / d)
    iters = 0

    def count(_):
        nonlocal iters
        iters += 1

    try:
        x, code = sla.cg(A, b, M=M, rtol=config.CG_TOL,
                         maxiter=config.CG_MAXITER, callback=count)
    except TypeError:  # scipy < 1.12 uses tol=
        x, code = sla.cg(A, b, M=M, tol=config.CG_TOL,
                         maxiter=config.CG_MAXITER, callback=count)
    if code != 0:
        raise RuntimeError(
            f"CG did not converge in {config.CG_MAXITER} iterations "
            f"(code {code}). Try a coarser grid or raise CG_MAXITER."
        )
    res = float(np.linalg.norm(b - A @ x) / np.linalg.norm(b))
    return x, SolveInfo(method="cg+jacobi", n_unknowns=n, iterations=iters,
                        residual=res)


def _face_current_density(V2: np.ndarray, mask2: np.ndarray, sigma: float,
                          h_m: float, t_m: float,
                          sig2d: np.ndarray | None = None,
                          rho: float | None = None) -> np.ndarray:
    """|J| (A/m^2) for one layer from face currents; V2 in volts.
    With a per-cell conductance map (buildup), face currents use the
    harmonic mean and J is referenced to the conductance-equivalent
    copper thickness t_eq = sigma_cell * rho (equals the geometric t for
    plain DC copper)."""
    ny, nx = mask2.shape
    face_x = mask2[:, :-1] & mask2[:, 1:]
    face_y = mask2[:-1, :] & mask2[1:, :]
    if sig2d is None:
        wx = wy = sigma
        teq = np.full((ny, nx), t_m)
    else:
        wx = 2.0 * sig2d[:, :-1] * sig2d[:, 1:] / (sig2d[:, :-1] + sig2d[:, 1:])
        wy = 2.0 * sig2d[:-1, :] * sig2d[1:, :] / (sig2d[:-1, :] + sig2d[1:, :])
        teq = sig2d * rho
    with np.errstate(invalid="ignore"):
        Ix = np.where(face_x, (V2[:, :-1] - V2[:, 1:]) * wx, 0.0)
        Iy = np.where(face_y, (V2[:-1, :] - V2[1:, :]) * wy, 0.0)
    IxP = np.zeros((ny, nx + 1))
    IxP[:, 1:nx] = Ix
    IyP = np.zeros((ny + 1, nx))
    IyP[1:ny, :] = Iy
    Jx = 0.5 * (IxP[:, :-1] + IxP[:, 1:])
    Jy = 0.5 * (IyP[:-1, :] + IyP[1:, :])
    Jmag = np.hypot(Jx, Jy) / (h_m * teq)
    Jmag[~mask2] = np.nan
    return Jmag


def _solve_equipotential(stack, e1, e2, edges):
    """Dirichlet terminals at 1 V / 0 V. Returns (Vflat_unit, R, I1, I2,
    mismatch, volts_per_amp, info). Fields at 1 V drive; scale by
    i_test * R to get volts at I_test."""
    if (layer := electrodes_touch(stack, e1, e2)) is not None:
        raise ElectrodeError(
            f"The terminals touch on {layer}. With the equipotential "
            f"contact model at least one cell of copper must separate "
            f"them; the uniform-injection model allows touching contacts."
        )
    state = np.zeros(stack.masks.size, dtype=np.uint8)
    state[stack.masks.ravel()] = 1
    state[e1.ravel()] = 2
    state[e2.ravel()] = 3

    A, rhs, idx = _assemble(state, edges, None)
    x, info = solve_system(A, rhs)

    Vflat = np.zeros(state.size)
    Vflat[state == 2] = 1.0
    Vflat[state == 1] = x

    Ie = edges.w * (Vflat[edges.a] - Vflat[edges.b])
    sa, sb = state[edges.a], state[edges.b]
    I1 = float(Ie[sa == 2].sum() - Ie[sb == 2].sum())
    I2 = float(Ie[sb == 3].sum() - Ie[sa == 3].sum())
    mismatch = abs(I1 - I2) / max(abs(I1), abs(I2), 1e-300)
    R = 1.0 / (0.5 * (I1 + I2))
    return Vflat, R, I1, I2, mismatch, R, info


def _solve_uniform(stack, e1, e2, edges):
    """Uniform orthogonal injection: every contact cell sources (sinks)
    1 A / N. Grounded at one V- cell. Returns like _solve_equipotential;
    fields are at 1 A drive, so volts_per_amp = 1."""
    n = stack.masks.size
    e1f, e2f = e1.ravel(), e2.ravel()
    n1, n2 = int(e1f.sum()), int(e2f.sum())

    inj = np.zeros(n)
    inj[e1f] = 1.0 / n1
    inj[e2f] = -1.0 / n2

    state = np.zeros(n, dtype=np.uint8)
    state[stack.masks.ravel()] = 1
    ground = int(np.flatnonzero(e2f)[0])
    state[ground] = 3                       # single Dirichlet 0 V reference;
    # its sink share is exactly the flux that exits through the reference,
    # so the grounded solution equals the pure-Neumann one

    A, rhs, idx = _assemble(state, edges, inj)
    x, info = solve_system(A, rhs)

    Vflat = np.zeros(n)
    Vflat[state == 1] = x

    v_plus = float(Vflat[e1f].mean())
    v_minus = float(Vflat[e2f].mean())
    R = (v_plus - v_minus) / 1.0
    Vflat = Vflat - v_minus                 # display reference: <V-> = 0

    # quality: KCL residual of the solved system
    res = info.residual
    if res is None:
        res = float(np.linalg.norm(A @ x - rhs)
                    / max(np.linalg.norm(rhs), 1e-300))
    return Vflat, R, 1.0, 1.0, res, 1.0, info


def _part_currents(parts, Ie, edges, e_flat, scale,
                   i_test, contact_model, n_terminal_cells):
    """Current through each contact part @ I_test. Equipotential: exact
    discrete flux out of the part's cells (same-terminal internal edges
    carry zero, opposite-terminal edges are forbidden). Uniform: the
    injection is prescribed, so a part carries exactly its cell share."""
    out = []
    for label, mask3 in parts:
        pf = mask3.ravel() & e_flat
        n = int(pf.sum())
        if contact_model == "uniform":
            amps = i_test * n / max(n_terminal_cells, 1)
        else:
            ina = pf[edges.a]
            inb = pf[edges.b]
            amps = abs(float(Ie[ina].sum() - Ie[inb].sum())) * scale
        out.append((label, amps))
    return out


def run_solve(problem: Problem, stack: RasterStack, e1: np.ndarray,
              e2: np.ndarray, i_test: float, freq_hz: float = 0.0,
              contact_model: str | None = None,
              parts1: list | None = None,
              parts2: list | None = None) -> Result:
    timings = {}
    L, ny, nx = stack.masks.shape
    h_m = stack.h_nm * 1e-9
    if contact_model is None:
        contact_model = config.CONTACT_MODEL

    # effective (AC) sheet conductances and barrel factor
    sigmas = [
        1.0 / skin.sheet_resistance_ac(
            problem.layers[li].thickness_nm * 1e-9, freq_hz,
            problem.rho_ohm_m, config.SKIN_SIDES)
        for li in range(L)
    ]
    rs_ratios = [
        skin.resistance_factor(problem.layers[li].thickness_nm * 1e-9,
                               freq_hz, problem.rho_ohm_m, config.SKIN_SIDES)
        for li in range(L)
    ]
    via_factor = skin.resistance_factor(problem.plating_nm * 1e-9, freq_hz,
                                        problem.rho_ohm_m, sides=2)
    sigma_buildup = 0.0
    if problem.buildups and stack.buildup is not None \
            and stack.buildup.any():
        sigma_buildup = 1.0 / skin.sheet_resistance_ac(
            problem.solder_thickness_nm * 1e-9, freq_hz,
            problem.solder_rho_ohm_m, config.SKIN_SIDES)
        if problem.extra_cu_nm > 0:
            sigma_buildup += 1.0 / skin.sheet_resistance_ac(
                problem.extra_cu_nm * 1e-9, freq_hz, problem.rho_ohm_m,
                config.SKIN_SIDES)
        eq_um = sigma_buildup * problem.rho_ohm_m * 1e6
        print(f"solder buildup: {problem.solder_thickness_nm / 1000:.0f} um "
              f"solder + {problem.extra_cu_nm / 1000:.0f} um Cu on "
              f"{int(stack.buildup.sum())} cells "
              f"(= {eq_um:.1f} um equivalent copper)")
    if freq_hz > 0:
        depth = skin.skin_depth_m(freq_hz, problem.rho_ohm_m)
        print(f"AC @ {freq_hz:g} Hz: skin depth {depth * 1e6:.0f} um, "
              f"per-layer Rs ratio "
              f"{', '.join(f'{r:.2f}' for r in rs_ratios)}, "
              f"via factor {via_factor:.2f}")

    t0 = time.perf_counter()
    edges = build_edges(stack, problem, sigmas, via_factor, sigma_buildup)
    if connected_restrict(stack, e1, e2, edges):
        edges = build_edges(stack, problem, sigmas, via_factor, sigma_buildup)
    if stack.buildup is not None:
        stack.buildup &= stack.masks
    for _, m in (parts1 or []) + (parts2 or []):
        m &= stack.masks                     # follow the component restriction
    timings["edges_s"] = time.perf_counter() - t0

    t0 = time.perf_counter()
    if contact_model == "equipotential":
        Vflat, R, I1, I2, mismatch, volts_per_amp, info = \
            _solve_equipotential(stack, e1, e2, edges)
    else:
        Vflat, R, I1, I2, mismatch, volts_per_amp, info = \
            _solve_uniform(stack, e1, e2, edges)
    timings["solve_s"] = time.perf_counter() - t0

    t0 = time.perf_counter()
    s = i_test * volts_per_amp              # unit-drive volts -> volts @ I_test

    # per-edge power @ I_test; distribute in-plane power to endpoint cells
    Pe = edges.w * ((Vflat[edges.a] - Vflat[edges.b]) * s) ** 2
    inplane = edges.via_index < 0
    Pflat = np.zeros(Vflat.size)
    np.add.at(Pflat, edges.a[inplane], 0.5 * Pe[inplane])
    np.add.at(Pflat, edges.b[inplane], 0.5 * Pe[inplane])
    Parea = Pflat.reshape(L, ny, nx) / (h_m * h_m)
    Parea[~stack.masks] = np.nan
    plane = ny * nx
    P_layers = [float(Pflat[li * plane:(li + 1) * plane].sum())
                for li in range(L)]
    P_vias = float(Pe[~inplane].sum())
    P_total = i_test ** 2 * R
    balance = abs((sum(P_layers) + P_vias) - P_total) / max(P_total, 1e-300)

    # via reports: max segment current + total power per via
    Ie = edges.w * (Vflat[edges.a] - Vflat[edges.b])   # amps at unit drive
    via_reports = []
    if problem.vias:
        vidx = edges.via_index
        for vi in np.unique(vidx[vidx >= 0]):
            sel = vidx == vi
            via = problem.vias[vi]
            via_reports.append(ViaReport(
                x_mm=via.x * 1e-6, y_mm=via.y * 1e-6, kind=via.kind,
                drill_mm=via.drill_nm * 1e-6,
                current_a=float(np.abs(Ie[sel]).max()) * s,
                power_w=float(Pe[sel].sum()),
            ))
        via_reports.sort(key=lambda v: v.current_a, reverse=True)

    # per-injection-area currents
    part_currents1 = _part_currents(
        parts1 or [], Ie, edges, e1.ravel(), s, i_test,
        contact_model, int(e1.sum()))
    part_currents2 = _part_currents(
        parts2 or [], Ie, edges, e2.ravel(), s, i_test,
        contact_model, int(e2.sum()))

    # embedded potential + per-layer current density @ I_test
    V3 = np.full((L, ny, nx), np.nan)
    V3[stack.masks] = Vflat.reshape(L, ny, nx)[stack.masks] * s
    J3 = np.stack([
        _face_current_density(
            np.nan_to_num(V3[li]), stack.masks[li], sigmas[li],
            h_m, problem.layers[li].thickness_nm * 1e-9,
            sig2d=_sigma_2d(stack, li, sigmas[li], sigma_buildup),
            rho=problem.rho_ohm_m)
        for li in range(L)
    ])
    timings["postprocess_s"] = time.perf_counter() - t0

    return Result(
        R_ohm=R, i_test=i_test, V=V3, Jmag=J3, Parea=Parea,
        layer_names=list(stack.layer_names),
        P_total=P_total, P_layers=P_layers, P_vias=P_vias,
        power_balance_rel=balance, via_reports=via_reports,
        I1_a=I1, I2_a=I2, mismatch_rel=mismatch,
        n_free=info.n_unknowns, solve_info=info,
        part_currents1=part_currents1, part_currents2=part_currents2,
        contact_model=contact_model,
        freq_hz=freq_hz,
        skin_depth_um=(skin.skin_depth_m(freq_hz, problem.rho_ohm_m) * 1e6
                       if freq_hz > 0 else None),
        rs_ratios=rs_ratios,
        timings=timings,
    )