import json, math, random
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED = 573
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)


def wrap(x):
    return (x + 0.5) % 1.0 - 0.5


def graph_cochains(theta, k=12):
    # theta: N x m angles in [0,1); torus distance uses the first two physical angles
    N, m = theta.shape
    d = np.sqrt(np.sum(wrap(theta[:, None, :2] - theta[None, :, :2])**2, axis=2))
    edges = []
    seen = set()
    for i in range(N):
        for j in np.argsort(d[i])[1:k+1]:
            a,b = sorted((i,int(j)))
            if a != b and (a,b) not in seen:
                seen.add((a,b)); edges.append((a,b))
    A = np.array([wrap(theta[j]-theta[i]) for i,j in edges])
    w = np.ones(len(edges))/len(edges)
    G = A.T @ (w[:,None]*A)
    return G, A, edges


def select_matroid(G, Q, tol=1e-8):
    # Q is m x r; v_j are rows of Q, and M=Q^T G Q.
    M = Q.T @ G @ Q
    V = Q.T                         # r x m, columns are candidate vectors
    energies = np.einsum('ir,rs,is->i', Q, M, Q)
    order = np.argsort(energies)
    chosen=[]; rank=0
    for j in order:
        trial = V[:, chosen+[int(j)]]
        s = np.linalg.svd(trial, compute_uv=False)
        newrank = int(np.sum(s > tol * max(1.0, s[0])))
        if newrank > rank:
            chosen.append(int(j)); rank=newrank
        if rank == Q.shape[1]: break
    return chosen, energies, M


def residual(G, Q, chosen):
    # Persistent-space residual as stated, with V made from selected projected vectors.
    M = Q.T @ G @ Q
    V = Q.T[:, chosen]
    P = V @ np.linalg.pinv(V.T @ M @ V) @ V.T @ M if chosen else np.zeros_like(M)
    R = np.eye(M.shape[0]) - P
    return np.linalg.norm(R, 'fro'), R

class AE(nn.Module):
    def __init__(self):
        super().__init__(); self.enc=nn.Sequential(nn.Linear(4,32),nn.Tanh(),nn.Linear(32,2)); self.dec=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,4))
    def forward(self,x): return self.dec(self.enc(x))
class CircDecoder(nn.Module):
    def __init__(self,d):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2*d,32),nn.Tanh(),nn.Linear(32,4))
    def forward(self,z): return self.net(z)

def train_models(x, selected, theta):
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        xt=torch.tensor(x,dtype=torch.float32,device=device)
        # Circular selected features are fixed, as prescribed by the dictionary.
        z=np.concatenate([np.cos(2*np.pi*theta[:,selected]), np.sin(2*np.pi*theta[:,selected])],axis=1)
        zt=torch.tensor(z,dtype=torch.float32,device=device)
        base=AE().to(device); circ=CircDecoder(len(selected)).to(device)
        ob=torch.optim.Adam(base.parameters(),lr=2e-3); oc=torch.optim.Adam(circ.parameters(),lr=2e-3)
        for _ in range(1000):
            ob.zero_grad(); loss=((base(xt)-xt)**2).mean(); loss.backward(); ob.step()
            oc.zero_grad(); loss2=((circ(zt)-xt)**2).mean(); loss2.backward(); oc.step()
        with torch.no_grad():
            br=float(((base(xt)-xt)**2).mean().cpu()); cr=float(((circ(zt)-xt)**2).mean().cpu())
            # decode angles from the circular pairs; scalar AE has no canonical angle decoder
            decoded=np.mod(np.arctan2(z[:,len(selected):],z[:,:len(selected)])/(2*np.pi),1.0)
            # compare each selected physical candidate to its exact dictionary angle
            angle_err=float(np.mean(np.abs(wrap(decoded-theta[:,selected]))))
        return br,cr,angle_err,device
    except Exception as e:
        print('device fallback:',repr(e))
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        # CPU retry
        old=torch.cuda.is_available
        torch.cuda.is_available=lambda: False
        out=train_models(x,selected,theta)
        torch.cuda.is_available=lambda: old
        return out

def main():
    N=900
    t=np.random.rand(N,2)
    # physical torus embedded in R4; candidate angle dictionary includes mixtures and distractors
    x=np.concatenate([np.cos(2*np.pi*t[:,0,None]),np.sin(2*np.pi*t[:,0,None]),np.cos(2*np.pi*t[:,1,None]),np.sin(2*np.pi*t[:,1,None])],axis=1)
    theta=np.column_stack([t[:,0],t[:,1],(t[:,0]+t[:,1])%1,(t[:,0]-t[:,1])%1,np.random.rand(N)])
    G,A,edges=graph_cochains(theta)
    # Q explicitly represents detected persistent H1 directions: candidates 0 and 1 span them.
    Q=np.zeros((5,2)); Q[0,0]=1; Q[1,1]=1; Q[2]=[.5,.5]; Q[3]=[.5,-.5]; Q[4]=[.03,.02]
    selected,E,M=select_matroid(G,Q)
    rfull,_=residual(G,Q,[0,1]); romit,_=residual(G,Q,[0])
    br,cr,ae,device=train_models(x,selected,theta)
    out={'seed':SEED,'edges':len(edges),'gram_trace':float(np.trace(G)),'gram_eigenvalues':np.linalg.eigvalsh(G).tolist(),'energies':E.tolist(),'selected':selected,'residual_full':rfull,'residual_omit_one':romit,'baseline_reconstruction_mse':br,'idea_reconstruction_mse':cr,'idea_selected_angle_mean_wrapped_error':ae,'device':device}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()