"""Unified evaluation pipeline: all methods, all per-user data saved.

CRASH-SAFE: Each example is appended to a JSONL file immediately after
computation. If the process is killed, all completed examples are preserved.
Already-complete methods are automatically skipped on re-run.

Usage:
    python scripts/run_all_methods.py --task review --setting user --device cuda:0
    python scripts/run_all_methods.py --task review --setting user --methods base,uph,lora
"""

import sys
import os
import json
import time
import numpy as np
import torch
from scipy import stats

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data.longlamp import load_longlamp, select_k_profile_items
from data.templates import build_query_prompt, build_prompt_with_examples
from data.style_features import compute_sfd, compute_feature_deltas
from models.qwen_wrapper import QwenWrapper
from models.cvh import UnconditionalHead
from adapt.cache_hidden import cache_support_hidden_states
from adapt.fit_theta import fit_theta
from baselines.peft_baseline import (
    PEFTBaseline, get_lora_config, get_tiny_lora_config, get_vera_config,
    get_prompt_tuning_config, get_prefix_tuning_config,
)
from baselines.bm25_top1 import bm25_select_top1
from baselines.dense_retrieval import DenseRetriever
from baselines.profile_based import generate_profile, build_profile_conditioned_prompt
from eval.metrics import compute_rouge, compute_meteor


ALL_METHODS = [
    'base', 'uph',
    'prompt_all_k', 'bm25_top1', 'dense_top1', 'profile_based',
    'lora', 'tiny_lora', 'vera',
    'prompt_tuning_5', 'prompt_tuning_10', 'prompt_tuning_20',
    'prefix_tuning_5', 'prefix_tuning_10',
]


def compute_per_user_metrics(pred, ref, support_texts):
    r = compute_rouge([pred], [ref])
    m = compute_meteor([pred], [ref])
    p = pred if pred.strip() else "empty"
    sfd_all = compute_sfd(p, support_texts, exclude_length=False)
    sfd_nolen = compute_sfd(p, support_texts, exclude_length=True)
    deltas = compute_feature_deltas(p, support_texts)
    return {
        'rouge1': r['rouge1'],
        'rougeL': r['rougeL'],
        'meteor': m,
        'sfd_all': sfd_all,
        'sfd_nolen': sfd_nolen,
        'length': len(pred.split()),
        'feature_deltas': {k: v['delta'] for k, v in deltas.items()},
    }


def generate_greedy(wrapper, prompt, max_new_tokens=512, min_new_tokens=128):
    chat_messages = [
        {"role": "system", "content": "You are a helpful writing assistant."},
        {"role": "user", "content": prompt},
    ]
    prompt_text = wrapper.tokenizer.apply_chat_template(
        chat_messages, tokenize=False, add_generation_prompt=True
    )
    input_ids = wrapper.tokenizer.encode(prompt_text, return_tensors="pt").to(wrapper.device)
    with torch.no_grad():
        outputs = wrapper.model.generate(
            input_ids,
            max_new_tokens=max_new_tokens, min_new_tokens=min_new_tokens,
            temperature=None, top_p=None, do_sample=False,
            pad_token_id=wrapper.tokenizer.pad_token_id,
        )
    return wrapper.tokenizer.decode(outputs[0, input_ids.shape[1]:], skip_special_tokens=True)


# ─── Incremental saving ──────────────────────────────────────────────

def get_method_dir(output_dir, task, setting, K, method_name, d=64):
    """Get the output directory for a method."""
    exp_dir = os.path.join(output_dir, f"{task}_{setting}_K{K}")
    method_label = f"uph_d{d}" if method_name == 'uph' and d != 64 else method_name
    return os.path.join(exp_dir, method_label), method_label


def is_method_complete(method_dir, N):
    """Check if a method already has a complete per_user.json."""
    path = os.path.join(method_dir, 'per_user.json')
    if not os.path.exists(path):
        return False
    try:
        with open(path) as f:
            data = json.load(f)
        return len(data.get('per_user', [])) >= N
    except:
        return False


def append_jsonl(path, entry):
    """Append one JSON entry to a JSONL file (crash-safe)."""
    with open(path, 'a') as f:
        f.write(json.dumps(entry, default=str) + '\n')


def read_jsonl(path):
    """Read all entries from a JSONL file."""
    entries = []
    if os.path.exists(path):
        with open(path) as f:
            for line in f:
                line = line.strip()
                if line:
                    entries.append(json.loads(line))
    return entries


def finalize_method(method_dir, method_label, per_user, task, setting, K, d=64):
    """Write final per_user.json from completed per-user list."""
    agg = {
        'rougeL': float(np.mean([u['metrics']['rougeL'] for u in per_user])),
        'meteor': float(np.mean([u['metrics']['meteor'] for u in per_user])),
        'sfd_nolen': float(np.mean([u['metrics']['sfd_nolen'] for u in per_user])),
        'avg_len': float(np.mean([u['metrics']['length'] for u in per_user])),
    }
    save_data = {
        'per_user': per_user,
        'aggregate': agg,
        'num_examples': len(per_user),
        'task': task, 'setting': setting, 'K': K,
        'method': method_label,
        'decode_policy': 'greedy, min=128, max=512',
    }
    if 'uph' in method_label:
        save_data['d'] = d
    path = os.path.join(method_dir, 'per_user.json')
    with open(path, 'w') as f:
        json.dump(save_data, f, indent=2, default=str)
    print(f"  Saved: {path} ({len(per_user)} examples)")


# ─── Method runners ──────────────────────────────────────────────────

class MethodRunner:
    def __init__(self, wrapper, device, dense_retriever=None, uph_d=64):
        self.wrapper = wrapper
        self.device = device
        self.dense_retriever = dense_retriever
        self.uph_d = uph_d

    def _make_entry(self, ex, ref, stexts, K, pred, timing, extra=None):
        metrics = compute_per_user_metrics(pred, ref, stexts)
        entry = {
            'example_id': ex['example_id'],
            'user_id': ex['user_id'],
            'prediction': pred,
            'reference': ref,
            'support_texts': stexts,
            'K': K,
            'metrics': metrics,
            **timing,
        }
        if extra:
            entry.update(extra)
        return entry

    def run(self, method_name, examples, support_sets, references, support_texts,
            N, method_dir, method_label, task, setting, K, d=64):
        """Run a method with incremental JSONL saving. Returns per_user list."""

        dispatch = {
            'base': self._run_base,
            'uph': self._run_uph,
            'prompt_all_k': self._run_prompt_all_k,
            'bm25_top1': self._run_bm25_top1,
            'dense_top1': self._run_dense_top1,
            'profile_based': self._run_profile_based,
            'lora': lambda *a, **kw: self._run_peft(*a, config=get_lora_config(rank=8), lr=1e-4, desc='LoRA r=8', **kw),
            'tiny_lora': lambda *a, **kw: self._run_peft(*a, config=get_tiny_lora_config(rank=1), lr=1e-4, desc='Tiny LoRA r=1', **kw),
            'vera': lambda *a, **kw: self._run_peft(*a, config=get_vera_config(rank=256), lr=1e-3, desc='VeRA r=256', **kw),
            'prompt_tuning_5': lambda *a, **kw: self._run_peft(*a, config=get_prompt_tuning_config(5), lr=1e-3, desc='PromptTuning L=5', steps=100, **kw),
            'prompt_tuning_10': lambda *a, **kw: self._run_peft(*a, config=get_prompt_tuning_config(10), lr=1e-3, desc='PromptTuning L=10', steps=100, **kw),
            'prompt_tuning_20': lambda *a, **kw: self._run_peft(*a, config=get_prompt_tuning_config(20), lr=1e-3, desc='PromptTuning L=20', steps=100, **kw),
            'prefix_tuning_5': lambda *a, **kw: self._run_peft(*a, config=get_prefix_tuning_config(5), lr=5e-4, desc='PrefixTuning L=5', steps=100, **kw),
            'prefix_tuning_10': lambda *a, **kw: self._run_peft(*a, config=get_prefix_tuning_config(10), lr=5e-4, desc='PrefixTuning L=10', steps=100, **kw),
        }

        if method_name not in dispatch:
            print(f"Unknown method: {method_name}")
            return []

        os.makedirs(method_dir, exist_ok=True)
        jsonl_path = os.path.join(method_dir, 'progress.jsonl')

        # Resume: check how many examples already done
        existing = read_jsonl(jsonl_path)
        start_idx = len(existing)

        if start_idx >= N:
            print(f"\n--- {method_name} --- SKIPPED (already {start_idx}/{N} done)")
            per_user = existing[:N]
        else:
            if start_idx > 0:
                print(f"\n--- {method_name} --- RESUMING from {start_idx}/{N}")
            else:
                print(f"\n--- {method_name} ---")

            per_user = dispatch[method_name](
                examples, support_sets, references, support_texts, N,
                jsonl_path=jsonl_path, start_idx=start_idx, existing=existing,
            )

        avg_rl = np.mean([u['metrics']['rougeL'] for u in per_user])
        avg_sfd = np.mean([u['metrics']['sfd_nolen'] for u in per_user])
        print(f"  Mean R-L: {avg_rl:.4f}, SFD_-len: {avg_sfd:.4f}")

        # Write final per_user.json
        finalize_method(method_dir, method_label, per_user, task, setting, K, d)
        return per_user

    # --- Individual method runners ---
    # All accept jsonl_path, start_idx, existing for resume support

    def _run_base(self, examples, support_sets, references, support_texts, N,
                  jsonl_path, start_idx, existing):
        per_user = list(existing)
        for i in range(start_idx, N):
            ex = examples[i]
            t0 = time.time()
            prompt = build_query_prompt(ex['query_input'], ex['task'])
            pred = generate_greedy(self.wrapper, prompt)
            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support_sets[i]),
                pred, {'gen_time': time.time() - t0}
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 40 == 0:
                print(f"    {i+1}/{N}")
        return per_user

    def _run_prompt_all_k(self, examples, support_sets, references, support_texts, N,
                          jsonl_path, start_idx, existing):
        per_user = list(existing)
        for i in range(start_idx, N):
            ex, support = examples[i], support_sets[i]
            t0 = time.time()
            prompt = build_prompt_with_examples(ex['query_input'], support, ex['task'])
            pred = generate_greedy(self.wrapper, prompt)
            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support),
                pred, {'gen_time': time.time() - t0}
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 40 == 0:
                print(f"    {i+1}/{N}")
        return per_user

    def _run_bm25_top1(self, examples, support_sets, references, support_texts, N,
                       jsonl_path, start_idx, existing):
        per_user = list(existing)
        for i in range(start_idx, N):
            ex, support = examples[i], support_sets[i]
            t0 = time.time()
            selected = bm25_select_top1(ex['query_input'], support)
            prompt = build_prompt_with_examples(ex['query_input'], selected, ex['task'])
            pred = generate_greedy(self.wrapper, prompt)
            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support),
                pred, {'gen_time': time.time() - t0}
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 40 == 0:
                print(f"    {i+1}/{N}")
        return per_user

    def _run_dense_top1(self, examples, support_sets, references, support_texts, N,
                        jsonl_path, start_idx, existing):
        if self.dense_retriever is None:
            self.dense_retriever = DenseRetriever(device='cpu')
        per_user = list(existing)
        for i in range(start_idx, N):
            ex, support = examples[i], support_sets[i]
            t0 = time.time()
            selected = self.dense_retriever.retrieve_top_k(ex['query_input'], support, k=1)
            prompt = build_prompt_with_examples(ex['query_input'], selected, ex['task'])
            pred = generate_greedy(self.wrapper, prompt)
            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support),
                pred, {'gen_time': time.time() - t0}
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 40 == 0:
                print(f"    {i+1}/{N}")
        return per_user

    def _run_profile_based(self, examples, support_sets, references, support_texts, N,
                           jsonl_path, start_idx, existing):
        per_user = list(existing)
        for i in range(start_idx, N):
            ex, support = examples[i], support_sets[i]
            t0 = time.time()
            profile = generate_profile(self.wrapper, support, ex['task'])
            prompt = build_profile_conditioned_prompt(ex['query_input'], profile, ex['task'])
            pred = generate_greedy(self.wrapper, prompt)
            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support),
                pred, {'gen_time': time.time() - t0},
                extra={'profile_summary': profile},
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 40 == 0:
                print(f"    {i+1}/{N}")
        return per_user

    def _run_uph(self, examples, support_sets, references, support_texts, N,
                 jsonl_path, start_idx, existing):
        d = self.uph_d
        H = self.wrapper.hidden_size
        uncond = UnconditionalHead(H, d=d, alpha=0.1, basis_seed=42).to(self.device)
        print(f"  UPH d={d}, params={d}, bytes={d*2}")
        lm_head_bias = None
        if hasattr(self.wrapper.model.lm_head, 'bias') and self.wrapper.model.lm_head.bias is not None:
            lm_head_bias = self.wrapper.model.lm_head.bias.data

        per_user = list(existing)
        for i in range(start_idx, N):
            ex, support = examples[i], support_sets[i]
            t0 = time.time()
            cached_h = cache_support_hidden_states(self.wrapper, support, ex['task'])
            if not cached_h:
                prompt = build_query_prompt(ex['query_input'], ex['task'])
                pred = generate_greedy(self.wrapper, prompt)
            else:
                theta = fit_theta(
                    cached_h=cached_h,
                    lm_head_weight=self.wrapper.lm_head_weight,
                    lm_head_bias=lm_head_bias,
                    head_module=uncond,
                    d=d, lr=0.05, steps=30, beta=0.05, lam=1e-4,
                    max_grad_norm=5.0, device=self.device,
                )
                prompt = build_query_prompt(ex['query_input'], ex['task'])
                pred = self.wrapper.generate_with_head_blended(
                    prompt, theta, uncond.forward_fn,
                    blend_gamma=0.5, max_new_tokens=512,
                    min_new_tokens=128, temperature=0.0,
                )
                del cached_h, theta
                torch.cuda.empty_cache()

            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support),
                pred, {'adapt_time': time.time() - t0}
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 40 == 0:
                avg_rl = np.mean([u['metrics']['rougeL'] for u in per_user])
                print(f"    {i+1}/{N} (avg R-L: {avg_rl:.4f})")
        return per_user

    def _run_peft(self, examples, support_sets, references, support_texts, N,
                  config, lr, desc, steps=30, jsonl_path=None, start_idx=0, existing=None):
        if existing is None:
            existing = []
        baseline = PEFTBaseline(self.wrapper, config)
        print(f"  {desc}: {baseline.n_params:,} params ({baseline.n_bytes:,} bytes), steps={steps}, lr={lr}")

        per_user = list(existing)
        for i in range(start_idx, N):
            ex, support = examples[i], support_sets[i]
            t0 = time.time()
            pred = baseline.adapt_and_generate(
                support_items=support,
                query_input=ex['query_input'],
                task=ex['task'],
                lr=lr, steps=steps,
                max_new_tokens=512, min_new_tokens=128,
            )
            entry = self._make_entry(
                ex, references[i], support_texts[i], len(support),
                pred, {'adapt_time': time.time() - t0},
                extra={'n_params': baseline.n_params, 'n_bytes': baseline.n_bytes},
            )
            per_user.append(entry)
            append_jsonl(jsonl_path, entry)
            if (i + 1) % 20 == 0:
                avg_rl = np.mean([u['metrics']['rougeL'] for u in per_user])
                avg_t = np.mean([u['adapt_time'] for u in per_user])
                print(f"    {i+1}/{N} (avg R-L: {avg_rl:.4f}, avg time: {avg_t:.1f}s)")

        baseline.cleanup()
        return per_user


# ─── Main ────────────────────────────────────────────────────────────

def paired_test(scores_a, scores_b, name_a, name_b, metric_name):
    a, b = np.array(scores_a), np.array(scores_b)
    diff = a - b
    mean_diff = np.mean(diff)
    t_stat, t_pval = stats.ttest_rel(a, b)
    try:
        w_stat, w_pval = stats.wilcoxon(a, b)
    except ValueError:
        w_stat, w_pval = float('nan'), float('nan')
    se = stats.sem(diff)
    ci_low, ci_high = mean_diff - 1.96 * se, mean_diff + 1.96 * se
    return {
        'mean_a': float(np.mean(a)), 'mean_b': float(np.mean(b)),
        'mean_diff': float(mean_diff),
        'ci_low': float(ci_low), 'ci_high': float(ci_high),
        't_pval': float(t_pval), 'w_pval': float(w_pval),
    }


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_eval', type=int, default=200)
    parser.add_argument('--task', type=str, default='review', choices=['review', 'topic'])
    parser.add_argument('--setting', type=str, default='user', choices=['user', 'temporal'])
    parser.add_argument('--methods', type=str, default='all',
                        help='Comma-separated methods or "all"')
    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--K', type=int, default=4)
    parser.add_argument('--d', type=int, default=64, help='UPH theta dimension')
    parser.add_argument('--output_dir', type=str, default='outputs/unified')
    args = parser.parse_args()

    N = args.num_eval
    task = args.task
    setting = args.setting
    K = args.K

    config_map = {
        ('review', 'user'): 'product_review_user',
        ('review', 'temporal'): 'product_review_temporal',
        ('topic', 'user'): 'topic_writing_user',
        ('topic', 'temporal'): 'topic_writing_temporal',
    }
    config_name = config_map[(task, setting)]

    if args.methods == 'all':
        methods = ALL_METHODS
    else:
        methods = [m.strip() for m in args.methods.split(',')]

    print(f"=== Unified Eval: {task}_{setting}, N={N}, K={K}, d={args.d} ===")
    print(f"Methods: {methods}")
    print(f"Decode: greedy, min=128, max=512")

    print("\nLoading data...")
    examples = load_longlamp(config_name, split='val')[:N]
    support_sets = [select_k_profile_items(ex['profile_items'], K, seed=0) for ex in examples]
    references = [ex['target_output'] for ex in examples]
    support_texts = [[s['support_output'] for s in ss] for ss in support_sets]

    print(f"Loading model on {args.device}...")
    wrapper = QwenWrapper('Qwen/Qwen2.5-1.5B-Instruct', device=args.device)

    runner = MethodRunner(wrapper, args.device, uph_d=args.d)
    all_per_user = {}

    for method in methods:
        method_dir, method_label = get_method_dir(
            args.output_dir, task, setting, K, method, args.d
        )

        # Skip if already complete
        if is_method_complete(method_dir, N):
            print(f"\n--- {method} --- COMPLETE (loading from disk)")
            with open(os.path.join(method_dir, 'per_user.json')) as f:
                data = json.load(f)
            all_per_user[method] = data['per_user'][:N]
            avg_rl = np.mean([u['metrics']['rougeL'] for u in all_per_user[method]])
            print(f"  Mean R-L: {avg_rl:.4f}")
            continue

        per_user = runner.run(
            method, examples, support_sets, references, support_texts,
            N, method_dir, method_label, task, setting, K, args.d,
        )
        all_per_user[method] = per_user

    # Summary table
    print("\n" + "=" * 90)
    print(f"{'Method':<15} {'R-L':<8} {'METEOR':<8} {'SFD_-len':<9} {'Len':<6}")
    print("-" * 90)
    for method in methods:
        if method not in all_per_user:
            continue
        pu = all_per_user[method]
        rl = np.mean([u['metrics']['rougeL'] for u in pu])
        mt = np.mean([u['metrics']['meteor'] for u in pu])
        sf = np.mean([u['metrics']['sfd_nolen'] for u in pu])
        ln = np.mean([u['metrics']['length'] for u in pu])
        print(f"{method:<15} {rl:<8.4f} {mt:<8.4f} {sf:<9.4f} {ln:<6.0f}")

    # Significance tests (UPH vs all others)
    sig_results = {}
    if 'uph' in all_per_user:
        print("\n" + "=" * 90)
        print("Significance (UPH vs each, paired t-test p-value)")
        print("=" * 90)
        uph_rl = [u['metrics']['rougeL'] for u in all_per_user['uph']]
        uph_sf = [u['metrics']['sfd_nolen'] for u in all_per_user['uph']]
        for method in methods:
            if method == 'uph' or method not in all_per_user:
                continue
            other_rl = [u['metrics']['rougeL'] for u in all_per_user[method]]
            other_sf = [u['metrics']['sfd_nolen'] for u in all_per_user[method]]
            rl_t = paired_test(uph_rl, other_rl, 'uph', method, 'R-L')
            sf_t = paired_test(uph_sf, other_sf, 'uph', method, 'SFD')
            sig_results[method] = {'rougeL': rl_t, 'sfd_nolen': sf_t}
            print(f"  vs {method:<12} R-L: diff={rl_t['mean_diff']:+.4f} p={rl_t['t_pval']:.2e}  "
                  f"SFD: diff={sf_t['mean_diff']:+.4f} p={sf_t['t_pval']:.2e}")

    # Save summary
    exp_dir = os.path.join(args.output_dir, f"{task}_{setting}_K{K}")
    summary = {}
    for method in methods:
        if method not in all_per_user:
            continue
        pu = all_per_user[method]
        summary[method] = {
            'rougeL': float(np.mean([u['metrics']['rougeL'] for u in pu])),
            'meteor': float(np.mean([u['metrics']['meteor'] for u in pu])),
            'sfd_nolen': float(np.mean([u['metrics']['sfd_nolen'] for u in pu])),
            'avg_len': float(np.mean([u['metrics']['length'] for u in pu])),
        }
    summary_path = os.path.join(exp_dir, 'summary.json')
    with open(summary_path, 'w') as f:
        json.dump({
            'aggregate': summary,
            'significance': sig_results,
            'num_examples': N, 'task': task, 'setting': setting, 'K': K,
            'methods': methods,
        }, f, indent=2, default=str)

    print(f"\nSummary: {summary_path}")


if __name__ == '__main__':
    main()