"""5.2 / 5.4 Base model vs. post-trained model: same family, same size, same prompts.

The base model (Qwen3.5-xB-Base) has only gone through pre-training; each prompt is fed as plain text for it to continue.
The post-trained model (Qwen3.5-xB) has gone through SFT and preference alignment; each prompt is fed in its chat format.
Each prompt is generated once at temperature 0, then sampled several times at temperature 0.7 to see whether the behavior is stable.

Usage:
    python compare.py                      # default: 2B, Chinese prompts
    python compare.py --size 0.8B          # machines with less memory
    python compare.py --size 9B --lang en  # larger model, English prompts (9B needs ~20 GB VRAM or memory)
    python compare.py --prompts my_prompts.json
The two models are loaded one after the other so only one model occupies memory at a time. Uses MLX on Mac if mlx-lm is installed; otherwise uses transformers (auto-selects CUDA / MPS / CPU).
"""
import argparse
import gc
import json
from datetime import datetime
from pathlib import Path


HERE = Path(__file__).resolve().parent


def pick_backend(choice):
    """Use MLX on Mac if mlx-lm is installed (much faster); otherwise use transformers and auto-select CUDA / MPS / CPU."""
    if choice in ('auto', 'mlx'):
        try:
            import mlx_lm  # noqa: F401
            return 'mlx'
        except ImportError:
            if choice == 'mlx':
                raise SystemExit('mlx-lm is not installed: uv pip install -r requirements-mac.txt')
    import torch
    if torch.cuda.is_available():
        return 'cuda'
    return 'mps' if torch.backends.mps.is_available() else 'cpu'


def load(repo, backend):
    """Return a function gen(text, temperature, n) -> n continuations."""
    if backend == 'mlx':
        import mlx.core as mx
        from mlx_lm import generate, load as mlx_load
        from mlx_lm.sample_utils import make_sampler
        model, tok = mlx_load(repo)

        def gen(text, temperature, n, max_new_tokens):
            sampler = make_sampler(temp=temperature, top_p=0.95 if temperature else 0.0)
            return [generate(model, tok, prompt=text, max_tokens=max_new_tokens, sampler=sampler) for _ in range(n)]
        return tok, gen, lambda: mx.clear_cache()

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(repo)
    model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.float32 if backend == 'cpu' else torch.bfloat16).to(backend)
    model.eval()

    def gen(text, temperature, n, max_new_tokens):
        ids = tok(text, return_tensors='pt').to(backend)
        kw = dict(do_sample=True, temperature=temperature, top_p=0.95, top_k=0) if temperature else dict(do_sample=False)
        out = model.generate(**ids, max_new_tokens=max_new_tokens, num_return_sequences=n, pad_token_id=tok.eos_token_id, **kw)
        return [tok.decode(o[ids['input_ids'].shape[1]:], skip_special_tokens=True) for o in out]

    def free():
        nonlocal model
        del model
        gc.collect()
        if backend == 'cuda':
            torch.cuda.empty_cache()
        elif backend == 'mps':
            torch.mps.empty_cache()
    return tok, gen, free


def run_model(repo, prompts, chat, backend, max_new_tokens, samples, seed):
    tok, gen, free = load(repo, backend)
    if backend == 'mlx':
        import mlx.core as mx
        mx.random.seed(seed)
    else:
        import torch
        torch.manual_seed(seed)
    out = {}
    for p in prompts:
        if chat:      # post-trained model: apply chat template with thinking mode disabled
            text = tok.apply_chat_template([{'role': 'user', 'content': p['text']}], tokenize=False,
                                           add_generation_prompt=True, enable_thinking=False)
        else:         # base model: pass raw text as-is
            text = p['text']
        greedy = gen(text, 0, 1, max_new_tokens)[0].strip()
        sampled = gen(text, 0.7, samples, max_new_tokens) if samples else []
        out[p['id']] = dict(greedy=greedy, samples=[x.strip() for x in sampled])
        print(f'  {repo} · {p["id"]} done')
    free()
    return out


def one_line(s, limit=400):
    s = s.replace('\n', ' ⏎ ')
    return s if len(s) <= limit else s[:limit] + '... (truncated)'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--size', default='2B', help='0.8B / 2B / 4B / 9B')
    ap.add_argument('--lang', choices=['zh', 'en'], default='zh')
    ap.add_argument('--prompts', help='Prompt file (default: prompts_<lang>.json)')
    ap.add_argument('--max-new-tokens', type=int, default=150)
    ap.add_argument('--samples', type=int, default=3, help='Number of samples per prompt at temperature 0.7')
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--backend', choices=['auto', 'mlx', 'torch'], default='auto')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    device = pick_backend(args.backend)
    prompts = json.loads(Path(args.prompts or HERE / f'prompts_{args.lang}.json').read_text(encoding='utf-8'))
    base_repo, inst_repo = f'Qwen/Qwen3.5-{args.size}-Base', f'Qwen/Qwen3.5-{args.size}'
    print(f'{base_repo} vs {inst_repo}, device {device}, {datetime.now():%Y-%m-%d %H:%M}')
    res = {'base': run_model(base_repo, prompts, False, device, args.max_new_tokens, args.samples, args.seed),
           'instruct': run_model(inst_repo, prompts, True, device, args.max_new_tokens, args.samples, args.seed)}

    lines = [f'# {base_repo} vs {inst_repo} ({device}, {datetime.now():%Y-%m-%d})', '',
             f'Generated once at temperature 0 (up to {args.max_new_tokens} tokens), then sampled {args.samples} times at temperature 0.7. ⏎ indicates a newline.']
    for p in prompts:
        lines += ['', f'## {p["kind"]}: {p["id"]}', '', f'Note: {p["note"]}', '', f'Prompt: {one_line(p["text"])}', '']
        for who, label in [('base', 'base'), ('instruct', 'post-trained')]:
            r = res[who][p['id']]
            lines.append(f'- **{label} (temperature 0)**: {one_line(r["greedy"])}')
            for i, s in enumerate(r['samples'], 1):
                lines.append(f'    - Sample {i}: {one_line(s, 200)}')
    report = '\n'.join(lines) + '\n'
    print('\n' + report)
    Path(args.out).mkdir(exist_ok=True)
    stem = f'{args.size}_{args.lang}_{device}_{datetime.now():%Y%m%d_%H%M}'
    (Path(args.out) / f'{stem}.md').write_text(report, encoding='utf-8')
    (Path(args.out) / f'{stem}.json').write_text(json.dumps(dict(prompts=prompts, results=res), ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'Report saved to {args.out}/{stem}.md')


if __name__ == '__main__':
    main()