import json, math, time
from itertools import combinations
import numpy as np
import torch
from torch import nn

SEED = 7
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if DEVICE == "cuda":
        torch.cuda.set_device(0)
        torch.zeros(1, device=DEVICE)
except Exception:
    DEVICE = "cpu"


def barycentric_matrix(n):
    d = n - 1
    H = np.zeros((n, n), dtype=np.float64)
    for i in range(n):
        for j in range(n):
            for r in range(i + 1):
                H[i, j] += r**j * (r + 1)**(d-j) * (-1)**(i-r) * math.comb(d+1, i-r)
    return H


def all_minors(A, max_order=None):
    n, m = A.shape
    q = min(n, m) if max_order is None else min(max_order, n, m)
    out = []
    for k in range(1, q + 1):
        for rows in combinations(range(n), k):
            for cols in combinations(range(m), k):
                out.append((k, rows, cols, float(np.linalg.det(A[np.ix_(rows, cols)]))))
    return out


class PositiveBidiagonalMixer(nn.Module):
    def __init__(self, n, depth=2):
        super().__init__()
        self.n, self.depth = n, depth
        self.ell = nn.Parameter(torch.full((depth, n-1), -4.0))
        self.u = nn.Parameter(torch.full((depth, n-1), -4.0))
        self.s = nn.Parameter(torch.zeros(n))

    def forward(self, x):
        y = x * torch.exp(self.s)
        for k in range(self.depth):
            z = torch.zeros_like(y)
            z[..., 1:] = y[..., :-1] * torch.nn.functional.softplus(self.ell[k])
            y = y + z
        for k in range(self.depth - 1, -1, -1):
            z = torch.zeros_like(y)
            z[..., :-1] = y[..., 1:] * torch.nn.functional.softplus(self.u[k])
            y = y + z
        return y

    def matrix(self):
        return self(torch.eye(self.n, device=self.s.device)).T


class DenseMap(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(n, n) / math.sqrt(n))
    def forward(self, x):
        return x @ self.weight.T


class LowRankMap(nn.Module):
    def __init__(self, n, rank):
        super().__init__()
        self.a = nn.Parameter(torch.randn(n, rank) / math.sqrt(n))
        self.b = nn.Parameter(torch.randn(rank, n) / math.sqrt(rank))
    def forward(self, x):
        return x @ (self.a @ self.b).T


def fit(model, xtr, ytr, xva, yva, steps=500, lr=0.03):
    model.to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    t0 = time.perf_counter()
    for _ in range(steps):
        opt.zero_grad(set_to_none=True)
        loss = ((model(xtr) - ytr)**2).mean()
        loss.backward()
        opt.step()
    with torch.no_grad():
        train = ((model(xtr) - ytr)**2).mean().item()
        val = ((model(xva) - yva)**2).mean().item()
    return train, val, time.perf_counter() - t0, sum(p.numel() for p in model.parameters())


def main():
    # Core formula and total-positivity sanity check.
    H = barycentric_matrix(5)
    hm = all_minors(H)
    negative = [v for v in hm if v[3] < -1e-8]
    # The practical parameterization has strictly positive entries, but not every
    # product of arbitrary positive bidiagonals is TP; test the claimed network map.
    torch.manual_seed(SEED)
    pos = PositiveBidiagonalMixer(5, depth=4).to(DEVICE)
    P = pos.matrix().detach().cpu().numpy()
    pm = all_minors(P)
    pnegative = [v for v in pm if v[3] < -1e-8]

    # Approximation/learning test: target is a dense linear map.
    torch.manual_seed(SEED)
    n, ns = 12, 1024
    x = torch.randn(ns, n, device=DEVICE)
    target = torch.randn(n, n, device=DEVICE) / math.sqrt(n)
    y = x @ target.T
    xtr, ytr, xva, yva = x[:768], y[:768], x[768:], y[768:]
    results = {}
    for name, model in [
        ("dense", DenseMap(n)),
        ("positive_K2", PositiveBidiagonalMixer(n, 2)),
        ("positive_K6", PositiveBidiagonalMixer(n, 6)),
        ("lowrank_r4", LowRankMap(n, 4)),
    ]:
        torch.manual_seed(SEED)
        tr, va, secs, params = fit(model, xtr, ytr, xva, yva)
        results[name] = {"train_mse": tr, "val_mse": va, "seconds": secs, "parameters": params}

    report = {
        "device": DEVICE,
        "barycentric_H_5": H.tolist(),
        "barycentric_minors": {"count": len(hm), "negative": len(negative), "minimum": min(v[3] for v in hm)},
        "positive_product_minors": {"count": len(pm), "negative": len(pnegative), "minimum": min(v[3] for v in pm)},
        "fit": results,
    }
    with open("results.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()