"""Multi-arch eval on HER2Match val: cleanfid + pytorch-fid + brown-frac.

Supports C7 (UNI-2H encoder), C3 (WaveletJointAEv3), A3 (SkipFreeAE3Stage).
Env-var configured for CI-friendly multi-run:
  ARCH=C7|C3|A3       — which model class to instantiate
  CKPT_LOCAL          — path to best.pt (state_dict must match ARCH)
  CKPT_TAG            — subfolder tag for outputs (e.g. 'c7_best', 'c3_dab_best')
  OUT_ROOT            — where to put generated tiles + results json
  HF_EXP_SUB          — HF sub-path under `augustander/her2match-src-ds-arch-sweep/` for upload
"""
import os, sys, json, time, glob, subprocess
from pathlib import Path
import numpy as np
import torch, torch.nn.functional as F
from PIL import Image

ARCH          = os.environ.get('ARCH', 'C7')
CKPT_LOCAL    = os.environ['CKPT_LOCAL']
CKPT_TAG      = os.environ.get('CKPT_TAG', 'best')
OUT_ROOT      = Path(os.environ.get('OUT_ROOT', f'/home/ubuntu/eval_{ARCH}_{CKPT_TAG}'))
HF_EXP_SUB    = os.environ.get('HF_EXP_SUB', 'multi_arch_eval')
HF_MODEL_REPO = 'augustander/her2match-src-ds-arch-sweep'
DATA_ROOT     = '/home/ubuntu/her2match_full/her2match_full'
GEN_DIR       = OUT_ROOT / 'generated_ihc'
GEN_DIR.mkdir(parents=True, exist_ok=True)
device = 'cuda'

sys.path.insert(0, '/home/ubuntu/work')
from src_ds_her2_thunder import (
    UniDwtC7, WaveletJointAEv3, SkipFreeAE3Stage,
    UNI2HSpatialLoss, StainDeconvDAB,
)

# ---- Load model per ARCH ----
def build_model(arch: str, ckpt: dict, uni_loss=None):
    if arch == 'C7':
        assert uni_loss is not None, 'C7 needs shared UNI-2H'
        m = UniDwtC7(
            uni_module=uni_loss.uni, wavelet='haar', taps=(13, 19),
            dec_ch=(512, 384, 256, 128, 64), skip_ch=(384, 256, 128),
            blocks_per_stage=(4, 3, 3, 2, 2), uni_target_size=224,
        )
    elif arch == 'C3':
        m = WaveletJointAEv3(base_ch=48, wavelet='haar')
    elif arch == 'A3':
        m = SkipFreeAE3Stage(base_ch=48, bottleneck_ch=256)
    else:
        raise ValueError(f'unknown ARCH={arch}')
    return m

# ---- ALWAYS load shared UNI (needed for C7 arch OR just for future perceptual use) ----
if ARCH == 'C7':
    print(f'[{ARCH}/{CKPT_TAG}] loading shared UNI-2H...', flush=True)
    uni_loss = UNI2HSpatialLoss(target_size=224).to(device).eval()
    for p in uni_loss.uni.parameters(): p.requires_grad_(False)
else:
    uni_loss = None

# ---- Build + load ----
G = build_model(ARCH, None, uni_loss=uni_loss).to(device).eval()
ckpt = torch.load(CKPT_LOCAL, map_location='cpu', weights_only=False)
missing, unexpected = G.load_state_dict(ckpt['model'], strict=False)
n_missing_uni = len([k for k in missing if k.startswith('uni.')])
n_missing_other = [k for k in missing if not k.startswith('uni.')]
print(f'[{ARCH}/{CKPT_TAG}] loaded ckpt step={ckpt.get("step","?")} '
      f'val_lpips={ckpt.get("val_lpips",-1):.4f}  '
      f'(missing_uni={n_missing_uni}, missing_other={n_missing_other}, unexpected={len(unexpected)})',
      flush=True)
for p in G.parameters(): p.requires_grad_(False)
stain_dab = StainDeconvDAB().to(device).eval()

# ---- Generate 3582 tiles ----
he_val = sorted(glob.glob(f'{DATA_ROOT}/HE/val/*.jpg'))
ihc_val_dir = f'{DATA_ROOT}/IHC/val'
BATCH = 16 if ARCH == 'C7' else 32
print(f'[{ARCH}/{CKPT_TAG}] generating {len(he_val)} tiles (batch={BATCH})...', flush=True)
t0 = time.time()
for i in range(0, len(he_val), BATCH):
    chunk = he_val[i:i+BATCH]
    imgs = []
    for f in chunk:
        im = Image.open(f).convert('RGB').resize((256, 256), Image.BICUBIC)
        x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.0).permute(2, 0, 1)
        imgs.append(x * 2 - 1)
    x = torch.stack(imgs).to(device)
    with torch.no_grad():
        y = G(x).clamp(-1, 1)
    y_np = ((y.cpu() + 1) * 127.5).clamp(0, 255).byte().permute(0, 2, 3, 1).numpy()
    for k, f in enumerate(chunk):
        Image.fromarray(y_np[k]).save(GEN_DIR / Path(f).name, quality=92)
    if (i // BATCH) % 20 == 0:
        done = i + len(chunk); rate = done / max(1e-6, time.time() - t0)
        eta = (len(he_val) - done) / max(1e-6, rate)
        print(f'  gen {done}/{len(he_val)}  ({rate:.1f} img/s  eta {eta:.0f}s)', flush=True)
print(f'[{ARCH}/{CKPT_TAG}] gen done in {time.time()-t0:.0f}s', flush=True)

# ---- cleanfid ----
print(f'[{ARCH}/{CKPT_TAG}] cleanfid...', flush=True)
try:
    from cleanfid import fid as cfid
    cleanfid_val = cfid.compute_fid(
        str(GEN_DIR), ihc_val_dir,
        device='cuda', mode='clean', batch_size=64, num_workers=4)
    print(f'[{ARCH}/{CKPT_TAG}] CLEANFID = {cleanfid_val:.3f}', flush=True)
except Exception as e:
    print(f'[{ARCH}/{CKPT_TAG}] cleanfid err: {e}', flush=True); cleanfid_val = None

# ---- pytorch-fid ----
print(f'[{ARCH}/{CKPT_TAG}] pytorch-fid...', flush=True)
try:
    r = subprocess.run(
        ['python', '-m', 'pytorch_fid', str(GEN_DIR), ihc_val_dir,
         '--device', 'cuda', '--batch-size', '50'],
        capture_output=True, text=True, timeout=1800)
    pytorch_fid_val = None
    for line in r.stdout.splitlines():
        if line.strip().startswith('FID:'):
            pytorch_fid_val = float(line.split()[1])
            break
    print(f'[{ARCH}/{CKPT_TAG}] PYTORCH_FID = {pytorch_fid_val}', flush=True)
    if pytorch_fid_val is None:
        print(f'[{ARCH}/{CKPT_TAG}] pytorch-fid stderr tail: {r.stderr[-500:]}', flush=True)
except Exception as e:
    print(f'[{ARCH}/{CKPT_TAG}] pytorch-fid err: {e}', flush=True); pytorch_fid_val = None

# ---- brown-frac ----
print(f'[{ARCH}/{CKPT_TAG}] brown-frac...', flush=True)
def dab_stats(paths, batch=64):
    means, fracs = [], []
    for i in range(0, len(paths), batch):
        chunk = paths[i:i+batch]
        imgs = []
        for f in chunk:
            im = Image.open(f).convert('RGB').resize((256, 256), Image.BICUBIC)
            x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.0).permute(2, 0, 1)
            imgs.append(x * 2 - 1)
        x = torch.stack(imgs).to(device)
        with torch.no_grad():
            dab = stain_dab(x)
        means.append(dab.mean(dim=(1,2)).cpu().numpy())
        fracs.append((dab > 0.3).float().mean(dim=(1,2)).cpu().numpy())
    return np.concatenate(means), np.concatenate(fracs)

gen_files = sorted(GEN_DIR.glob('*.jpg'))
tgt_files = sorted(Path(ihc_val_dir).glob('*.jpg'))
common = sorted({p.name for p in gen_files} & {p.name for p in tgt_files})
gen_files = [GEN_DIR / n for n in common]
tgt_files = [Path(ihc_val_dir) / n for n in common]
gm, gf = dab_stats(gen_files)
tm, tf = dab_stats(tgt_files)
her2p_mask = tf > 0.05
her2p_n = int(her2p_mask.sum())

results = {
    'arch': ARCH, 'ckpt_tag': CKPT_TAG, 'ckpt_local': CKPT_LOCAL,
    'ckpt_step': int(ckpt.get('step', -1)),
    'ckpt_val_lpips': float(ckpt.get('val_lpips', -1)),
    'cleanfid': cleanfid_val, 'pytorch_fid': pytorch_fid_val,
    'n_tiles': len(common), 'n_her2p_tiles': her2p_n,
    'overall_gen_bf': float(gf.mean()), 'overall_tgt_bf': float(tf.mean()),
    'overall_gen_dab': float(gm.mean()), 'overall_tgt_dab': float(tm.mean()),
    'her2p_gen_bf': float(gf[her2p_mask].mean()) if her2p_n else None,
    'her2p_tgt_bf': float(tf[her2p_mask].mean()) if her2p_n else None,
    'her2p_under_predict_pct': float((gf[her2p_mask] < 0.5 * tf[her2p_mask]).mean()) if her2p_n else None,
}
print(f'[{ARCH}/{CKPT_TAG}] RESULTS:', json.dumps(results, indent=2), flush=True)
out_json = OUT_ROOT / f'eval_{ARCH}_{CKPT_TAG}.json'
out_json.write_text(json.dumps(results, indent=2))

if 'HF_TOKEN' not in os.environ and os.path.exists('/tmp/hf_token'):
    os.environ['HF_TOKEN'] = open('/tmp/hf_token').read().strip()
from huggingface_hub import HfApi
HfApi().upload_file(
    path_or_fileobj=str(out_json),
    path_in_repo=f'{HF_EXP_SUB}/eval_{ARCH}_{CKPT_TAG}.json',
    repo_id=HF_MODEL_REPO, repo_type='model',
    commit_message=f'multi-eval: {ARCH} {CKPT_TAG} cleanfid+pytorch-fid+brown-frac')
print(f'=== EVAL_DONE {ARCH}/{CKPT_TAG} ===', flush=True)