import json, math, sys, time
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from plyfile import PlyData, PlyElement
sys.path.insert(0, '/home/dev/polymorphic/vendor/AnySplat')
from src.model.model.anysplat import AnySplat
from src.utils.image import process_image

ROOT = Path('/home/dev/polymorphic')
OUT = ROOT / 'fourier_scene_apt'
FRAMES = sorted((ROOT / 'exp_apt' / 'frames').glob('*.jpg'))
OUT.mkdir(parents=True, exist_ok=True)
(OUT / 'images').mkdir(exist_ok=True)
assert len(FRAMES) == 24, len(FRAMES)
device = 'cuda'
imgs = torch.stack([process_image(str(f)) for f in FRAMES]).unsqueeze(0).to(device)
model = AnySplat.from_pretrained('lhjiang/anysplat').to(device).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)

# Align the all-frame predicted target cameras into the context-only reconstruction frame,
# matching the held-out AnySplat evaluation protocol.
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:]
fovxs = []
def blender_matrix(c2w):
    # FourierSplatting's Blender loader flips camera Y/Z. Pre-flip so that
    # loading recovers the original AnySplat camera-to-world matrix.
    return (c2w @ np.diag([1., -1., -1., 1.])).tolist()

train_frames, test_frames = [], []
for i, f in enumerate(FRAMES):
    im = Image.open(f).convert('RGB')
    w, h = im.size
    side = min(w, h)
    left, top = (w-side)//2, (h-side)//2
    im = im.crop((left, top, left+side, top+side)).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()
        rec = {'file_path':f'images/{name[:-4]}','transform_matrix':blender_matrix(c2w),'index':i}
        train_frames.append(rec)
    else:
        c2w0 = p_all['extrinsic'][0,i].float().cpu().numpy()
        c2w = np.eye(4, dtype=np.float32)
        c2w[:3,:3] = R @ c2w0[:3,:3]
        c2w[:3,3] = scale * R @ c2w0[:3,3] + trans
        intr = p_all['intrinsic'][0,i].float().cpu().numpy()
        rec = {'file_path':f'images/{name[:-4]}','transform_matrix':blender_matrix(c2w),'index':i}
        test_frames.append(rec)
    fovxs.append(2.0 * math.atan(1.0 / (2.0 * 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))

# Seed from context-only AnySplat primitives; no target-frame geometry enters initialization.
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),'pose_alignment_translation_norm':float(np.linalg.norm(trans)),'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)