"""Run official image-only models and save their native dense outputs.

Run inside the GPU holder using envs/mapanything/bin/python run.py MODEL.
The separate site builder applies identical display filtering to each model.
"""
import argparse
import json
import os
from pathlib import Path
import subprocess
import time

os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')
os.environ.setdefault('OMP_NUM_THREADS', '8')
import numpy as np
import torch

ROOT = Path(__file__).resolve().parent
R2S = ROOT.parents[2]
CHECKPOINTS = R2S / 'sceneforge/vendor/checkpoints/geometry_compare'

def array(x):
    return x.detach().float().cpu().numpy() if isinstance(x, torch.Tensor) else np.asarray(x)

def main():
    global ROOT
    ap = argparse.ArgumentParser()
    ap.add_argument('model', choices=['mapanything', 'da3', 'vggt-omega'])
    ap.add_argument('--work', type=Path, default=ROOT, help='Experiment directory containing selection.json')
    a = ap.parse_args()
    ROOT = a.work.resolve()
    torch.set_num_threads(8)
    torch.manual_seed(0)
    np.random.seed(0)
    selection = json.loads((ROOT / 'selection.json').read_text())
    paths = [str(ROOT.parent / 'images' / f['name']) for f in selection['frames']]
    out = ROOT / 'out' / a.model
    out.mkdir(parents=True, exist_ok=True)
    t0 = time.perf_counter()
    print('MODEL', a.model, 'VIEWS', len(paths), 'GPU', torch.cuda.get_device_name(), flush=True)
    if a.model == 'mapanything':
        from mapanything.models import MapAnything
        from mapanything.utils.image import load_images
        model = MapAnything.from_pretrained(str(CHECKPOINTS / 'mapanything')).cuda().eval()
        views = load_images(paths)
        torch.cuda.synchronize()
        loaded = time.perf_counter()
        torch.cuda.reset_peak_memory_stats()
        with torch.inference_mode():
            pred = model.infer(views, memory_efficient_inference=True, minibatch_size=1,
                               use_amp=True, amp_dtype='bf16', apply_mask=True,
                               mask_edges=False, apply_confidence_mask=False)
        torch.cuda.synchronize()
        inferred = time.perf_counter()
        peak = torch.cuda.max_memory_allocated() / 1e9
        depth = np.stack([array(p['depth_z'])[0, ..., 0] for p in pred])
        conf = np.stack([array(p['conf'])[0] for p in pred])
        c2w = np.stack([array(p['camera_poses'])[0] for p in pred])
        K = np.stack([array(p['intrinsics'])[0] for p in pred])
        rgb = np.stack([(array(p['img_no_norm'])[0] * 255).clip(0, 255).astype('uint8') for p in pred])
        mask = np.stack([array(p['mask'])[0, ..., 0].astype(bool) for p in pred])
        # Preserve the native point maps: MapAnything also predicts non-pinhole rays.
        xyz = np.stack([array(p['pts3d'])[0] for p in pred])
        np.save(out / 'points.npy', xyz)
        model_name = 'facebook/map-anything'
        settings = dict(memory_efficient_inference=True, minibatch_size=1, mask_edges=False,
                        normalization='dinov2', resize='official fixed_mapping 518')
        scale = 'Model-predicted metric scale; unverified against physical measurements.'
    elif a.model == 'da3':
        from depth_anything_3.api import DepthAnything3
        model = DepthAnything3.from_pretrained(str(CHECKPOINTS / 'da3')).cuda().eval()
        torch.cuda.synchronize()
        loaded = time.perf_counter()
        torch.cuda.reset_peak_memory_stats()
        pred = model.inference(paths, process_res=504, use_ray_pose=True,
                               ref_view_strategy='saddle_balanced')
        torch.cuda.synchronize()
        inferred = time.perf_counter()
        peak = torch.cuda.max_memory_allocated() / 1e9
        depth, conf, K, rgb = pred.depth, pred.conf, pred.intrinsics, pred.processed_images
        ext = np.broadcast_to(np.eye(4), (len(paths), 4, 4)).copy()
        ext[:, :3, :] = pred.extrinsics[:, :3, :]
        c2w = np.linalg.inv(ext)
        mask = np.isfinite(depth) & (depth > 0)
        model_name = 'depth-anything/DA3NESTED-GIANT-LARGE-1.1'
        settings = dict(process_res=504, process_res_method='upper_bound_resize',
                        use_ray_pose=True, ref_view_strategy='saddle_balanced')
        scale = 'Nested model-predicted metric scale; unverified against physical measurements.'
    else:
        from vggt_omega.models import VGGTOmega
        from vggt_omega.utils.load_fn import load_and_preprocess_images
        from vggt_omega.utils.pose_enc import encoding_to_camera
        model = VGGTOmega().cuda().eval()
        state = torch.load(CHECKPOINTS / 'vggt-omega/vggt_omega_1b_512.pt', map_location='cpu', weights_only=True)
        print(model.load_state_dict(state.get('model', state.get('state_dict', state)), strict=True), flush=True)
        del state
        images = load_and_preprocess_images(paths, image_resolution=512).cuda()
        torch.cuda.synchronize()
        loaded = time.perf_counter()
        torch.cuda.reset_peak_memory_stats()
        with torch.inference_mode():
            pred = model(images)
            extrinsics, intrinsics = encoding_to_camera(pred['pose_enc'], pred['images'].shape[-2:])
        torch.cuda.synchronize()
        inferred = time.perf_counter()
        peak = torch.cuda.max_memory_allocated() / 1e9
        depth = array(pred['depth'])[0, ..., 0]
        conf = array(pred['depth_conf'])[0]
        K = array(intrinsics)[0]
        ext = np.broadcast_to(np.eye(4), (len(paths), 4, 4)).copy()
        ext[:, :3, :] = array(extrinsics)[0, :, :3, :]
        c2w = np.linalg.inv(ext)
        rgb = (array(images).transpose(0, 2, 3, 1) * 255).clip(0, 255).astype('uint8')
        mask = np.isfinite(depth) & (depth > 0)
        model_name = 'facebook/VGGT-Omega / vggt_omega_1b_512.pt'
        settings = dict(image_resolution=512, resize='official balanced', autocast='bf16')
        scale = 'Relative scale; viewer alignment supplies an approximate display scale.'
    print('INFERRED', depth.shape, 'seconds', inferred-loaded, 'peak GB', peak, flush=True)
    for n, x in dict(depth=depth, confidence=conf, c2w=c2w, intrinsics=K, rgb=rgb, mask=mask).items():
        np.save(out / f'{n}.npy', x)
    vendor = {'mapanything':'map-anything', 'da3':'depth-anything-3', 'vggt-omega':'vggt-omega'}[a.model]
    commit = subprocess.check_output(['git','-C',str(R2S / 'sceneforge/vendor' / vendor),'rev-parse','HEAD'],text=True).strip()
    result = dict(model=a.model, checkpoint=model_name, commit=commit, n_views=len(paths),
                  image_height=int(depth.shape[1]), image_width=int(depth.shape[2]),
                  input='RGB images only', settings=settings, scale=scale,
                  inference_seconds=round(inferred-loaded,3), load_preprocess_seconds=round(loaded-t0,3),
                  peak_gpu_allocated_gb=round(peak,3), gpu=torch.cuda.get_device_name(),
                  torch=torch.__version__, total_seconds=round(time.perf_counter()-t0,3))
    (out / 'result.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2), flush=True)

if __name__ == '__main__':
    main()