import json
import random
import sys
import numpy as np
import torch
from torch import nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import train_model, evaluate, sweep_baseline, make_report
from tree_polynomial_track import get_dataset, META

SEEDS = tuple(range(8))
# The union is used for both methods, satisfying learning-rate parity.
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 24
NTR, NTE = 400, 200
YS = np.array([-3., -2., -1., -.5, .5, 1., 2., 3.], dtype=np.float64)


def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        try:
            torch.cuda.manual_seed_all(seed)
        except Exception:
            pass


def polynomial_values(tree_x):
    """Evaluate P(root; 1,y) for each encoded fixed-size rooted tree.

    Parent ids are stored as parent_id/5 and edge lengths in column 1.
    This is the exact product recursion, evaluated independently at eight y's.
    """
    n = len(tree_x)
    out = np.ones((n, len(YS)), dtype=np.float64)
    for i in range(n):
        parent = np.rint(tree_x[i, :, 0] * 5).astype(int)
        z = np.ones((6, len(YS)), dtype=np.float64)
        # Every non-root node is a leaf or has already been evaluated when used.
        for v in (5, 4, 3, 2, 1):
            children = np.flatnonzero(parent == v)
            if len(children) == 0:
                z[v] = 1.0
            else:
                z[v] = np.prod(tree_x[i, children, 1, None] * YS[None, :] + z[children], axis=0)
        children = np.flatnonzero(parent == 0)
        z[0] = 1.0 if len(children) == 0 else np.prod(
            tree_x[i, children, 1, None] * YS[None, :] + z[children], axis=0)
        out[i] = z[0]
    return out


def representation(ds, kind):
    q = dict(ds)
    if kind == "baseline":
        # Standard sum-style scalar message passing: root receives the sum of
        # edge lengths and terminal messages, which is blind to this topology.
        a = q["xtr"]
        b = q["xte"]
        q["xtr"] = a[:, :, 1].sum(1, keepdims=True)
        q["xte"] = b[:, :, 1].sum(1, keepdims=True)
    else:
        # Finite compressed branch-length polynomial fingerprint.
        q["xtr"] = np.log(np.abs(polynomial_values(q["xtr"])) + 1e-10).astype(np.float32)
        q["xte"] = np.log(np.abs(polynomial_values(q["xte"])) + 1e-10).astype(np.float32)
    q["xtr"] = torch.as_tensor(q["xtr"], dtype=torch.float32)
    q["xte"] = torch.as_tensor(q["xte"], dtype=torch.float32)
    q["ytr"] = torch.as_tensor(q["ytr"], dtype=torch.long)
    q["yte"] = torch.as_tensor(q["yte"], dtype=torch.long)
    q["input_shape"] = tuple(q["xtr"].shape[1:])
    q["out_dim"] = 2
    return q


def run(kind, lr, seed):
    seed_all(seed)
    ds = representation(get_dataset(seed, NTR, NTE), kind)
    model = nn.Sequential(nn.Flatten(), nn.Linear(int(np.prod(ds["input_shape"])), 32),
                          nn.ReLU(), nn.Linear(32, 2))
    _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=float(lr), batch=128,
                               log=lambda *args, **kwargs: None)
    return float(metric)


def base_factory(cfg):
    return lambda seed: run("baseline", cfg["lr"], seed)


def idea_factory(cfg):
    return lambda seed: run("idea", cfg["lr"], seed)


def mechanism_signature():
    """Retest the claimed topology sensitivity through trained model outputs."""
    seed_all(9001)
    ds0 = get_dataset(9001, 200, 200)
    db = representation(ds0, "baseline")
    di = representation(ds0, "idea")
    # Same branch-length multiset, one example of each topology.
    xb = torch.cat([db["xte"][0:1], db["xte"][1:2]])
    xi = torch.cat([di["xte"][0:1], di["xte"][1:2]])
    def fit(ds):
        model = nn.Sequential(nn.Flatten(), nn.Linear(int(np.prod(ds["input_shape"])), 32),
                              nn.ReLU(), nn.Linear(32, 2))
        net, _, _ = train_model(model, ds, epochs=EPOCHS, lr=3e-3, batch=128,
                                log=lambda *a, **k: None)
        return net
    mb, mi = fit(db), fit(di)
    devb = next(mb.parameters()).device
    devi = next(mi.parameters()).device
    with torch.no_grad():
        pb = torch.softmax(mb(xb.to(devb)), 1).cpu()
        pi = torch.softmax(mi(xi.to(devi)), 1).cpu()
    observed_baseline = float(torch.linalg.vector_norm(pb[0] - pb[1]))
    observed_idea = float(torch.linalg.vector_norm(pi[0] - pi[1]))
    # Prediction: polynomial representation should yield larger trained-output
    # separation than the sum representation on this structural pair.
    return {"prediction": "trained fingerprint outputs separate topology pair more than sum baseline",
            "observed_baseline_output_distance": observed_baseline,
            "observed_idea_output_distance": observed_idea,
            "predicted_ratio_lower_bound": 1.0,
            "observed_ratio": observed_idea / max(observed_baseline, 1e-12),
            "confirmed": bool(observed_idea > observed_baseline)}


def main():
    grid = [{"lr": lr} for lr in LRS]
    base = sweep_baseline(base_factory, grid, seeds=(0, 1, 2, 3))
    trials = [{"cfg": c, "result": evaluate(idea_factory(c), SEEDS)} for c in grid]
    best = min(trials, key=lambda z: z["result"]["mean"])
    report = make_report("rooted_tree_branch_classification", "mlp_tiny", base,
                         best["result"], {
                             "custom_track": {"name": META["name"], "file": "tree_polynomial_track.py", "domain": META["domain"]},
                             "idea_config": best["cfg"], "idea_sweep": trials,
                             "mechanism_signature": mechanism_signature()})
    with open("bench_report.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()