#!/usr/bin/env python
"""
Evaluate Enhanced DoVLA-Attention checkpoint with SAME eval script as baseline.

This uses the REAL evaluation metric: selected_success_rate
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.data.datasets import CILDataset
from dovla_cil.models.dovla_attention_enhanced import DoVLAAttentionEnhanced
from dovla_cil.utils.io import write_json


def evaluate_enhanced_checkpoint(
    checkpoint_path: str | Path,
    dataset_dir: str | Path,
    output_path: str | Path | None = None,
    device: str = "auto",
) -> dict:
    """Evaluate enhanced model using SAME protocol as baseline."""

    resolved_device = (
        "cuda" if device == "auto" and torch.cuda.is_available() else "cpu"
        if device == "auto"
        else device
    )

    checkpoint = torch.load(checkpoint_path, map_location=resolved_device, weights_only=False)
    args = checkpoint["args"]

    # Recreate model from args
    model = DoVLAAttentionEnhanced(
        obs_dim=70,  # Padded
        action_dim=32,  # Padded
        hidden_dim=args["hidden_dim"],
        n_heads=args["n_heads"],
        n_layers=args["n_layers"],
        num_tasks=6,
        use_contrastive=True,
        use_graph=True,
        use_task_adaptive=True
    ).to(resolved_device)

    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    dataset = CILDataset(dataset_dir)

    # Use ALL groups for fair comparison with baseline
    val_group_ids = list(dataset.group_ids)

    selected_success = 0
    oracle_success = 0
    top1_correct = 0

    task_map = {
        "PickCube-v1": 0,
        "PushCube-v1": 1,
        "PullCube-v1": 2,
        "StackCube-v1": 3,
        "LiftPegUpright-v1": 4,
        "PegInsertionSide-v1": 5
    }

    def pad(vec, target):
        if len(vec) >= target:
            return vec[:target]
        return vec + [0.0] * (target - len(vec))

    with torch.no_grad():
        for group_id in val_group_ids:
            records = dataset.get_group(group_id)
            if not records:
                continue

            # Prepare observations (pad to 70)
            obs_list = []
            for r in records:
                obs_data = r.observation_inline
                if "state" in obs_data:
                    obs = list(obs_data["state"])
                else:
                    obs = []
                    for v in obs_data.values():
                        if isinstance(v, list):
                            obs.extend(v)
                        elif isinstance(v, (int, float)):
                            obs.append(v)
                obs = pad([float(x) for x in obs], 70)
                obs_list.append(obs)

            obs = torch.tensor([obs_list[0]], dtype=torch.float32, device=resolved_device)

            # Prepare actions (pad to 32)
            actions = torch.tensor(
                [pad(r.action_chunk.flat_values, 32) for r in records],
                dtype=torch.float32,
                device=resolved_device
            ).unsqueeze(0)

            # Task ID
            task_id = torch.tensor([task_map.get(records[0].task_id, 0)], device=resolved_device)

            # Forward
            scores_matrix, _ = model(obs, actions, task_id, None)

            # scores_matrix is (1, K, K) pairwise
            # Aggregate to per-action scores (sum of pairwise wins)
            scores = scores_matrix[0].sum(dim=1).cpu().tolist()

            # Select best action
            selected = max(range(len(records)), key=lambda i: (scores[i], -i))

            # Check success
            utilities = [r.reward.score for r in records]
            best_utility = max(utilities)

            is_selected_success = int(records[selected].reward.terminal_success)
            has_oracle_success = int(any(r.reward.terminal_success for r in records))
            is_top1 = int(abs(utilities[selected] - best_utility) < 1e-6)

            selected_success += is_selected_success
            oracle_success += has_oracle_success
            top1_correct += is_top1

    group_count = len(val_group_ids)
    result = {
        "checkpoint": str(checkpoint_path),
        "dataset": str(dataset_dir),
        "num_groups": group_count,
        "selected_success_rate": selected_success / group_count,
        "oracle_success_rate": oracle_success / group_count,
        "top1_action_selection": top1_correct / group_count,
    }

    if output_path:
        write_json(result, output_path)

    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--device", default="auto")

    args = parser.parse_args(argv)

    result = evaluate_enhanced_checkpoint(
        args.checkpoint, args.dataset, args.out, args.device
    )

    print(f"Selected success rate: {result['selected_success_rate']:.4f}")
    print(f"Top-1 selection: {result['top1_action_selection']:.4f}")
    print(f"Oracle success: {result['oracle_success_rate']:.4f}")

    return 0


if __name__ == "__main__":
    sys.exit(main())