import json, math, sys
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from plyfile import PlyData, PlyElement
PROJECT_ROOT = Path(__file__).resolve().parents[1]
ANYSPLAT_ROOT = PROJECT_ROOT / 'src' / 'AnySplat'
sys.path.insert(0, str(ANYSPLAT_ROOT))
from src.model.model.anysplat import AnySplat
from src.utils.image import process_image

if len(sys.argv) != 3:
    raise SystemExit('usage: make_fourier_scene_generic.py FRAMES_DIR OUTPUT_DIR')
ROOT = PROJECT_ROOT
FRAMES = sorted(Path(sys.argv[1]).glob('*.jpg'))
OUT = Path(sys.argv[2])
OUT.mkdir(parents=True, exist_ok=True)
(OUT/'images').mkdir(exist_ok=True)
assert len(FRAMES) >= 8 and len(FRAMES) % 2 == 0, f'Need even frame count >=8, got {len(FRAMES)}'
imgs = torch.stack([process_image(str(f)) for f in FRAMES]).unsqueeze(0).cuda()
model = AnySplat.from_pretrained('lhjiang/anysplat').cuda().eval()
for p in model.parameters(): p.requires_grad_(False)
ctx_i = list(range(0,len(FRAMES),2)); tgt_i = list(range(1,len(FRAMES),2))
with torch.inference_mode():
    g_ctx, p_ctx = model.inference((imgs[:,ctx_i]+1)*0.5)
    g_all, p_all = model.inference((imgs+1)*0.5)
src = p_all['extrinsic'][0,ctx_i,:3,3].double().cpu().numpy()
dst = p_ctx['extrinsic'][0,:,:3,3].double().cpu().numpy()
mu_s,mu_d = src.mean(0),dst.mean(0)
xs,xd = src-mu_s,dst-mu_d
cov = xd.T@xs/len(src)
U,S,Vt = np.linalg.svd(cov)
D=np.eye(3); D[2,2]=np.sign(np.linalg.det(U)*np.linalg.det(Vt))
R=U@D@Vt
scale=(S*np.diag(D)).sum()/(xs**2).sum(axis=1).mean()
trans=mu_d-scale*R@mu_s
H,W=imgs.shape[-2:]
flip=np.diag([1.,-1.,-1.,1.])
train_frames=[]; test_frames=[]; fovxs=[]
for i,f in enumerate(FRAMES):
    im=Image.open(f).convert('RGB'); w,h=im.size; side=min(w,h)
    im=im.crop(((w-side)//2,(h-side)//2,(w+side)//2,(h+side)//2)).resize((W,H),Image.Resampling.LANCZOS)
    name=f'{i+1:04d}.png'; im.save(OUT/'images'/name)
    if i in ctx_i:
        j=ctx_i.index(i); c2w=p_ctx['extrinsic'][0,j].float().cpu().numpy()
        intr=p_ctx['intrinsic'][0,j].float().cpu().numpy()
        train_frames.append({'file_path':f'images/{name[:-4]}','transform_matrix':(c2w@flip).tolist(),'index':i})
    else:
        c0=p_all['extrinsic'][0,i].float().cpu().numpy()
        c2w=np.eye(4,dtype=np.float32); c2w[:3,:3]=R@c0[:3,:3]; c2w[:3,3]=scale*R@c0[:3,3]+trans
        intr=p_all['intrinsic'][0,i].float().cpu().numpy()
        test_frames.append({'file_path':f'images/{name[:-4]}','transform_matrix':(c2w@flip).tolist(),'index':i})
    fovxs.append(2*math.atan(1/(2*float(intr[0,0]))))
fovx=float(np.median(fovxs))
(OUT/'transforms_train.json').write_text(json.dumps({'camera_angle_x':fovx,'frames':train_frames},indent=2))
(OUT/'transforms_test.json').write_text(json.dumps({'camera_angle_x':fovx,'frames':test_frames},indent=2))
means=g_ctx.means[0].float(); scales=g_ctx.scales[0].float(); opacity=g_ctx.opacities[0].float().reshape(-1); harm=g_ctx.harmonics[0].float()
score=opacity*scales.max(dim=1).values.clamp_min(1e-7); n=min(120_000,len(score)); idx=torch.topk(score,n,sorted=False).indices
xyz=means[idx].cpu().numpy().astype(np.float32); rgb=(0.5+0.28209479177387814*harm[idx,:,0]).clamp(0,1).cpu().numpy()
verts=np.empty(n,dtype=[('x','f4'),('y','f4'),('z','f4'),('nx','f4'),('ny','f4'),('nz','f4'),('red','u1'),('green','u1'),('blue','u1')])
verts['x'],verts['y'],verts['z']=xyz.T; verts['nx'],verts['ny'],verts['nz']=0,0,0
verts['red'],verts['green'],verts['blue']=(rgb*255).astype(np.uint8).T
PlyData([PlyElement.describe(verts,'vertex')]).write(str(OUT/'points3d.ply'))
meta={'source':'AnySplat context-only geometry initialization','total_video_frames':len(FRAMES),'context_views':len(train_frames),'heldout_views':len(test_frames),'seed_primitives':n,'image_hw':[H,W],'camera_fov_x':fovx,'pose_alignment_scale':float(scale),'heldout_geometry_used_for_initialization':False}
(OUT/'seed_metadata.json').write_text(json.dumps(meta,indent=2)); print(json.dumps(meta,indent=2),flush=True)