#!/usr/bin/env python3
"""Evaluate held-out loss after applying HEAPr pruning masks in memory."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

import numpy as np

from heapr.constants import DEFAULT_GROUP_WIDTH
from heapr.eval.loss import evaluate_token_cache
from heapr.model_utils import (
    build_max_memory,
    load_causal_lm,
    validate_model_device_placement,
)
from heapr.prune import (
    apply_atomic_mask_to_model,
    apply_group_mask_to_model,
    atomic_mask_from_scores,
    group_mask_from_scores,
)
from heapr.utils import require_torch, write_json


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--cache-path", required=True)
    parser.add_argument("--scores-path", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--mode", choices=["atomic", "group"], default="atomic")
    parser.add_argument("--group-indices")
    parser.add_argument("--group-width", type=int, default=DEFAULT_GROUP_WIDTH)
    parser.add_argument("--ratios", nargs="+", type=float, default=[0.10, 0.20, 0.40])
    parser.add_argument("--revision")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--max-chunks", type=int)
    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")
    parser.add_argument("--cache-implementation", default="static")
    parser.add_argument("--no-cache", 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 == "group" and not args.group_indices:
        raise ValueError("--group-indices is required for grouped pruning")

    scores = np.load(args.scores_path)
    group_indices = np.load(args.group_indices) if args.group_indices else None
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    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

    for ratio in args.ratios:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            for device_idx in range(torch.cuda.device_count()):
                try:
                    torch.cuda.reset_peak_memory_stats(device_idx)
                except RuntimeError:
                    pass

        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=not args.no_cache,
            cache_implementation=args.cache_implementation,
        )
        validate_model_device_placement(
            model,
            allow_cpu_offload=args.allow_cpu_offload,
            requested_gpu_count=requested_gpu_count,
        )

        if args.mode == "atomic":
            keep_mask = atomic_mask_from_scores(scores, ratio)
            apply_atomic_mask_to_model(model, keep_mask)
            mask_summary = {
                "mode": "atomic",
                "ratio": ratio,
                "mask_shape": list(keep_mask.shape),
                "kept": int(keep_mask.sum()),
                "pruned": int((~keep_mask).sum()),
            }
        else:
            keep_mask = group_mask_from_scores(scores, ratio)
            apply_group_mask_to_model(
                model,
                keep_mask,
                group_width=args.group_width,
                group_indices=group_indices,
            )
            mask_summary = {
                "mode": "group",
                "ratio": ratio,
                "mask_shape": list(keep_mask.shape),
                "kept_groups": int(keep_mask.sum()),
                "pruned_groups": int((~keep_mask).sum()),
                "group_width": int(args.group_width),
                "group_indices_shape": list(group_indices.shape) if group_indices is not None else None,
            }

        metrics = evaluate_token_cache(
            model,
            args.cache_path,
            batch_size=args.batch_size,
            max_chunks=args.max_chunks,
            use_cache=not args.no_cache,
            cache_implementation=args.cache_implementation,
        )
        metrics["mask_summary"] = mask_summary
        output_path = output_dir / f"{args.mode}_pruned_{int(ratio * 100):02d}pct_loss.json"
        write_json(output_path, metrics)
        print(output_path)
        del model


if __name__ == "__main__":
    main()