#!/usr/bin/env python3
"""ngspice-hier-op -- incremental ("hierarchical") operating-point solver.

Produce AN operating point for a netlist that a plain `.op` cannot solve, at any
cost, WITHOUT analysing what the circuit does.  Not the perfect OP -- a starting
point the designer takes over from.  The baseline alternative is no OP at all.

Method: never hand the solver the whole problem at once.  Take a subset that
converges, save its state, add ONE block back, re-solve seeded from that state.
Every solve is then a small perturbation of an already-solved circuit.

Per block, stop at the first rung that yields a state:

  a  .op            seeded with the incoming state as .nodeset
  b  .tran ... uic  seeded with the incoming state as .ic, t = 50n
       -> save endpoint -> .op from that endpoint
       -> OP fails: double t and repeat, up to 4 doublings (50n 100n 200n 400n)
  c  ladder exhausted -> hand over the best transient endpoint anyway, labelled
       'transient' rather than 'op'.  Never return nothing.

`tran uic` is the escalation because it is ROBUST -- it always produces a state.
It is NOT a functional check; nothing here inspects or interprets behaviour.
Non-clocked blocks settle in ~50n or never; for clocked blocks the time is
unknowable without the designer, so the doubling ladder is a bounded search.

Backtracking: if a block fails BOTH rungs, the previous step's transient may have
been too short.  Re-run the IMMEDIATELY preceding step at the longer durations and
retry the failing block once; still failing -> stop.  One backtrack, not a search
(going further back would invalidate every step folded after it).

YOUR FILES ARE NEVER MODIFIED.  Blocks are masked with the `disable` command, in
the netlist ngspice holds in memory; the seed and the analysis are appended in a
generated file passed alongside your deck.

Usage:
  ngspice-hier-op DECK [options]
  hier_op DECK              (alias inside ngspice: `shell ngspice-hier-op`)
"""
import argparse, json, os, re, subprocess, sys, time

DOUBLINGS = 4
T0 = 50e-9
RSHUNT = 1e9
MIN_DEVICES = 300
PRIM = set("mrcdqivejfhglkbstw")

FAIL_MARKS = (
    "simulation(s) aborted", "could not be simulated successfully",
    "doAnalyses:", "Timestep too small", "no writable vector found",
    "incomplete or empty netlist", "circuit not parsed",
    "Unable to find definition of model",
)


def say(*a):
    print(*a, flush=True)


# ---------- netlist parsing (read-only; the deck is never written) ----------

def logical_lines(path):
    buf = None
    for raw in open(path, errors="replace"):
        ln = raw.rstrip("\n")
        if ln.lstrip().startswith("+") and buf is not None:
            buf += " " + ln.lstrip()[1:].strip()
            continue
        if buf is not None:
            yield buf
        buf = ln
    if buf is not None:
        yield buf


def top_instances(cir):
    out, depth = [], 0
    for ln in logical_lines(cir):
        st = ln.lstrip(); low = st.lower()
        if low.startswith(".subckt"): depth += 1; continue
        if low.startswith(".ends"):   depth -= 1; continue
        if depth == 0 and st[:1] in ("x", "X"):
            toks = [t for t in st.split() if "=" not in t]
            if len(toks) >= 3:
                out.append((toks[0], toks[-1], toks[1:-1]))
    return out


def device_counts(cir):
    defs, cur, depth, top = {}, None, 0, []
    for ln in logical_lines(cir):
        st = ln.lstrip(); low = st.lower()
        if low.startswith(".subckt"):
            depth += 1
            if depth == 1:
                cur = {"kids": [], "prims": 0}
                defs[st.split()[1]] = cur
            continue
        if low.startswith(".ends"):
            depth -= 1
            if depth == 0: cur = None
            continue
        c = st[:1].lower()
        if not c or c in ("*", "."):
            continue
        if c == "x":
            toks = [t for t in st.split() if "=" not in t]
            if len(toks) >= 3:
                (cur["kids"] if cur else top).append((toks[0], toks[-1]))
        elif c in PRIM and cur:
            cur["prims"] += 1
    memo = {}
    def total(name):
        if name in memo: return memo[name]
        d = defs.get(name)
        if not d: return 0
        memo[name] = 0
        n = d["prims"] + sum(total(t) for _i, t in d["kids"])
        memo[name] = n
        return n
    return {i: total(t) for i, t in top}


def port_nets(cir, inst):
    for name, _t, nets in top_instances(cir):
        if name.lower() == inst.lower():
            return {n.lower() for n in nets}
    return set()


# ---------- state files ----------

_ST = re.compile(r"^\s*\.(ic|nodeset)\s+v\((.+?)\)\s*=\s*(\S+)", re.I)


def read_state(path):
    out = []
    for ln in open(path, errors="replace"):
        m = _ST.match(ln)
        if m:
            out.append((m.group(2).strip(), m.group(3)))
    return out


def write_state(path, entries, form):
    with open(path, "w") as f:
        for net, val in entries:
            f.write(f".{form} v({net}) = {val}\n")
    return len(entries)


def finite(v):
    try:
        x = float(v)
    except ValueError:
        return False
    return x == x and x not in (float("inf"), float("-inf"))


def emit_both(src, base, drop=()):
    """Every step leaves BOTH forms -- .nodeset seeds a .op, .ic seeds a tran uic.
    Feeding the wrong one is a SILENT no-op, so never make that a judgement call.

    Non-finite entries are dropped.  A node the solver could not resolve comes
    back as nan/inf, and seeding the next step with `.ic v(x) = -nan` is strictly
    worse than saying nothing about that node: the value propagates into the
    matrix instead of leaving the solver free to work it out.
    """
    d = {n.lower() for n in drop}
    kept, bad = [], []
    for n, v in read_state(src):
        if n.lower() in d:
            continue
        (kept if finite(v) else bad).append((n, v))
    write_state(base + ".ic", kept, "ic")
    write_state(base + ".nodeset", kept, "nodeset")
    return len(kept), len(bad)


# ---------- running ----------

def run_attempt(a, tag, drop, seed, analysis):
    """One ngspice run: generated file carries the seed + .control; deck untouched.

    Everything the generated file names is a BARE lowercase filename, and ngspice
    is run with its working directory set to outdir.  ngspice lowercases netlist
    content, so an absolute path written into a .control is silently mangled the
    moment any component contains an upper-case letter (mktemp -d names do) and
    wrnodev then writes nowhere.
    """
    tag = tag.lower()
    inc_name, state_name = tag + ".inc", tag + ".state"
    inc = os.path.join(a.outdir, inc_name)
    state = os.path.join(a.outdir, state_name)
    log = os.path.join(a.outdir, tag + ".log")
    if os.path.exists(state):
        os.remove(state)
    with open(inc, "w") as f:
        f.write("* generated by ngspice-hier-op -- your deck is NOT modified\n")
        if a.rshunt:
            # Removing a block leaves the nets it drove with nothing holding
            # them, so they go singular and come back as nan.  rshunt ties every
            # node to ground through a large resistance; it is structural to this
            # method rather than a tuning knob.  --rshunt 0 turns it off.
            f.write(f".option rshunt={a.rshunt:g}\n")
        # pstran adds a pseudo-transient rung below the built-in ladder, giving
        # a seeded .op one more way to land before we fall through to rung (b).
        # Independent of rshunt: --rshunt 0 must not silently take this with it.
        f.write(".option pstran\n")
        if seed:
            f.write(f".include {os.path.basename(seed)}\n")
        f.write(".control\n")
        if drop:
            f.write("disable " + " ".join(drop) + "\nreset\n")
        for c in analysis:
            f.write(c + "\n")
        # `quit` so the deck's OWN analysis cards (a .op/.tran the user already
        # has) do not run after our block: they would double every attempt's
        # cost, and a failure of theirs lands in the same log we judge ours by,
        # marking a perfectly good state as FAILED.
        f.write(f"wrnodev {state_name}\nquit\n.endc\n")
    t0 = time.time()
    with open(log, "w") as fh:
        subprocess.run([a.ngspice, "-b", a.deck, inc_name], cwd=a.outdir,
                       stdout=fh, stderr=subprocess.STDOUT)
    secs = time.time() - t0

    # A seeded node with no matrix diagonal (e.g. its only element was stripped
    # by topology reduction) cannot be seeded under KLU: unlike SPARSE, KLU
    # cannot create the element, and ngspice aborts with a misleading
    # "doAnalyses: out of memory".  Seeding is the whole method here, so fall
    # back to SPARSE for this run only -- the deck's own solver choice is left
    # alone everywhere it works.
    txt = open(log, errors="replace").read()
    if seed and "cannot create a new element" in txt:
        say("    (KLU cannot seed an element-less node; retrying with .option sparse)")
        body = open(inc).read().split("\n")
        body.insert(1, ".option sparse")
        open(inc, "w").write("\n".join(body))
        with open(log, "w") as fh:
            subprocess.run([a.ngspice, "-b", a.deck, inc_name], cwd=a.outdir,
                           stdout=fh, stderr=subprocess.STDOUT)
        secs = time.time() - t0
        txt = open(log, errors="replace").read()

    ok, why = True, "converged"
    if not (os.path.exists(state) and os.path.getsize(state) > 0 and read_state(state)):
        ok, why = False, "no state written"
    else:
        for m in FAIL_MARKS:
            if m in txt:
                ok, why = False, m
                break
    say(f"    [{tag.split('__')[-1]}] {'OK ' if ok else 'FAIL'} {secs:7.1f}s  {why}")
    return ok, state, secs, why


def fmt(t):
    return f"{t*1e9:g}n"


def fold(a, blk, drop, state_in, tag):
    """One block: walk rung a -> b(xN doublings) -> c.  Returns (label, base, rec)."""
    rec = {"block": blk, "drop": list(drop), "attempts": []}
    seed_ns = seed_ic = None
    if state_in:
        # the added block's port nets sat at whatever its ABSENCE imposed -- drop
        # them so the incoming state does not fight the block being added
        nets = port_nets(a.deck, blk) if blk else set()
        base = os.path.join(a.outdir, tag + "__seed")
        n, bad = emit_both(state_in + ".ic", base, drop=nets)
        seed_ns, seed_ic = base + ".nodeset", base + ".ic"
        say(f"   seed: {n} nodes ({len(nets)} of {blk}'s port nets dropped)"
            if blk else f"   seed: {n} nodes")
        if bad:
            say(f"   !! {bad} node(s) came back non-finite (nan/inf) and were NOT seeded.")
            say( "      Removing a block leaves the nets it drove floating; with no")
            say( "      `.option rshunt` nothing ties them down, so they go singular.")

    ok, st, secs, why = run_attempt(a, tag + "__a", drop, seed_ns, ["op"])
    rec["attempts"].append({"rung": "a", "t": None, "ok": ok, "secs": secs, "why": why})
    if ok:
        out = os.path.join(a.outdir, tag)
        emit_both(st, out)
        rec.update(result="op", rung="a")
        return "op", out, rec

    best, t = None, a.t0
    for i in range(a.doublings):
        ok, st, secs, why = run_attempt(
            a, f"{tag}__b{i}_tran{fmt(t)}", drop, seed_ic,
            [f"tran {t/100:.4g} {t:.4g} uic"])
        rec["attempts"].append({"rung": "b.tran", "t": t, "ok": ok, "secs": secs, "why": why})
        if not ok:
            say(f"    tran produced no state at {fmt(t)}; not doubling further")
            break
        best = st
        mid = os.path.join(a.outdir, f"{tag}__b{i}_end")
        emit_both(st, mid)
        ok2, st2, secs2, why2 = run_attempt(
            a, f"{tag}__b{i}_op", drop, mid + ".nodeset", ["op"])
        rec["attempts"].append({"rung": "b.op", "t": t, "ok": ok2, "secs": secs2, "why": why2})
        if ok2:
            out = os.path.join(a.outdir, tag)
            emit_both(st2, out)
            rec.update(result="op", rung=f"b@{fmt(t)}")
            return "op", out, rec
        t *= 2

    if best:
        say("    ladder exhausted; delivering the transient endpoint as the state")
        out = os.path.join(a.outdir, tag)
        emit_both(best, out)
        rec.update(result="transient", rung="c")
        return "transient", out, rec

    rec.update(result="both-failed", rung=None)
    return None, None, rec


def last_tran_t(rec):
    ts = [x["t"] for x in rec.get("attempts", []) if x["rung"] == "b.tran" and x["ok"] and x["t"]]
    return max(ts) if ts else None


def main():
    ap = argparse.ArgumentParser(
        prog="ngspice-hier-op",
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""examples:
  ngspice-hier-op chip.cir --list         show the fold order, run nothing
  ngspice-hier-op chip.cir                fold every block, chain the states
  ngspice-hier-op chip.cir --t0 5e-6      start clocked blocks at 5us, not 50n

inside ngspice (or a .control block):
  hier_op chip.cir

output (in --outdir, default ./hier_op):
  <step>.ic        state as .ic       -- seeds a `.tran ... uic`
  <step>.nodeset   same state         -- seeds a `.op`
  <step>.json      which rung produced it, and what each attempt cost
  <step>__*.log    the ngspice run behind every attempt

Each state is labelled `op` (a converged operating point) or `transient` (an
endpoint delivered because the ladder was exhausted).  A `transient` state is a
usable starting point, not an operating point -- the label is the difference.""")
    ap.add_argument("deck", help="your netlist; it is NEVER modified")
    ap.add_argument("--ngspice", default="ngspice", help="ngspice binary")
    ap.add_argument("--outdir", default="hier_op", help="where states/logs are written")
    ap.add_argument("--min-devices", type=int, default=MIN_DEVICES,
                    help="ignore top-level instances smaller than this; below it the "
                         "segmentation is too fine to be worth a step (default %(default)s)")
    ap.add_argument("--rshunt", type=float, default=RSHUNT,
                    help="resistance from every node to ground, added to each run; "
                         "removing blocks leaves floating nodes, so this is "
                         "structural here (0 disables, default %(default)g)")
    ap.add_argument("--t0", type=float, default=T0, help="first transient length")
    ap.add_argument("--doublings", type=int, default=DOUBLINGS)
    ap.add_argument("--list", action="store_true", help="show fold order and exit")
    a = ap.parse_args()

    a.deck = os.path.abspath(a.deck)
    if not os.path.exists(a.deck):
        sys.exit(f"no such deck: {a.deck}")
    # A relative --ngspice would be resolved against whatever directory each run
    # happens to start in, not the one the user typed it from.  A bare name is
    # left alone so PATH lookup still works.
    if os.sep in a.ngspice:
        a.ngspice = os.path.abspath(a.ngspice)
    os.makedirs(a.outdir, exist_ok=True)
    a.outdir = os.path.abspath(a.outdir)

    dc = device_counts(a.deck)
    cand = sorted(((n, c) for n, c in dc.items() if c >= a.min_devices),
                  key=lambda kv: kv[1])
    if not cand:
        sys.exit(f"no top-level instances with >= {a.min_devices} devices in {a.deck}")
    if a.list:
        say(f"fold candidates (>= {a.min_devices} devices), smallest-first:")
        for n, c in cand:
            say(f"  {c:8d}  {n}")
        return 0

    order = [n for n, _c in cand]
    say(f"deck   : {a.deck}  (never modified)")
    say(f"blocks : {len(order)} to fold, smallest-first: {', '.join(order)}")

    out = list(order)
    say("\n--- base: all fold candidates disabled, no seed")
    label, state, rec = fold(a, None, out, None, "base")
    json.dump(rec, open(os.path.join(a.outdir, "base.json"), "w"), indent=2)
    if not label:
        say("!! the base subset itself produced no state -- nothing to build on")
        return 2
    say(f"   -> base state [{label}]")

    folded, prev, backtracked, stopped = [], None, False, None
    i = 0
    while i < len(order):
        blk = order[i]
        rest = [b for b in out if b != blk]
        tag = f"s{i+1}_{blk}"
        say(f"\n--- [{i+1}/{len(order)}] fold {blk}")
        step = {"blk": blk, "tag": tag, "state_in": state, "drop": rest}
        label, newstate, rec = fold(a, blk, rest, state, tag)
        json.dump(rec, open(os.path.join(a.outdir, tag + ".json"), "w"), indent=2)

        if label:
            step["rec"] = rec
            state, out = newstate, rest
            folded.append((blk, label, rec.get("rung")))
            say(f"   -> state [{label}] via rung {rec.get('rung')}")
            prev = step
            i += 1
            continue

        say(f"    {blk} failed BOTH .op and tran uic")
        if backtracked:
            stopped = f"{blk} still fails after one backtrack"; break
        if prev is None:
            stopped = f"{blk} fails and there is no preceding step to lengthen"; break
        t_used = last_tran_t(prev["rec"])
        used = sum(1 for x in prev["rec"]["attempts"] if x["rung"] == "b.tran")
        if not t_used:
            stopped = (f"{blk} fails and {prev['blk']} came from .op -- "
                       f"no transient to lengthen"); break
        if used >= a.doublings:
            stopped = f"{blk} fails and {prev['blk']} already used all its doublings"; break
        say(f"    backtracking to {prev['blk']}: re-running from {fmt(t_used*2)}")
        a2 = argparse.Namespace(**vars(a)); a2.t0 = t_used * 2
        a2.doublings = a.doublings - used
        l2, s2, r2 = fold(a2, prev["blk"], prev["drop"], prev["state_in"],
                          prev["tag"] + "_bt")
        backtracked = True
        if not l2:
            stopped = f"backtrack of {prev['blk']} produced no state"; break
        state = s2; prev["rec"] = r2
        say(f"    {prev['blk']} re-solved longer; retrying {blk}")

    say("\n=== summary ===")
    for blk, res, rung in folded:
        say(f"  folded {blk:14s} [{res} via {rung}]")
    say(f"  final state : {state}.ic  /  {state}.nodeset")
    say(f"  still out   : {', '.join(out) or 'none'}")
    if stopped:
        say(f"  STOPPED     : {stopped}")
        say("  (every state produced up to here is still on disk and usable)")
    return 1 if stopped else 0


if __name__ == "__main__":
    sys.exit(main())
