"""
emergence_audit plugin — AdS/deep-learning (Hashimoto, Sugishita, Tanaka & Tomiya 2018).
X = the bulk metric function h(eta) read off the layer weights. Uses the faithful toy in
verification/emergent-geometry/reproductions/.
"""
import os
import sys
import numpy as np
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
REPRO = os.path.join(HERE, "..", "..", "emergent-geometry", "reproductions")
sys.path.insert(0, REPRO)
from adsdl_lib import (ETA, DETA, make_dataset, make_dataset_nobulk, h_true, h_bump,   # noqa: E402
                       AdSDLNet, regulariser)
from einstein_residual import nearest_solution                                      # noqa: E402

EPOCHS = 60
eta = ETA[:-1]

generators = {
    "real":        lambda s: ("real", make_dataset(h_true, 1000, np.random.default_rng(100 + s))),
    "nonsolution": lambda s: ("bump", make_dataset(h_bump, 1000, np.random.default_rng(200 + s))),
    "stateless":   lambda s: ("nobulk", make_dataset_nobulk(1000, np.random.default_rng(300 + s))),
}
priors = {"smooth_1e-3": (1e-3, 0.0), "einstein_1e-4": (0.0, 1e-4)}
fit_tolerance = 0.25


def _einstein_reg(h):
    hp = (h[:-1] - h[1:]) / DETA
    hm = 0.5 * (h[:-1] + h[1:])
    return ((-9.0 + hp + hm ** 2) ** 2).sum()


def fit(data, prior, seed):
    tag, (X, y) = data
    c_s, c_e = prior
    net = AdSDLNet(seed=seed)
    g = torch.Generator().manual_seed(seed)
    Xt, yt = torch.tensor(X), torch.tensor(y)
    opt = torch.optim.Adam(net.parameters(), lr=1e-2)
    n = len(X)
    for _ in range(EPOCHS):
        perm = torch.randperm(n, generator=g)
        for b in range(0, n, 10):
            idx = perm[b:b + 10]
            opt.zero_grad()
            out = net.output(Xt[idx, 0], Xt[idx, 1])
            loss = (out - yt[idx]).abs().sum() + c_s * regulariser(net.h) + c_e * _einstein_reg(net.h)
            loss.backward(); opt.step()
    with torch.no_grad():
        out = net.output(Xt[:, 0], Xt[:, 1])
        fitq = float((out - yt).abs().mean())
    return net.h.detach().numpy().copy(), fitq


def feature(h):
    return np.asarray(h)


def family_distance(h):
    return nearest_solution(eta, h, 3)["distance"]


def generator_distance(h, data):
    tag = data[0]
    ref = {"real": h_true, "bump": h_bump}.get(tag)
    if ref is None:
        return float("nan")
    return float(np.sqrt(np.mean((h - ref(eta)) ** 2)))


def condition_number():
    # sensitivity of the horizon output to the ten weights, at the true metric (Jacobian cond)
    net = AdSDLNet(seed=0)
    with torch.no_grad():
        net.h.copy_(torch.tensor(h_true(eta)))
    X, _ = make_dataset(h_true, 200, np.random.default_rng(0))
    phi, pi = torch.tensor(X[:, 0]), torch.tensor(X[:, 1])
    J = torch.autograd.functional.jacobian(lambda h: _forward_with(net, h, phi, pi), net.h.detach())
    return float(np.linalg.cond(J.numpy()))


def _forward_with(net, h, phi, pi):
    for n in range(len(ETA) - 1):
        phi_new = phi - DETA * pi
        pi_new = pi + DETA * (h[n] * pi + phi - phi ** 3)   # m^2 = -1, lambda = 1
        phi, pi = phi_new, pi_new
    return pi