"""Balanced AC + rectangular-voltage GNN + physics-informed training.

Pure Python standard library. Three-phase 1 MVA / 0.4 kV LL bases; PQ net
consumption is positive. Series-only lines; one fixed V=1+j0 reference. The
inverter is fixed PQ. The learner predicts V without calling solve_ac().
Thermal/voltage limits classify states; they are NOT penalties in the loss.
"""
import math

H, K, V_SCALE, S_SCALE = 6, 4, 0.05, 0.2
I_BASE = 1000/(math.sqrt(3)*0.4)
DEFAULTS = dict(load_scale=1.2, bus5_kw=100., pv4_kw=50., z_scale=1.,
                topology="looped", rating_scale=1., vmin=.95, vmax=1.05)
TRAINING_DEFAULTS = dict(epochs=400, lr=.003, **{"lambda": 1.}, label_fraction=.25, noise=.002)
SIZES = dict(Win=H*3, bin=H, Ws=H*H, Wn=H*H, Wskip=H*3, bh=H, We=H*2, Wo=2*H, bo=2)
KEYS = list(SIZES)


def zeros():
    return {key: [0.]*n for key, n in SIZES.items()}


def copy_weights(theta):
    return {key: theta[key][:] for key in KEYS}


class Random:
    """Same seeded 32-bit LCG and Box-Muller transform as the browser."""
    def __init__(self, seed):
        self.state = int(seed) & 0xffffffff

    def uniform(self):
        self.state = (1664525*self.state+1013904223) & 0xffffffff
        return (self.state+0.5)/4294967296

    def normal(self):
        return math.sqrt(-2*math.log(self.uniform()))*math.cos(2*math.pi*self.uniform())


def initialize():
    rng, theta = Random(7), zeros()
    scales = dict(Win=.3, Ws=.04, Wn=.04, Wskip=.08, We=.02, Wo=.04)
    for key in ("Win", "Ws", "Wn", "Wskip", "We", "Wo"):
        theta[key] = [rng.normal()*scales[key] for _ in theta[key]]
    for j in range(H):
        theta["Ws"][j*H+j] += .5
        theta["Wn"][j*H+j] += .3
    return theta


def check_weights(theta):
    if not theta or not all(key in theta and len(theta[key]) == n
                            and all(math.isfinite(v) and abs(v) < 100 for v in theta[key])
                            for key, n in SIZES.items()):
        raise ValueError("Invalid model weights")
    return theta


def config(parameters=None):
    c = dict(DEFAULTS, **(parameters or {}))
    for key, lo, hi in (("load_scale", 0, 2.5), ("bus5_kw", 0, 300), ("pv4_kw", 0, 300),
                        ("z_scale", .3, 2), ("rating_scale", .3, 2), ("vmin", .8, 1), ("vmax", 1, 1.2)):
        if not math.isfinite(c[key]) or not lo <= c[key] <= hi:
            raise ValueError("Invalid " + key)
    if c["topology"] not in ("looped", "feeder", "radial") or c["vmin"] >= c["vmax"]:
        raise ValueError("Invalid topology or voltage bounds")
    return c


def build_graph(parameters=None, outage="base"):
    """Rebuild active arcs AND Ybus after every operating/topology change.

    Relative to the preceding five-bus lesson, line 1-3 adds a second path to
    the same external reference. 'looped' closes both 1-3 and 3-5; 'feeder'
    opens 1-3; 'radial' opens both. Only base-case active lines are N-1 candidates.
    """
    c = config(parameters)
    p = [0., .08*c["load_scale"], .1*c["load_scale"],
         .06*c["load_scale"]-c["pv4_kw"]/1000, c["bus5_kw"]/1000*c["load_scale"]]
    q = [0., .03*c["load_scale"], .04*c["load_scale"],
         .025*c["load_scale"], .4*c["bus5_kw"]/1000*c["load_scale"]]
    nodes = [dict(id=i+1, p=p[i], q=q[i], ref=i == 0,
                  x=[p[i]/.2, q[i]/.1, 1 if i == 0 else 0]) for i in range(5)]
    raw = [("12", 0, 1, .04, .025, .5, True), ("23", 1, 2, .06, .035, .22, True),
           ("24", 1, 3, .05, .03, .22, True), ("45", 3, 4, .055, .035, .18, True),
           ("35", 2, 4, .08, .05, .18, c["topology"] != "radial"),
           ("13", 0, 2, .07, .045, .45, c["topology"] == "looped")]
    edges = [dict(id=name, **{"from": f, "to": t}, r=r*c["z_scale"], x=x*c["z_scale"],
                  limit=limit*c["rating_scale"], base_active=active, active=active and name != outage)
             for name, f, t, r, x, limit, active in raw]
    if outage != "base" and not any(e["id"] == outage and e["base_active"] for e in edges):
        raise ValueError("Outage must be an in-service base-case line")
    g, b, links, degree = [[0.]*5 for _ in range(5)], [[0.]*5 for _ in range(5)], [], [0]*5
    for e in edges:
        den = e["r"]**2+e["x"]**2
        e["g"], e["b"] = e["r"]/den, -e["x"]/den
        if not e["active"]:
            continue
        f, t = e["from"], e["to"]
        for i, j, sign in ((f, f, 1), (t, t, 1), (f, t, -1), (t, f, -1)):
            g[i][j] += sign*e["g"]
            b[i][j] += sign*e["b"]
        for source, target in ((f, t), (t, f)):
            links.append(dict(source=source, target=target,
                              gate=1/(1+math.hypot(e["r"], e["x"])/.05), e=[e["r"]/.05, e["x"]/.05]))
            degree[target] += 1
    seen = {0}
    for _ in range(5):
        for e in edges:
            if e["active"] and (e["from"] in seen or e["to"] in seen):
                seen.update((e["from"], e["to"]))
    return dict(case=c, outage=outage, nodes=nodes, edges=edges, g=g, b=b, links=links,
                degree=degree, unsupplied=[v["id"] for i, v in enumerate(nodes) if i not in seen])


def power(graph, voltage):
    """S_inj = V * conjugate(YV), using real/imaginary scalar arithmetic."""
    ir, ii = [0.]*5, [0.]*5
    for i in range(5):
        for j in range(5):
            ir[i] += graph["g"][i][j]*voltage[j][0]-graph["b"][i][j]*voltage[j][1]
            ii[i] += graph["b"][i][j]*voltage[j][0]+graph["g"][i][j]*voltage[j][1]
    return dict(ir=ir, ii=ii, p=[u*ir[i]+w*ii[i] for i, (u, w) in enumerate(voltage)],
                q=[w*ir[i]-u*ii[i] for i, (u, w) in enumerate(voltage)])


def residual(graph, voltage, with_jacobian=False):
    """PQ residuals only: reference P/Q supply is unknown, not forced to zero."""
    calculated = power(graph, voltage)
    rp = [calculated["p"][i]+graph["nodes"][i]["p"] for i in range(1, 5)]
    rq = [calculated["q"][i]+graph["nodes"][i]["q"] for i in range(1, 5)]
    values = rp+rq
    out = dict(rp=rp, rq=rq, values=values, max_pu=max(map(abs, values)),
               rms_pu=math.sqrt(sum(v*v for v in values)/8),
               loss=sum((v/S_SCALE)**2 for v in values)/8)
    if with_jacobian:
        jacobian = [[0.]*10 for _ in range(8)]
        for i in range(1, 5):
            for j in range(5):
                G, B = graph["g"][i][j], graph["b"][i][j]
                u, w, delta = *voltage[i], int(i == j)
                jacobian[i-1][j] = delta*calculated["ir"][i]+u*G+w*B
                jacobian[i-1][5+j] = delta*calculated["ii"][i]-u*B+w*G
                jacobian[4+i-1][j] = -delta*calculated["ii"][i]+w*G-u*B
                jacobian[4+i-1][5+j] = delta*calculated["ir"][i]-w*B-u*G
        out["jacobian"] = jacobian
    return out


def linear_solve(matrix, rhs):
    a, n = [list(row)+[v] for row, v in zip(matrix, rhs)], len(rhs)
    for k in range(n):
        pivot = max(range(k, n), key=lambda i: abs(a[i][k]))
        if abs(a[pivot][k]) < 1e-12:
            raise ArithmeticError("Singular Jacobian")
        a[k], a[pivot] = a[pivot], a[k]
        for i in range(k+1, n):
            factor = a[i][k]/a[k][k]
            for j in range(k, n+1):
                a[i][j] -= factor*a[k][j]
    answer = [0.]*n
    for i in range(n-1, -1, -1):
        answer[i] = (a[i][n]-sum(a[i][j]*answer[j] for j in range(i+1, n)))/a[i][i]
    return answer


def newton_direction(r):
    return linear_solve([row[1:5]+row[6:10] for row in r["jacobian"]], [-v for v in r["values"]])


def trial_step(voltage, step, alpha):
    return [[1., 0.]]+[[voltage[i][0]+alpha*step[i-1], voltage[i][1]+alpha*step[4+i-1]] for i in range(1, 5)]


def solve_ac(graph, max_iterations=30):
    """Damped rectangular Newton reference. A failure returns unknown, not unsafe."""
    if graph["unsupplied"]:
        return dict(ok=False, reason="unsupplied", unsupplied=graph["unsupplied"], v=None)
    voltage = [[1., 0.] for _ in range(5)]
    for k in range(max_iterations+1):
        r = residual(graph, voltage, True)
        if r["max_pu"] < 1e-10:
            return dict(ok=True, v=voltage, residual_pu=r["max_pu"], iterations=k)
        if k == max_iterations:
            break
        try:
            step = newton_direction(r)
        except ArithmeticError:
            break
        accepted, alpha = False, 1.
        while alpha >= 1/128:
            trial = trial_step(voltage, step, alpha)
            if all(math.hypot(*u) > .1 for u in trial) and residual(graph, trial)["max_pu"] < r["max_pu"]:
                voltage, accepted = trial, True
                break
            alpha /= 2
        if not accepted:
            break
    return dict(ok=False, reason="nonconvergence", unsupplied=[], v=None)


def diagnostics(graph, voltage):
    vm = [math.hypot(*u) for u in voltage]
    angles = [math.atan2(u[1], u[0])*180/math.pi for u in voltage]
    branches = []
    for e in graph["edges"]:
        if not e["active"]:
            branches.append(dict(id=e["id"], active=False, current_pu=0., current_a=0., loading_pct=0., loss_pu=0.))
            continue
        dr = voltage[e["from"]][0]-voltage[e["to"]][0]
        di = voltage[e["from"]][1]-voltage[e["to"]][1]
        current = math.hypot(e["g"]*dr-e["b"]*di, e["b"]*dr+e["g"]*di)
        branches.append(dict(id=e["id"], active=True, current_pu=current, current_a=current*I_BASE,
                             loading_pct=current/e["limit"]*100, loss_pu=e["r"]*current*current))
    violations = [dict(kind="voltage", id=i+1) for i, value in enumerate(vm)
                  if value < graph["case"]["vmin"]-1e-9 or value > graph["case"]["vmax"]+1e-9]
    violations += [dict(kind="current", id=e["id"]) for e in branches if e["active"] and e["loading_pct"] > 100+1e-7]
    return dict(v=voltage, vm=vm, angles=angles, branches=branches, min_v=min(vm),
                max_loading=max(e["loading_pct"] for e in branches), residual=residual(graph, voltage),
                violations=violations, status="unsafe" if violations else "safe")


def refine(graph, input_voltage):
    """Optional ONE accepted Newton correction, explicitly separate from the GNN."""
    voltage = [u[:] for u in input_voltage]
    r = residual(graph, voltage, True)
    before = r["max_pu"]
    if before < 1e-10:
        return dict(v=voltage, accepted=True, before_pu=before, after_pu=before)
    try:
        step = newton_direction(r)
    except ArithmeticError:
        return dict(v=voltage, accepted=False, before_pu=before, after_pu=before)
    alpha = 1.
    while alpha >= 1/128:
        trial = trial_step(voltage, step, alpha)
        after = residual(graph, trial)["max_pu"]
        if all(math.hypot(*u) > .1 for u in trial) and after < before:
            return dict(v=trial, accepted=True, before_pu=before, after_pu=after)
        alpha /= 2
    return dict(v=voltage, accepted=False, before_pu=before, after_pu=before)


def predict(graph, theta, cache=False):
    """Feature/edge -> four tied message layers -> complex voltage; no AC solve.

    h0=tanh(Win*x+bin), m=gate*h_sender+We*[r/.05,x/.05]
    h_new=tanh(Ws*h+Wn*mean(messages)+Wskip*x+bh)
    V_PQ=[1+.05*y0, .05*y1], y=Wo*hK+bo; reference V=1+j0 is imposed exactly.
    Matrix rows are flattened to lists for transparent cross-language parity.
    """
    check_weights(theta)
    features = [v["x"] for v in graph["nodes"]]
    layers = [[[math.tanh(theta["bin"][a]+sum(theta["Win"][a*3+f]*x for f, x in enumerate(row)))
                for a in range(H)] for row in features]]
    aggregates = []
    for k in range(K):
        old, agg = layers[k], [[0.]*H for _ in range(5)]
        for edge in graph["links"]:
            for a in range(H):
                agg[edge["target"]][a] += (edge["gate"]*old[edge["source"]][a]
                    + theta["We"][a*2]*edge["e"][0]+theta["We"][a*2+1]*edge["e"][1])/graph["degree"][edge["target"]]
        aggregates.append(agg)
        new = []
        for i, row in enumerate(old):
            state = []
            for a in range(H):
                z = theta["bh"][a]
                for f in range(H):
                    z += theta["Ws"][a*H+f]*row[f]+theta["Wn"][a*H+f]*agg[i][f]
                for f in range(3):
                    z += theta["Wskip"][a*3+f]*features[i][f]
                state.append(math.tanh(z))
            new.append(state)
        layers.append(new)
    voltage = [[1., 0.]]
    for row in layers[K][1:]:
        voltage.append([1+V_SCALE*(theta["bo"][0]+sum(theta["Wo"][f]*x for f, x in enumerate(row))),
                        V_SCALE*(theta["bo"][1]+sum(theta["Wo"][H+f]*x for f, x in enumerate(row)))])
    return dict(v=voltage, layers=layers, aggregates=aggregates) if cache else voltage


def objective(graph, voltage, observed=None, physics_weight=1., with_gradient=False):
    """Dimensionless supervised and AC residual MSE, with explicit voltage gradients.

    P/Q mask excludes the reference. Data mask excludes unlabeled graphs.
    Limits do not appear here: unsafe operating points still obey AC equations.
    """
    r = residual(graph, voltage, with_gradient and physics_weight != 0)
    data, dv = 0., [[0., 0.] for _ in range(5)]
    if observed is not None:
        for i in range(1, 5):
            for f in range(2):
                error = (voltage[i][f]-observed[i][f])/V_SCALE
                data += error*error/8
                if with_gradient:
                    dv[i][f] += (voltage[i][f]-observed[i][f])/(4*V_SCALE*V_SCALE)
    if with_gradient and physics_weight:
        for i in range(8):
            coefficient = physics_weight*r["values"][i]/(4*S_SCALE*S_SCALE)
            for j in range(1, 5):
                dv[j][0] += coefficient*r["jacobian"][i][j]
                dv[j][1] += coefficient*r["jacobian"][i][5+j]
    return dict(data=data, physics=r["loss"], total=data+physics_weight*r["loss"], gradient=dv, residual=r)


def backward(graph, theta, passed, dv):
    """Chain rule: dL/dV -> readout -> every message layer -> input encoder."""
    grad, dh = zeros(), [[0.]*H for _ in range(5)]
    for i in range(1, 5):
        for f in range(2):
            d = V_SCALE*dv[i][f]
            grad["bo"][f] += d
            for a in range(H):
                grad["Wo"][f*H+a] += d*passed["layers"][K][i][a]
                dh[i][a] += d*theta["Wo"][f*H+a]
    for k in reversed(range(K)):
        previous, da = [[0.]*H for _ in range(5)], [[0.]*H for _ in range(5)]
        for i in range(5):
            for a in range(H):
                d = dh[i][a]*(1-passed["layers"][k+1][i][a]**2)
                grad["bh"][a] += d
                for f in range(H):
                    grad["Ws"][a*H+f] += d*passed["layers"][k][i][f]
                    grad["Wn"][a*H+f] += d*passed["aggregates"][k][i][f]
                    previous[i][f] += d*theta["Ws"][a*H+f]
                    da[i][f] += d*theta["Wn"][a*H+f]
                for f in range(3):
                    grad["Wskip"][a*3+f] += d*graph["nodes"][i]["x"][f]
        for edge in graph["links"]:
            for a in range(H):
                d = da[edge["target"]][a]/graph["degree"][edge["target"]]
                previous[edge["source"]][a] += edge["gate"]*d
                grad["We"][a*2] += d*edge["e"][0]
                grad["We"][a*2+1] += d*edge["e"][1]
        dh = previous
    for i in range(5):
        for a in range(H):
            d = dh[i][a]*(1-passed["layers"][0][i][a]**2)
            grad["bin"][a] += d
            for f in range(3):
                grad["Win"][a*3+f] += d*graph["nodes"][i]["x"][f]
    return grad


def loss_gradient(sample, theta, physics_weight=1., data_weight=1.):
    passed = predict(sample["graph"], theta, True)
    data = objective(sample["graph"], passed["v"], sample["observed"], 0., True)
    physics = objective(sample["graph"], passed["v"], None, physics_weight, True)
    dv = [[data_weight*v+physics["gradient"][i][f] for f, v in enumerate(row)]
          for i, row in enumerate(data["gradient"])]
    return dict(data=data["data"], physics=physics["physics"], loss=data_weight*data["data"]+physics_weight*physics["physics"],
                gradient=backward(sample["graph"], theta, passed, dv))


def dataset(groups, seed, options=None, training=False):
    c = dict(TRAINING_DEFAULTS, **(options or {}))
    rng, noise_rng, samples, failed = Random(seed), Random(seed ^ 0x9e3779b9), [], 0
    labeled_groups = math.floor(groups*c["label_fraction"]+0.5)
    for group in range(groups):
        operating = dict(load_scale=.6+.9*rng.uniform(), bus5_kw=60+100*rng.uniform(),
                         pv4_kw=100*rng.uniform(), z_scale=.8+.4*rng.uniform())
        base = build_graph(operating)
        outages = ["base"]+[e["id"] for e in base["edges"] if e["base_active"]]
        for outage in outages:
            graph = build_graph(operating, outage)
            labeled = not training or group < labeled_groups
            truth = solve_ac(graph) if labeled else None
            if truth and not truth["ok"]:
                failed += 1
                continue
            observed = [[v+c["noise"]*noise_rng.normal() if i and training else v for v in u]
                        for i, u in enumerate(truth["v"])] if labeled else None
            samples.append(dict(group=group, graph=graph, truth=truth["v"] if truth else None, observed=observed))
    return dict(samples=samples, groups=groups, labeled_groups=labeled_groups if training else groups, failed=failed)


def evaluate(data, theta):
    mae, physics, unsafe, false_safe, correct = 0., 0., 0, 0, 0
    for sample in data["samples"]:
        pred = predict(sample["graph"], theta)
        ref, out = diagnostics(sample["graph"], sample["truth"]), diagnostics(sample["graph"], pred)
        mae += sum(abs(out["vm"][i]-ref["vm"][i]) for i in range(1, 5))
        physics += out["residual"]["loss"]
        if ref["status"] == "unsafe":
            unsafe += 1
            false_safe += out["status"] == "safe"
        correct += ref["status"] == out["status"]
    n = len(data["samples"])
    return dict(voltage_mae_pu=mae/(4*n), physics_mse=physics/n,
                residual_rms_pu=S_SCALE*math.sqrt(physics/n), false_safe=false_safe,
                unsafe=unsafe, correct=correct, cases=n)


def train_model(options=None, progress=None):
    """Matched Adam experiments; labeled graphs and collocation graphs have separate means.

    Train/validation split operating groups BEFORE expanding each into base +
    six outages. Unlabeled clean voltages are not consumed by loss_gradient.
    Gradient norm is clipped at 10; Adam beta1=.9, beta2=.999, epsilon=1e-8.
    """
    c = dict(TRAINING_DEFAULTS, **(options or {}))
    if (not isinstance(c["epochs"], int) or not 0 <= c["epochs"] <= 1200
            or not math.isfinite(c["lr"]) or not 0 < c["lr"] <= .02
            or c["label_fraction"] not in (0, .25, .5, 1)
            or not math.isfinite(c["lambda"]) or not 0 <= c["lambda"] <= 10
            or not math.isfinite(c["noise"]) or not 0 <= c["noise"] <= .01):
        raise ValueError("Invalid training settings")
    training, validation = dataset(8, 17, c, True), dataset(3, 83, c, False)
    n = len(training["samples"])
    nlabel = sum(s["observed"] is not None for s in training["samples"])
    theta, moment, variance, history = initialize(), zeros(), zeros(), []

    def record(epoch):
        data, physics = 0., 0.
        for sample in training["samples"]:
            loss = objective(sample["graph"], predict(sample["graph"], theta), sample["observed"], c["lambda"])
            if sample["observed"] is not None:
                data += loss["data"]/(nlabel or 1)
            physics += loss["physics"]/n
        history.append(dict(epoch=epoch, data=data, physics=physics, total=data+c["lambda"]*physics,
                            validation=evaluate(validation, theta)))

    record(0)
    for epoch in range(1, c["epochs"]+1):
        grad = zeros()
        for sample in training["samples"]:
            if c["lambda"] == 0 and sample["observed"] is None:
                continue
            out = loss_gradient(sample, theta, c["lambda"], n/nlabel if nlabel else 0)
            for key, size in SIZES.items():
                for j in range(size):
                    grad[key][j] += out["gradient"][key][j]/n
        norm = math.sqrt(sum(sum(x*x for x in grad[key]) for key in KEYS))
        clip = min(1., 10/(norm or 1))
        for key, size in SIZES.items():
            for j in range(size):
                g = grad[key][j]*clip
                moment[key][j] = .9*moment[key][j]+.1*g
                variance[key][j] = .999*variance[key][j]+.001*g*g
                theta[key][j] -= c["lr"]*(moment[key][j]/(1-.9**epoch))/(math.sqrt(variance[key][j]/(1-.999**epoch))+1e-8)
        check_weights(theta)
        if epoch % 20 == 0 or epoch == c["epochs"]:
            record(epoch)
            if progress:
                progress(epoch, c["epochs"])
    return dict(ok=True, weights=theta, config=c, history=history, validation=evaluate(validation, theta),
                training_cases=n, labeled_cases=nlabel, training_groups=8, validation_groups=3,
                failed_labels=training["failed"]+validation["failed"], parameter_count=sum(SIZES.values()),
                trained=c["lambda"] > 0 or nlabel > 0)


def screen(parameters=None, models=None, include_reference=True, correct_physics=False):
    """Enumerate base + each active line outage. No redispatch or dynamics.

    Unsafe: known unsupplied buses or a violated voltage/current criterion.
    Unknown: failed AC reference or missing learned weights. Low residual alone
    is not a safety certificate. The optional Newton correction is reported.
    """
    c = config(parameters)
    base = build_graph(c)
    outages = ["base"]+[e["id"] for e in base["edges"] if e["base_active"]]
    cases = []
    for outage in outages:
        graph = build_graph(c, outage)
        row = dict(outage=outage, unsupplied=graph["unsupplied"], graph=graph)
        if graph["unsupplied"]:
            for kind in ("reference", "data", "physics"):
                row[kind] = dict(status="unsafe", reason="unsupplied")
        else:
            ref = solve_ac(graph) if include_reference else None
            row["reference"] = (diagnostics(graph, ref["v"]) if ref["ok"] else dict(status="unknown", reason=ref["reason"])) if ref else None
            for kind in ("data", "physics"):
                if not models or not models.get(kind) or models[kind].get("trained") is False:
                    row[kind] = dict(status="unknown", reason="untrained")
                    continue
                model = models[kind]
                raw = predict(graph, model.get("weights", model))
                if kind == "physics" and correct_physics:
                    correction = refine(graph, raw)
                    row[kind] = dict(diagnostics(graph, correction["v"]), raw_v=raw, correction=correction)
                else:
                    row[kind] = diagnostics(graph, raw)
        cases.append(row)

    def summarize(kind):
        if any(row.get(kind) and row[kind]["status"] == "unsafe" for row in cases):
            return "unsafe"
        return "unknown" if any(not row.get(kind) or row[kind]["status"] == "unknown" for row in cases) else "safe"

    return dict(ok=True, case=c, cases=cases,
                summary={kind: summarize(kind) for kind in ("reference", "data", "physics")})