"""One candidate per process; synthetic tensors only, no camera/event calls."""
import json
import pathlib
import sys
import time
import traceback

import torch

root = pathlib.Path('/work')
name = sys.argv[1]
torch.set_num_threads(2)
torch.manual_seed(7)
device = torch.device('cuda:0')
torch.cuda.set_per_process_memory_fraction(.25,device)
torch.cuda.reset_peak_memory_stats()
report = dict(candidate=name,synthetic_input=True,real_video_evaluation=False,
              torch=torch.__version__,cuda=torch.version.cuda,gpu=torch.cuda.get_device_name(0))
start = time.monotonic()
try:
    # Only verified, official downloaded checkpoints are loaded in this isolated container.
    checkpoint = torch.load(root/'weights'/(name+'.pth'),map_location='cpu',weights_only=False)
    if name in ('stgcnpp','posec3d'):
        import mmaction.models
        from mmaction.registry import MODELS
        from mmengine.config import Config
        from mmengine.registry import init_default_scope
        init_default_scope('mmaction')
        configs = {
          'stgcnpp':'skeleton/stgcnpp/stgcnpp_8xb16-joint-u100-80e_ntu60-xsub-keypoint-2d.py',
          'posec3d':'skeleton/posec3d/slowonly_r50_8xb16-u48-240e_ntu60-xsub-keypoint.py'}
        cfg = Config.fromfile('/opt/mmaction2/configs/'+configs[name])
        model = MODELS.build(cfg.model)
        model.load_state_dict(checkpoint['state_dict'],strict=True)
        report['checkpoint_key_count'] = len(checkpoint['state_dict'])
        shape = (1,1,2,100,17,3) if name == 'stgcnpp' else (1,1,17,48,64,64)
        sample = torch.rand(shape,device=device)
        model.eval().to(device)
        with torch.inference_mode():
            output = model(sample,mode='tensor',stage='head')
        assert output.shape == (1,60) and torch.isfinite(output).all()
        report.update(input_shape=list(shape),output_shape=list(output.shape),strict_load=True)
    elif name == 'ava_slowonly':
        from ava_adapter import AvaAdapter,check_roi_contract
        check_roi_contract(device)
        model = AvaAdapter(checkpoint).eval().to(device)
        sample = torch.rand((1,3,4,256,256),device=device)
        boxes = torch.tensor([[0,32,24,180,240],[0,100,48,240,244]],device=device,dtype=torch.float32)
        with torch.inference_mode():
            output = model(sample,boxes)
        assert output.shape == (2,81) and torch.isfinite(output).all()
        report.update(input_shape=list(sample.shape),boxes_shape=list(boxes.shape),
                      output_shape=list(output.shape),strict_load=True,
                      checkpoint_key_count=model.checkpoint_key_count,
                      background_index=0,action_classes=80,
                      compatibility='Official backbone/head + torchvision RoIAlign; not full MMDetection wrapper',
                      roi_coordinate_contract='PASS',mmcv_numerical_equivalence='NOT_TESTED')
    elif name == 'yowov2':
        sys.path.insert(0,'/opt/yowov2')
        from config.yowo_v2_config import yowo_v2_config
        from models.yowo.yowo import YOWO
        model = YOWO(yowo_v2_config['yowo_v2_large'],device,num_classes=80,
                     conf_thresh=.1,nms_thresh=.5,topk=40,trainable=False,multi_hot=True)
        model.load_state_dict(checkpoint['model'],strict=True)
        model.eval().to(device)
        sample = torch.rand((1,3,16,224,224),device=device)
        with torch.inference_mode():
            batch_output = model(sample)
        assert isinstance(batch_output,list) and len(batch_output)==1
        output = batch_output[0]
        import numpy as np
        assert output.ndim == 2 and output.shape[1] == 85 and np.isfinite(output).all()
        report.update(input_shape=list(sample.shape),output_shape=list(output.shape),
                      strict_load=True,checkpoint_key_count=len(checkpoint['model']),
                      output_schema='xyxy[4], objectness[1], action_sigmoid[80]',
                      postprocess=dict(conf_thresh=.1,nms_thresh=.5,topk=40))
    else:
        raise ValueError('Unknown candidate')
    torch.cuda.synchronize()
    report['status'] = 'SMOKE_PASS'
except Exception as error:
    report.update(status='BLOCKED',error_type=type(error).__name__,error=str(error),traceback=traceback.format_exc())
finally:
    report['wall_seconds'] = round(time.monotonic()-start,3)
    report['peak_cuda_allocated_bytes'] = torch.cuda.max_memory_allocated()
    report['peak_cuda_reserved_bytes'] = torch.cuda.max_memory_reserved()
    target = root/'reports'/('smoke_'+name+'.json')
    if target.exists():
        with (root/'reports'/('smoke_history_'+name+'.jsonl')).open('a') as history:
            history.write(json.dumps(json.loads(target.read_text()))+'\n')
    target.write_text(json.dumps(report,indent=2))
    print(json.dumps(report))
raise SystemExit(0 if report['status']=='SMOKE_PASS' else 1)