#!/usr/bin/env python3
"""Generate SWE-bench patch predictions from baseline or in-memory pruned Laguna."""

from __future__ import annotations

import argparse
import json
import sys
from contextlib import nullcontext
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np
from tqdm.auto import tqdm

from heapr.grouped_model import RepackedExpandedGroupedLagunaContext
from heapr.model_utils import (
    build_max_memory,
    load_causal_lm,
    load_tokenizer,
    validate_model_device_placement,
)
from heapr.prune import apply_atomic_mask_to_model, apply_group_mask_to_model, atomic_mask_from_scores
from heapr.utils import require_torch


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--output-path", required=True)
    parser.add_argument("--dataset", default="princeton-nlp/SWE-bench_Lite_oracle")
    parser.add_argument("--split", default="test")
    parser.add_argument(
        "--mode",
        choices=["baseline", "atomic", "native-group", "repacked-group"],
        default="baseline",
    )
    parser.add_argument("--scores-path")
    parser.add_argument("--ratio", type=float)
    parser.add_argument("--group-mask")
    parser.add_argument("--group-indices")
    parser.add_argument(
        "--child-budget-mode",
        choices=["backfill", "layer-scaled", "parent-weighted"],
        default="backfill",
    )
    parser.add_argument("--max-instances", type=int)
    parser.add_argument("--max-new-tokens", type=int, default=200)
    parser.add_argument("--revision")
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--gpu-memory-per-device")
    parser.add_argument("--max-gpu-memory")
    parser.add_argument("--max-cpu-memory")
    parser.add_argument("--offload-folder")
    parser.add_argument("--allow-cpu-offload", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.offload_folder and not args.allow_cpu_offload:
        raise ValueError("--offload-folder requires --allow-cpu-offload")
    if args.mode == "atomic" and (args.scores_path is None or args.ratio is None):
        raise ValueError("--scores-path and --ratio are required for atomic mode")
    if args.mode in {"native-group", "repacked-group"} and (
        args.group_mask is None or args.group_indices is None
    ):
        raise ValueError("--group-mask and --group-indices are required for grouped modes")

    from datasets import load_dataset
    from swebench.inference.make_datasets.utils import extract_diff

    max_memory = build_max_memory(
        gpu_memory_per_device=args.gpu_memory_per_device,
        max_gpu_memory=args.max_gpu_memory,
        max_cpu_memory=args.max_cpu_memory,
        allow_cpu_offload=args.allow_cpu_offload,
    )
    torch = require_torch()
    requested_gpu_count = torch.cuda.device_count() if args.gpu_memory_per_device else None
    model = load_causal_lm(
        args.model_id,
        revision=args.revision,
        dtype=args.dtype,
        max_memory=max_memory,
        offload_folder=args.offload_folder if args.allow_cpu_offload else None,
        use_cache=True,
    )
    validate_model_device_placement(
        model,
        allow_cpu_offload=args.allow_cpu_offload,
        requested_gpu_count=requested_gpu_count,
    )
    tokenizer = load_tokenizer(args.model_id, revision=args.revision)

    context = nullcontext(model)
    if args.mode == "atomic":
        apply_atomic_mask_to_model(
            model,
            atomic_mask_from_scores(np.load(args.scores_path), args.ratio),
        )
    elif args.mode == "native-group":
        apply_group_mask_to_model(
            model,
            np.load(args.group_mask),
            group_indices=np.load(args.group_indices),
        )
    elif args.mode == "repacked-group":
        context = RepackedExpandedGroupedLagunaContext(
            model,
            keep_group_masks=np.load(args.group_mask),
            group_indices=np.load(args.group_indices),
            child_budget_mode=args.child_budget_mode,
        )

    dataset = load_dataset(args.dataset, split=args.split)
    if args.max_instances is not None:
        dataset = dataset.select(range(min(args.max_instances, len(dataset))))

    output_path = Path(args.output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    input_device = model.device
    with context, output_path.open("w", encoding="utf-8") as output_file, torch.no_grad():
        for instance in tqdm(dataset, desc="SWE-bench predictions"):
            inputs = tokenizer(instance["text"], return_tensors="pt")
            inputs = {key: value.to(input_device) for key, value in inputs.items()}
            generated = model.generate(
                **inputs,
                do_sample=False,
                max_new_tokens=args.max_new_tokens,
                use_cache=True,
            )
            completion = tokenizer.decode(
                generated[0, inputs["input_ids"].shape[-1] :],
                skip_special_tokens=False,
            )
            prediction = {
                "instance_id": instance["instance_id"],
                "full_output": completion,
                "model_patch": extract_diff(completion),
                "model_name_or_path": args.model_id,
            }
            output_file.write(json.dumps(prediction, sort_keys=True) + "\n")
            output_file.flush()
    print(output_path)


if __name__ == "__main__":
    main()