"""
AdS/DL (Hashimoto, Sugishita, Tanaka & Tomiya 2018, arXiv:1802.08313) at toy scale.

Static homogeneous scalar phi(eta) in ds^2 = -f dt^2 + d eta^2 + g dx^2 obeys (their eq. 5)
    d pi/d eta + h(eta) pi - m^2 phi - V'(phi) = 0,   pi = d phi/d eta,
    h(eta) = d/d eta log sqrt(f g^(d-1)).
Integrated INWARD from the boundary (eta = 1.0) to eta_fin = 0.1 in steps of 0.1:
    phi <- phi - d eta * pi
    pi  <- pi  + d eta * (h pi - m^2 phi - V'(phi))
True metric: AdS4-Schwarzschild, h(eta) = 3 coth(3 eta); m^2 = -1, V = phi^4/4 (lambda = 1).
Horizon output F = pi(eta_fin); smoothed indicator t(F) = [tanh(100(F-0.1)) - tanh(100(F+0.1)) + 2]/2,
positive data (true regular solutions, |F| < 0.1) labelled 0, negatives 1; loss sum |t - label|
+ c_reg sum_n eta_n^4 (h_{n+1} - h_n)^2.  Vacuum Einstein for this ansatz: -9 + h' + h^2 = 0 (their S.34).
"""
import numpy as np
import torch

torch.set_default_dtype(torch.float64)

ETA = np.round(np.arange(1.0, 0.1 - 1e-9, -0.1), 10)   # 1.0, 0.9, ..., 0.1  (10 slices)
DETA = 0.1
M2, LAM = -1.0, 1.0


def h_true(eta):
    return 3.0 / np.tanh(3.0 * eta)


def h_bump(eta, amp=2.0, center=0.55, width=0.1):
    """A metric function that is NOT a vacuum-Einstein solution."""
    return h_true(eta) + amp * np.exp(-0.5 * ((eta - center) / width) ** 2)


def propagate_np(h_vals, phi, pi):
    """h_vals: h at ETA[0..8] (the 9 steps from 1.0 to 0.1). Returns pi at eta_fin."""
    for n in range(len(ETA) - 1):
        h = h_vals[n]
        phi_new = phi - DETA * pi
        pi_new = pi + DETA * (h * pi - M2 * phi - LAM * phi ** 3)
        phi, pi = phi_new, pi_new
    return phi, pi


def make_dataset(h_fn, n_each=1000, rng=None, eps=0.1):
    """Rejection-sample n_each positives (|pi(eta_fin)| < eps) and negatives."""
    rng = rng or np.random.default_rng(0)
    h_vals = h_fn(ETA)
    pos, neg = [], []
    while len(pos) < n_each or len(neg) < n_each:
        phi = rng.uniform(0, 1.5, 20000)
        pi = rng.uniform(-0.2, 0.2, 20000)
        _, pf = propagate_np(h_vals, phi, pi)
        ok = np.abs(pf) < eps
        for x, y, o in zip(phi, pi, ok):
            if o and len(pos) < n_each:
                pos.append((x, y))
            elif (not o) and len(neg) < n_each:
                neg.append((x, y))
            if len(pos) >= n_each and len(neg) >= n_each:
                break
    X = np.array(pos + neg)
    y = np.array([0.0] * n_each + [1.0] * n_each)
    return X, y


def make_dataset_nobulk(n_each=1000, rng=None, band=0.03):
    """Positives = a wiggly band in the (phi, pi) plane generated by NO bulk at all:
    |pi - g(phi)| < band,  g(phi) = 0.12 sin(4 phi) - 0.04."""
    rng = rng or np.random.default_rng(0)
    g = lambda p: 0.12 * np.sin(4 * p) - 0.04
    pos, neg = [], []
    while len(pos) < n_each or len(neg) < n_each:
        phi = rng.uniform(0, 1.5, 20000)
        pi = rng.uniform(-0.2, 0.2, 20000)
        ok = np.abs(pi - g(phi)) < band
        for x, yv, o in zip(phi, pi, ok):
            if o and len(pos) < n_each:
                pos.append((x, yv))
            elif (not o) and len(neg) < n_each:
                neg.append((x, yv))
            if len(pos) >= n_each and len(neg) >= n_each:
                break
    X = np.array(pos + neg)
    y = np.array([0.0] * n_each + [1.0] * n_each)
    return X, y


class AdSDLNet(torch.nn.Module):
    def __init__(self, seed=0):
        super().__init__()
        g = torch.Generator().manual_seed(seed)
        init = torch.tensor(1.0 / ETA[:-1]) + torch.randn(len(ETA) - 1, generator=g)
        self.h = torch.nn.Parameter(init)

    def forward(self, phi, pi):
        for n in range(len(ETA) - 1):
            phi_new = phi - DETA * pi
            pi_new = pi + DETA * (self.h[n] * pi - M2 * phi - LAM * phi ** 3)
            phi, pi = phi_new, pi_new
        return pi                                         # F = pi(eta_fin)

    def output(self, phi, pi):
        F = self.forward(phi, pi)
        return (torch.tanh(100 * (F - 0.1)) - torch.tanh(100 * (F + 0.1)) + 2) / 2


def regulariser(h):
    eta = torch.tensor(ETA[:-1])
    return (eta[:-1] ** 4 * (h[1:] - h[:-1]) ** 2).sum()


def train(net, X, y, c_reg=1e-3, epochs=100, batch=10, lr=1e-2, stop=2e-4, seed=0):
    g = torch.Generator().manual_seed(seed)
    X_t = torch.tensor(X); y_t = torch.tensor(y)
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    n = len(X)
    for ep in range(epochs):
        perm = torch.randperm(n, generator=g)
        tot = 0.0
        for b in range(0, n, batch):
            idx = perm[b:b + batch]
            opt.zero_grad()
            out = net.output(X_t[idx, 0], X_t[idx, 1])
            loss = (out - y_t[idx]).abs().sum() + c_reg * regulariser(net.h)
            loss.backward()
            opt.step()
            tot += float((out - y_t[idx]).abs().sum())
        mean_loss = tot / n
        if mean_loss < stop:
            break
    with torch.no_grad():
        out = net.output(X_t[:, 0], X_t[:, 1])
        acc = float(((out > 0.5).double() == y_t).double().mean())
        final = float((out - y_t).abs().mean())
    return final, acc, ep + 1


def einstein_residual(h_vals):
    """-9 + h' + h^2 at the slice midpoints; h' via finite difference along decreasing eta."""
    h = np.asarray(h_vals)
    hp = (h[:-1] - h[1:]) / DETA        # dh/d eta  (eta decreases with index)
    hm = 0.5 * (h[:-1] + h[1:])
    return -9.0 + hp + hm ** 2