"""Five-bus graph, explicit message passing, and a seven-parameter GNN.

Only the standard library is used. A flat-voltage linearization generates
voltage labels; it is not full nonlinear AC power flow. All P/Q values are
three-phase totals on a 1 MVA base; impedances use a 0.4 kV LL voltage base.
Positive node P/Q denote net consumption. PV at bus 4 is fixed PQ, Q=0.
"""
import math

DEFAULTS = dict(load_scale=1., bus5_kw=100., pv4_kw=50., z_scale=1.,
                steps=3, aggregation="mean", edge_aware=True,
                open12=False, open23=False, open24=False, open45=False, tie35=False)
INITIAL_THETA = dict(alpha=0.5, beta=0.6, gamma=0.4, w=[0.2, 0.2, -0.1], b=0.)


def copy_theta(theta):
    return dict(theta, w=theta["w"][:])


def config(parameters=None):
    c = dict(DEFAULTS, **(parameters or {}))
    for key, lo, hi in (("load_scale", 0, 3), ("bus5_kw", 0, 300),
                        ("pv4_kw", 0, 300), ("z_scale", 0.1, 3)):
        if not math.isfinite(c[key]) or not lo <= c[key] <= hi:
            raise ValueError("Invalid " + key)
    if (not isinstance(c["steps"], int) or isinstance(c["steps"], bool)
            or not 1 <= c["steps"] <= 4 or c["aggregation"] not in ("sum", "mean")):
        raise ValueError("Invalid message-passing configuration")
    for key in ("edge_aware", "open12", "open23", "open24", "open45", "tie35"):
        if not isinstance(c[key], bool):
            raise ValueError("Expected a boolean: " + key)
    return c


def check_theta(theta):
    values = [theta["alpha"], theta["beta"], theta["gamma"], theta["b"], *theta["w"]]
    if len(theta["w"]) != 3 or not all(math.isfinite(v) and abs(v) <= 100 for v in values):
        raise ValueError("Invalid GNN weights")
    return copy_theta(theta)


def build_graph(parameters=None):
    """Encode specified net PQ, reference flag, and physical line parameters.

    Bus labels 1..5 map to indices 0..4. Solved voltages and bus IDs are not
    features. Slack P/Q=0 are placeholders, not its solved supply.
    """
    c = config(parameters)
    p = [v/1000 for v in (0, 80*c["load_scale"], 100*c["load_scale"],
                          60*c["load_scale"]-c["pv4_kw"], c["bus5_kw"]*c["load_scale"])]
    q = [v/1000 for v in (0, 30*c["load_scale"], 40*c["load_scale"],
                          25*c["load_scale"], 0.4*c["bus5_kw"]*c["load_scale"])]
    nodes = [dict(id=i+1, p_pu=p[i], q_pu=q[i], reference=int(i == 0),
                  x=[p[i]/0.2, q[i]/0.1, int(i == 0)]) for i in range(5)]
    raw = [("12", 0, 1, .04, .025, not c["open12"]),
           ("23", 1, 2, .06, .035, not c["open23"]),
           ("24", 1, 3, .05, .03, not c["open24"]),
           ("45", 3, 4, .055, .035, not c["open45"]),
           ("35", 2, 4, .08, .05, c["tie35"])]
    edges = [dict(id=name, **{"from": f, "to": t}, r=r*c["z_scale"],
                  x=x*c["z_scale"], active=active) for name, f, t, r, x, active in raw]
    return dict(nodes=nodes, edges=edges)


def arcs(graph, edge_aware):
    """Two computational directions per active undirected line; no self-loops."""
    result = []
    for e in graph["edges"]:
        if e["active"]:
            gate = 1/(1+math.hypot(e["r"], e["x"])/0.05) if edge_aware else 1.
            for source, target in ((e["from"], e["to"]), (e["to"], e["from"])):
                result.append(dict(edge=e["id"], source=source, target=target, gate=gate))
    return result


def distances(graph, source):
    d = [None]*len(graph["nodes"])
    d[source] = 0
    for _ in graph["nodes"]:
        for e in graph["edges"]:
            if not e["active"]:
                continue
            for u, v in ((e["from"], e["to"]), (e["to"], e["from"])):
                if d[u] is not None and (d[v] is None or d[v] > d[u]+1):
                    d[v] = d[u]+1
    return d


def forward(graph, parameters=None, theta=None):
    """Synchronous layers: every message reads the previous layer only.

    m[j->i] = gate[e] * h[j]
    aggregate = incoming sum or mean (zero for an empty neighborhood)
    h_new[i] = tanh(alpha*h[i] + beta*aggregate[i] + gamma*x[i])
    alpha/beta/gamma and the 3-channel readout are tied across all layers/nodes.
    """
    c, theta = config(parameters), check_theta(theta or INITIAL_THETA)
    n, links = len(graph["nodes"]), arcs(graph, c["edge_aware"])
    degree = [0]*n
    for a in links:
        degree[a["target"]] += 1
    features = [v["x"][:] for v in graph["nodes"]]
    layers, messages, aggregates = [features], [], []
    for k in range(c["steps"]):
        old, agg, msg = layers[k], [[0.]*3 for _ in range(n)], []
        for a in links:
            value = [a["gate"]*v for v in old[a["source"]]]
            for f in range(3):
                agg[a["target"]][f] += value[f]
            msg.append(dict(a, value=value))
        if c["aggregation"] == "mean":
            agg = [[v/degree[i] if degree[i] else 0. for v in row] for i, row in enumerate(agg)]
        layers.append([[math.tanh(theta["alpha"]*v + theta["beta"]*agg[i][f]
                                  + theta["gamma"]*features[i][f])
                        for f, v in enumerate(row)] for i, row in enumerate(old)])
        messages.append(msg)
        aggregates.append(agg)
    output = [theta["b"]+sum(w*h for w, h in zip(theta["w"], row)) for row in layers[-1]]
    return dict(layers=layers, messages=messages, aggregates=aggregates, degree=degree,
                output=output, raw_voltage=[1. if graph["nodes"][i]["reference"] else 1.-0.05*v
                                           for i, v in enumerate(output)])


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 linear reference")
        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 reference(graph):
    """Y_LL * deltaV = -conj(S_load), flat source V=1 and magnitude ≈1+Re(deltaV).

    Series lines only; losses and second-order terms are omitted. The connected
    graph may have loops. If any bus is unsupplied, the whole comparison pauses.
    This is not a voltage reference for a sustained island with its own source.
    """
    source = next(i for i, v in enumerate(graph["nodes"]) if v["reference"])
    reachable = distances(graph, source)
    unsupplied = [v["id"] for i, v in enumerate(graph["nodes"]) if reachable[i] is None]
    if unsupplied:
        return dict(ok=False, unsupplied=unsupplied, voltage=None, reason="island")
    pq = [i for i in range(len(graph["nodes"])) if i != source]
    index, n = {v: i for i, v in enumerate(pq)}, len(pq)
    g, b = [[0.]*n for _ in pq], [[0.]*n for _ in pq]
    for e in graph["edges"]:
        if not e["active"]:
            continue
        den = e["r"]**2+e["x"]**2
        gr, bi = e["r"]/den, -e["x"]/den
        f, t = e["from"], e["to"]
        for u, v, sign in ((f, f, 1), (t, t, 1), (f, t, -1), (t, f, -1)):
            if u in index and v in index:
                g[index[u]][index[v]] += sign*gr
                b[index[u]][index[v]] += sign*bi
    matrix = [g[i]+[-v for v in b[i]] for i in range(n)]+[b[i]+g[i] for i in range(n)]
    rhs = [-graph["nodes"][i]["p_pu"] for i in pq]+[graph["nodes"][i]["q_pu"] for i in pq]
    delta = linear_solve(matrix, rhs)
    voltage = [1.]*len(graph["nodes"])
    for i, bus in enumerate(pq):
        voltage[bus] = 1+delta[i]
    return dict(ok=True, unsupplied=[], voltage=voltage, delta_re=delta[:n], delta_im=delta[n:], pq=pq)


def solve(parameters=None):
    c = config(parameters)
    graph, theta = build_graph(c), check_theta(c.get("theta", INITIAL_THETA))
    passed, ref = forward(graph, c, theta), reference(graph)
    adjacency = [[int(any(e["active"] and ((e["from"] == i and e["to"] == j)
                      or (e["to"] == i and e["from"] == j)) for e in graph["edges"]))
                  for j in range(5)] for i in range(5)]
    links = arcs(graph, c["edge_aware"])
    mae = sum(abs(a-b) for i, (a, b) in enumerate(zip(passed["raw_voltage"], ref["voltage"]))
              if not graph["nodes"][i]["reference"])/4 if ref["ok"] else None
    return dict(ok=True, case=c, graph=graph, theta=theta, **passed, reference=ref,
                adjacency=adjacency, edge_index=[[a["source"] for a in links], [a["target"] for a in links]],
                prediction=passed["raw_voltage"] if ref["ok"] else None, mae_pu=mae)


def dataset(count, seed):
    state = int(seed) & 0xffffffff

    def random():
        nonlocal state
        state = (1664525*state+1013904223) & 0xffffffff
        return (state+0.5)/4294967296

    result = []
    for _ in range(count):
        c = dict(load_scale=0.6+0.8*random(), bus5_kw=60+100*random(),
                 pv4_kw=100*random(), z_scale=0.8+0.4*random())
        graph = build_graph(c)
        result.append(dict(graph=graph, target=reference(graph)["voltage"]))
    return result


def loss_gradient(graph, target, parameters, theta):
    """Backpropagate through readout and ALL message-passing layers.

    Seven learned scalars: alpha, beta, gamma, w[0:3], b. Contributions to tied
    parameters sum over every node, edge, and layer where they were used.
    Loss excludes the reference bus; voltage drops are scaled by 0.05 pu.
    """
    c, passed = config(parameters), forward(graph, parameters, theta)
    n, count = len(graph["nodes"]), len(graph["nodes"])-1
    g, loss, dh = dict(alpha=0., beta=0., gamma=0., w=[0.]*3, b=0.), 0., [[0.]*3 for _ in range(n)]
    for i in range(n):
        if graph["nodes"][i]["reference"]:
            continue
        error = passed["output"][i]-(1-target[i])/0.05
        d = 2*error/count
        loss += error*error/count
        g["b"] += d
        for f in range(3):
            g["w"][f] += d*passed["layers"][-1][i][f]
            dh[i][f] = d*theta["w"][f]
    for k in reversed(range(c["steps"])):
        previous = [[0.]*3 for _ in range(n)]
        dz = [[v*(1-passed["layers"][k+1][i][f]**2) for f, v in enumerate(row)] for i, row in enumerate(dh)]
        for i in range(n):
            for f in range(3):
                g["alpha"] += dz[i][f]*passed["layers"][k][i][f]
                g["beta"] += dz[i][f]*passed["aggregates"][k][i][f]
                g["gamma"] += dz[i][f]*graph["nodes"][i]["x"][f]
                previous[i][f] += theta["alpha"]*dz[i][f]
        for a in passed["messages"][k]:
            divisor = passed["degree"][a["target"]] if c["aggregation"] == "mean" else 1
            for f in range(3):
                previous[a["source"]][f] += theta["beta"]*a["gate"]*dz[a["target"]][f]/divisor
        dh = previous
    return dict(loss=loss, gradient=g)


def evaluate(data, parameters, theta):
    loss, mae = 0., 0.
    for sample in data:
        passed = forward(sample["graph"], parameters, theta)
        for i, node in enumerate(sample["graph"]["nodes"]):
            if not node["reference"]:
                error = passed["raw_voltage"][i]-sample["target"][i]
                loss += (error/0.05)**2
                mae += abs(error)
    n = len(data)*(len(data[0]["graph"]["nodes"])-1)
    return dict(mse=loss/n, mae_pu=mae/n)


def train_model(parameters=None):
    """Train all 7 parameters on 32 radial cases; validate on 16 new cases.

    The data distribution is fixed and independent of the current lab operating
    controls. Hyperparameters steps/aggregation/edge_aware select the model.
    Topology-changing cases are not part of this training/validation distribution.
    """
    c = config(parameters)
    epochs, lr = c.get("epochs", 400), c.get("lr", 0.03)
    if not isinstance(epochs, int) or not 0 <= epochs <= 1000 or not math.isfinite(lr) or not 0 < lr <= .2:
        raise ValueError("Invalid training settings")
    training, validation = dataset(32, 17), dataset(16, 83)
    theta, history = copy_theta(INITIAL_THETA), []

    def record(epoch):
        history.append(dict(epoch=epoch, train=evaluate(training, c, theta), validation=evaluate(validation, c, theta)))

    record(0)
    for epoch in range(1, epochs+1):
        g = dict(alpha=0., beta=0., gamma=0., w=[0.]*3, b=0.)
        for sample in training:
            d = loss_gradient(sample["graph"], sample["target"], c, theta)["gradient"]
            for key in ("alpha", "beta", "gamma", "b"):
                g[key] += d[key]/len(training)
            for f in range(3):
                g["w"][f] += d["w"][f]/len(training)
        for key in ("alpha", "beta", "gamma", "b"):
            theta[key] -= lr*g[key]
        for f in range(3):
            theta["w"][f] -= lr*g["w"][f]
        check_theta(theta)
        if epoch % 10 == 0 or epoch == epochs:
            record(epoch)
    return dict(ok=True, theta=theta, history=history, training_count=32, validation_count=16,
                epochs=epochs, lr=lr, config={key: c[key] for key in ("steps", "aggregation", "edge_aware")}, parameter_count=7)