import argparse
import csv
import json
from pathlib import Path


def parser(mode):
    p = argparse.ArgumentParser(description=f'Metalorian {mode} sampling')
    p.add_argument('--target_idx', type=int, required=True, help='Metal index; Cu=3, Cd=1, Zn=13, Ni=8')
    p.add_argument('--seq_len', type=int, default=30)
    p.add_argument('--batch', type=int, default=1)
    p.add_argument('--desired', type=int, default=1)
    p.add_argument('--max_rounds', type=int, default=150)
    p.add_argument('--grad_steps', type=int, default=4)
    p.add_argument('--guidance_scale', type=float, default=6.0)
    p.add_argument('--penalty', type=float, default=0.1)
    p.add_argument('--tau', type=float, default=1.0)
    p.add_argument('--label_clamp', type=float, default=0.8)
    p.add_argument('--best_decode', type=int, default=16, help='Candidate decodes; 0 uses argmax')
    p.add_argument('--decode_temp', type=float, default=1.0)
    p.add_argument('--device', default='cuda:0')
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--output', type=Path, required=True)
    if mode == 'prototype':
        p.add_argument('--proto', required=True)
        p.add_argument('--t_start', type=int, default=15)
    p.set_defaults(guide_on='x0', no_accum=False, track_best=False, accept_all=False,
                   amp=False, out_prefix=None, cys_logit_bias=0.0, cys_target=0.0,
                   cys_weight=0.0, min_cys=0, max_aa_frac=1.0, relax_after=10**9,
                   init_noise=0.3)
    return p


def collect(args, sampler_class, mode):
    from .sampling import FLAGS, METAL_LABELS
    if not 0 <= args.target_idx < 15 or not 1 <= args.seq_len <= 1022:
        raise ValueError('target_idx must be 0–14; seq_len must be 1–1022')
    if min(args.batch, args.desired, args.max_rounds) < 1:
        raise ValueError('batch, desired and max_rounds must be positive')
    if mode == 'prototype':
        if not 0 < args.t_start < 50:
            raise ValueError('t_start must be 1–49')
        if not args.proto or set(args.proto)-set('ACDEFGHIKLMNPQRSTVWY'):
            raise ValueError('proto must contain standard uppercase amino acids')
    import random
    import numpy as np
    import torch
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    sampler = sampler_class(args, device=args.device)
    rows, seen, candidates = [], set(), 0
    for round_index in range(1, args.max_rounds+1):
        sequences, _ = sampler.run_round()
        sequences = [s for s in sequences if s]
        candidates += len(sequences)
        hits, _ = sampler.accept(sequences, None) if sequences else ([], None)
        for hit in hits:
            if hit['sequence'] in seen:
                continue
            seen.add(hit['sequence'])
            rows.append(dict(metal=METAL_LABELS[args.target_idx], **hit))
            if len(rows) == args.desired:
                break
        print(f'round={round_index} candidates={candidates} unique_accepted={len(rows)}/{args.desired}', flush=True)
        if len(rows) >= args.desired:
            break
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open('w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['metal','sequence','score','n_cys','n_his'])
        writer.writeheader()
        writer.writerows(rows)
    metadata = dict(mode=mode, arguments={k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()},
                    candidates=candidates, unique_accepted=len(rows), rounds=round_index,
                    target_threshold=sampler.threshold, diffusion_steps=int(FLAGS.T),
                    beta_start=float(FLAGS.beta_1), beta_end=float(FLAGS.beta_T),
                    predictor='checkpoints/metalatte_predictor.ckpt', diffusion='checkpoints/metalorian_diffusion.pt')
    args.output.with_suffix('.metadata.json').write_text(json.dumps(metadata, indent=2)+'\n')
    print(f'Wrote {len(rows)} sequences to {args.output}; requested {args.desired}.')