"""Single-image inference using the unmodified official DoUNet."""
import argparse
import json
import sys
from pathlib import Path


def main():
    p = argparse.ArgumentParser()
    source = p.add_mutually_exclusive_group(required=True)
    source.add_argument('--repo-id')
    source.add_argument('--ckpt', type=Path)
    p.add_argument('--revision', default='main', help='Prefer the uploaded commit SHA')
    p.add_argument('--image', required=True, type=Path)
    p.add_argument('--device', default='auto')
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--disable-amp', action='store_true')
    p.add_argument('--trust-checkpoint', action='store_true', help='Allow pickle loading for a trusted official checkpoint only')
    a = p.parse_args()
    import torch
    from PIL import Image
    from huggingface_hub import hf_hub_download
    sys.path.insert(0, str(Path(__file__).resolve().parent / 'upstream'))
    from models.dou import DoUNet, load_clip_to_cpu
    from utils.options import TrainCfg
    from dataset.dataset import translate_duplicate
    from torchvision import transforms
    torch.manual_seed(a.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(a.seed)
    device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu') if a.device == 'auto' else a.device)
    path = a.ckpt or Path(hf_hub_download(a.repo_id, 'best.pth', revision=a.revision))
    cfg = TrainCfg(model_name='ViT-L/14', classnames=('real', 'generated'), learn_logit_scale=False, amp=not a.disable_amp)
    backbone = load_clip_to_cpu(cfg).eval()
    for param in backbone.parameters():
        param.requires_grad_(False)
    model = DoUNet(cfg, backbone).to(device)
    ckpt = torch.load(path, map_location='cpu', weights_only=False)
    if not isinstance(ckpt, dict) or not isinstance(ckpt.get('model'), dict):
        raise ValueError("Expected a checkpoint dictionary containing ckpt['model']")
    # Upstream uses strict=False silently. Reject missing/unexpected tensors here
    # to avoid presenting an incompletely loaded detector as a valid result.
    report = model.load_state_dict(ckpt['model'], strict=False)
    if report.missing_keys or report.unexpected_keys:
        raise RuntimeError(f'Checkpoint mismatch: {report}. Resolve against the official release; do not ignore it.')
    model.eval()
    # Obtain the exact official evaluation transform without scanning a dataset.
    transform = transforms.Compose([
        transforms.Lambda(
            lambda img: translate_duplicate(img, 224)),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            [0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711]),
    ])
    with Image.open(a.image) as img:
        x = transform(img.convert('RGB')).unsqueeze(0).to(device)
    with torch.inference_mode(), torch.amp.autocast('cuda', enabled=cfg.amp and device.type == 'cuda'):
        probs = model(x).softmax(dim=-1)[0].float().cpu()
    print(json.dumps({'label': ['real', 'generated'][probs.argmax().item()],
                      'real_score': probs[0].item(), 'generated_score': probs[1].item(),
                      'seed': a.seed, 'note': 'Softmax scores are not calibrated probabilities of authenticity.'}))


if __name__ == '__main__':
    main()