import json, math, random
import numpy as np
import torch
import torch.nn as nn

SEED = 2359
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type == 'cuda': torch.cuda.get_device_properties(0)
except Exception:
    device = torch.device('cpu')


def exact_gaussian_ot(x, m1, s1, m2, s2):
    # In Euclidean space, for diagonal covariances, T(x)=m2+D2 D1^-1 (x-m1).
    return (x - m1) * (s2 / s1) + m2


def gaussian_sanity():
    g = torch.Generator().manual_seed(SEED)
    m1 = torch.tensor([-0.7, 0.4]); s1 = torch.tensor([1.3, 0.7])
    m2 = torch.tensor([1.1, -0.8]); s2 = torch.tensor([0.8, 1.4])
    x = torch.randn(20000, 2, generator=g) * s1 + m1
    y = exact_gaussian_ot(x, m1, s1, m2, s2)
    mean_err = float((y.mean(0) - m2).norm())
    cov = torch.cov(y.T)
    cov_err = float((cov - torch.diag(s2*s2)).norm())
    delta = torch.tensor([2.0, -1.5]); d2 = float(delta @ delta)
    scaling = []
    for L in [1, 2, 4, 8, 16]:
        observed = L * 0.5 * float((delta / L @ (delta / L)))
        predicted = d2 / (2 * L)
        scaling.append({'L': L, 'predicted': predicted, 'observed': observed,
                        'relative_error': abs(predicted-observed)/predicted})
    return {'gaussian_mean_error': mean_err, 'gaussian_covariance_error': cov_err,
            'translation_cost_scaling': scaling}


def moons(n, noise=0.08, seed=0):
    rng = np.random.RandomState(seed)
    t = rng.rand(n) * math.pi
    a = np.stack([np.cos(t), np.sin(t)], 1)
    b = np.stack([1-np.cos(t), 1-np.sin(t)-0.45], 1)
    z = np.concatenate([a, b], 0)[:n]
    z += noise * rng.randn(*z.shape)
    return torch.tensor(z, dtype=torch.float32)


def mmd_rbf(x, y):
    dxx = torch.cdist(x, x).square(); dyy = torch.cdist(y, y).square(); dxy = torch.cdist(x, y).square()
    ans = 0.
    for h in (0.3, 0.6, 1.0, 1.8, 3.0):
        ans = ans + torch.exp(-dxx/(2*h*h)).mean() + torch.exp(-dyy/(2*h*h)).mean() - 2*torch.exp(-dxy/(2*h*h)).mean()
    return ans / 5


def sliced_wasserstein(x, y, projections=32):
    th = torch.arange(projections, device=x.device, dtype=x.dtype) * (math.pi / projections)
    dirs = torch.stack([torch.cos(th), torch.sin(th)], 1)
    a = torch.sort(x @ dirs.T, 0).values
    b = torch.sort(y @ dirs.T, 0).values
    return float((a-b).abs().mean().item())


class PotentialLayer(nn.Module):
    # T(x)=x+grad psi(x), the Euclidean Exp_x(grad psi) adaptation.
    def __init__(self, hidden=24):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
        # Small initial displacement makes the composition near identity.
        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)

    def forward(self, x, create_graph=True):
        x = x.requires_grad_(True)
        psi = self.net(x).sum()
        v = torch.autograd.grad(psi, x, create_graph=create_graph, retain_graph=create_graph)[0]
        return x + v


class ResidualLayer(nn.Module):
    def __init__(self, hidden=24):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2))
        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)

    def forward(self, x): return x + self.net(x)


class OTFlow(nn.Module):
    def __init__(self, potential=True, layers=2):
        super().__init__()
        cls = PotentialLayer if potential else ResidualLayer
        self.layers = nn.ModuleList([cls() for _ in range(layers)])
        self.potential = potential

    def forward(self, x, create_graph=True):
        for layer in self.layers:
            x = layer(x, create_graph) if self.potential else layer(x)
        return x


def train_and_measure(potential, source, target, steps=260):
    model = OTFlow(potential=potential).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=0.008)
    x, ytar = source.to(device), target.to(device)
    for step in range(steps):
        y = model(x, True)
        # OT displacement plus the implementation's distribution discrepancy.
        loss = 0.5*(y-x).square().sum(1).mean() + 12*mmd_rbf(y, ytar)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
    # Potential maps require input derivatives at evaluation time; disable only parameter grads.
    with torch.enable_grad():
        y = model(x.detach(), False)
        mm = float(mmd_rbf(y, ytar).item())
        sw = sliced_wasserstein(y, ytar)
        disp = float(0.5*(y-x.detach()).square().sum(1).mean().item())
    # Jacobian determinant on a small held-out batch tests local invertibility.
    q = source[:48].to(device).requires_grad_(True)
    dets = []
    for i in range(2):
        row = torch.autograd.grad(model(q, True)[:, i].sum(), q, retain_graph=True)[0]
        if i == 0: jac0 = row
        else: jac1 = row
    J = torch.stack([jac0, jac1], 2)
    dets = torch.linalg.det(J).detach()
    return {'mmd': mm, 'sliced_wasserstein': sw, 'displacement_cost': disp,
            'jacobian_positive_fraction': float((dets > 0).float().mean().item()),
            'jacobian_min': float(dets.min().item()), 'jacobian_max': float(dets.max().item())}


def main():
    global device
    sanity = gaussian_sanity()
    source = moons(192, seed=SEED+1)
    gen = torch.Generator().manual_seed(SEED+2)
    target = torch.randn(192, 2, generator=gen) * 1.05
    try:
        baseline = train_and_measure(False, source, target)
        idea = train_and_measure(True, source, target)
    except Exception:
        device = torch.device('cpu')
        baseline = train_and_measure(False, source, target)
        idea = train_and_measure(True, source, target)
    result = {'device': str(device), 'sanity_check': sanity,
              'mini_experiment': {'baseline_unconstrained_residual': baseline, 'ot_gradient_potential': idea},
              'settings': {'seed': SEED, 'samples': 192, 'layers': 2, 'steps': 260, 'mmd_weight': 12}}
    with open('results.json', 'w') as f: json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()