#!/usr/bin/env python
"""
Evaluate DoVLA-Hybrid checkpoint with DIRECT selection.
"""
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_hybrid import DoVLAHybrid
from dovla_cil.utils.io import write_json


def evaluate_hybrid_checkpoint(
    checkpoint_path: str | Path,
    dataset_dir: str | Path,
    output_path: str | Path | None = None,
    device: str = "auto",
) -> dict:
    """Evaluate Hybrid with direct selection."""

    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"]

    model = DoVLAHybrid(
        obs_dim=70,
        action_dim=32,
        lang_dim=0,
        d_model=args["d_model"],
        n_heads=args["n_heads"],
        n_layers=args["n_layers"],
        d_ff=args["d_ff"],
        dropout=0.1
    ).to(resolved_device)

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

    dataset = CILDataset(dataset_dir)
    val_group_ids = list(dataset.group_ids)

    selected_success = 0
    oracle_success = 0
    top1_correct = 0

    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

            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)

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

            # DIRECT prediction
            pred_rewards, pred_success_probs = model(obs, actions, lang=None)

            # Hybrid selection
            hybrid_scores = (pred_success_probs * pred_rewards).squeeze(0).cpu().tolist()
            selected = max(range(len(records)), key=lambda i: hybrid_scores[i])

            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,
        "approach": "hybrid_direct_scoring"
    }

    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_hybrid_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())