"""Geometry-in -> results-out pipeline shared by the KiCad entrypoint and
the offline standalone runner."""
from __future__ import annotations

from pathlib import Path

import numpy as np

from . import config, plots, progress, raster, report, solver, trim
from .errors import ElectrodeError, UserFacingError
from .geometry import Problem
from .solver import Result


def run(problem: Problem, outdir: Path | None, show: bool = True,
        i_test: float | None = None, freq_hz: float = 0.0,
        contact_model: str | None = None, overlay=None,
        trim_pct: float | None = None, trim_abs: float | None = None,
        trim_push=None, v_nominal: float | None = None) -> Result:
    """overlay: optional callback(stack, result) run after the solve
    (EXPERIMENTAL in-KiCad overlays); its failures are non-fatal.
    trim_pct / trim_abs: mark copper below this threshold (% of the
    mean |J| / absolute A/mm2; at most one, both None = off): per-layer
    areas are printed, polygons saved to
    <outdir>/low_current_copper.json and handed to trim_push, an
    optional callback(trim_result) that pushes them into the board
    (failures non-fatal). Problems with terminals run in PDN mode:
    i_test/contact_model are ignored there (draws come from the
    terminals, the contact models are fixed) and v_nominal is the
    default supply open-circuit voltage."""
    pdn = bool(problem.terminals)
    if pdn and (problem.electrodes1 or problem.electrodes2):
        raise ElectrodeError(
            "The problem carries both classic V+/V- electrodes and PDN "
            "terminals - exactly one terminal scheme must be used."
        )
    if not pdn:
        if i_test is None:
            i_test = config.TEST_CURRENT_A
        if i_test <= 0:
            raise UserFacingError(
                f"Test current must be > 0 A (got {i_test:g}).")
    h = raster.choose_cell_size(problem.copper_bbox(), len(problem.layers))
    progress.stage(f"rasterizing {len(problem.layers)} layer(s) at cell "
                   f"size {h / 1000:.1f} um ...")
    stack = raster.rasterize_stack(problem, h)
    print(f"grid {stack.shape2d[1]}x{stack.shape2d[0]}x{stack.nlayers}, "
          f"{int(stack.masks.sum())} copper cells, {len(problem.vias)} "
          f"via/pad barrel(s)")

    if pdn:
        if contact_model is not None:
            print("PDN mode: contact models are fixed (Thevenin supplies "
                  "/ uniform-injection loads) - ignoring the setting")
        tmasks = raster.terminal_masks(stack, problem)
        tparts = raster.terminal_partition(stack, problem)
        draw = sum(t.i_draw_a for t in problem.terminals
                   if t.role == "load")
        progress.stage(f"solving PDN, {draw:g} A total draw"
                       + (f", {freq_hz:g} Hz" if freq_hz > 0 else " DC")
                       + " ...")
        result = solver.run_solve_pdn(problem, stack, tmasks, tparts,
                                      freq_hz, v_nominal)
        for s_ in result.supplies:
            print(f"  {s_.label}: {s_.i_a:.4g} A @ {s_.v_contact:.4g} V "
                  f"(v_oc {s_.v_oc:g} V, r_out {s_.r_out_ohm:g} ohm, "
                  f"P_int {s_.p_internal_w:.3g} W)")
        for l_ in result.loads:
            print(f"  {l_.label}: {l_.i_a:.4g} A, V {l_.v_mean:.4g} V "
                  f"(min {l_.v_min:.4g}), P {l_.p_w:.4g} W")
        # figures reuse the two-terminal color scheme: supplies as V+,
        # loads as V- (masks already follow the solve's restriction)
        e1 = np.zeros_like(stack.masks)
        e2 = np.zeros_like(stack.masks)
        for t, m in zip(problem.terminals, tmasks):
            if t.role == "supply":
                e1 |= m
            else:
                e2 |= m
    else:
        e1, e2 = raster.electrode_masks(stack, problem)
        parts1, parts2 = raster.electrode_partition(stack, problem)

        progress.stage(f"solving @ {i_test:g} A"
                       + (f", {freq_hz:g} Hz" if freq_hz > 0 else " DC")
                       + " ...")
        result = solver.run_solve(problem, stack, e1, e2, i_test, freq_hz,
                                  contact_model, parts1, parts2)
        for prefix, pcs in (("P", result.part_currents1),
                            ("N", result.part_currents2)):
            for i, (label, amps) in enumerate(pcs):
                print(f"  {prefix}{i + 1} ({label}): {amps:.4g} A "
                      f"({100 * amps / i_test:.1f}%)")

    if outdir is not None:
        outdir.mkdir(parents=True, exist_ok=True)
        report.write_summary(outdir, problem, stack, result)
    print(report.result_line(result, problem, stack))

    if overlay is not None:
        try:
            overlay(stack, result)
        except Exception as e:
            print(f"overlay push failed: {e}")

    if trim_pct is not None or trim_abs is not None:
        tr = trim.compute(result, stack, pct=trim_pct, abs_a_mm2=trim_abs)
        print(trim.summary_line(tr))
        if outdir is not None:
            trim.write_json(outdir, tr)
        if trim_push is not None:
            try:
                trim_push(tr)
            except Exception as e:
                print(f"trim push failed: {e}")

    progress.stage("rendering figures ...")
    figs = [
        (plots.fig_raster(stack, e1, e2, problem, result), "1_raster_map"),
        (plots.fig_potential(result, stack, e1, e2, problem), "2_potential"),
        (plots.fig_current(result, stack, e1, e2, problem),
         "3_current_density"),
        (plots.fig_power(result, stack, e1, e2, problem), "4_power_density"),
    ]
    if result.mode == "pdn" and result.pairs:
        figs.append((plots.fig_pdn_pairs(result), "5_source_sink_pairs"))
    plots.save_and_show(figs, outdir, show=show)   # closes the window itself
    return result