"""Community comparison — compare a dataset against the benchmark database.

Finds similar entries by task_type and robot_type, filters by domain
(simulation vs real-world), then produces a structured comparison showing
how the user's dataset stacks up against same-domain peers.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path

logger = logging.getLogger(__name__)

GROUND_TRUTH_PATH = Path(__file__).parent / "data" / "ground_truth_condensed.json"

# Known simulation benchmark prefixes / keywords.  Entries that match any of
# these are tagged as simulation; everything else is assumed real-world.
_SIM_KEYWORDS = frozenset([
    "robomimic", "pusht", "metaworld", "maniskill", "libero",
    "ravens", "softgym", "furniture_bench",
])


def _is_sim_entry(entry: dict) -> bool:
    """Return True if a benchmark entry is from a simulation environment."""
    eid = entry.get("id", "").lower()
    dataset = entry.get("dataset", "").lower()
    summary = entry.get("features_summary", "").lower()

    # Explicit keyword match
    for kw in _SIM_KEYWORDS:
        if kw in eid or kw in dataset:
            return True

    # "sim" as a standalone token (e.g. "aloha_sim_*")
    for text in (eid, dataset):
        parts = text.replace("/", "_").replace("-", "_").split("_")
        if "sim" in parts:
            return True

    # Keyword in description
    if "simulation" in summary or "simulated" in summary:
        return True

    # Scripted demos for envs with scripted counterparts are sim
    if "scripted" in eid and any(kw in eid for kw in ("aloha", "transfer", "insertion")):
        return True

    return False


def _is_sim_dataset(dataset_name: str) -> bool:
    """Return True if the *user's* dataset is simulation."""
    lower = dataset_name.lower()
    parts = lower.replace("/", "_").replace("-", "_").split("_")
    if "sim" in parts:
        return True
    for kw in _SIM_KEYWORDS:
        if kw in lower:
            return True
    if "scripted" in lower:
        return True
    return False


def _difficulty_tier(success_rate: float, is_sim: bool) -> str:
    """Compute difficulty tier from success rate + domain."""
    if success_rate >= 0.9 and is_sim:
        return "easy"
    if success_rate >= 0.7 or (success_rate >= 0.5 and is_sim):
        return "medium"
    return "hard"


@dataclass
class PeerEntry:
    """A single benchmark peer for display."""

    id: str
    dataset: str
    num_episodes: int
    success_rate: float
    policy: str
    task_type: str
    is_sim: bool = False
    difficulty_tier: str = "medium"


@dataclass
class CommunityComparison:
    """Result of comparing a dataset against the benchmark database."""

    similar_successful: list[PeerEntry]  # peers with >70% success
    similar_all: list[PeerEntry]  # all matched peers (for stats)
    your_episodes: int
    peer_avg_episodes: float
    your_episode_percentile: str  # "bottom 25%", "middle 50%", etc.
    your_coverage: float | None  # overall coverage-like metric
    peer_avg_coverage: float | None
    actionable_tip: str  # one-liner recommendation
    # Domain-aware fields
    your_domain: str = "real"  # "sim" or "real"
    same_domain_count: int = 0
    cross_domain_count: int = 0
    includes_cross_domain: bool = False
    filtering_note: str = ""


def compute_community_comparison(
    dataset_name: str,
    task_type: str | None,
    robot_type: str | None,
    num_episodes: int,
    coverage_score: float | None = None,
    policy_type: str | None = None,
) -> CommunityComparison | None:
    """Find similar benchmark entries and build a comparison.

    Prefers same-domain matches (sim vs sim, real vs real).  Falls back to
    cross-domain matches when fewer than 3 same-domain peers exist.

    Returns None if the benchmark database can't be loaded or no similar
    entries are found.
    """
    try:
        with open(GROUND_TRUTH_PATH) as f:
            ground_truth: list[dict] = json.load(f)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.warning("Failed to load benchmark database: %s", e)
        return None

    if not ground_truth:
        return None

    # Determine user's domain
    user_is_sim = _is_sim_dataset(dataset_name)
    user_domain = "sim" if user_is_sim else "real"

    # Tag every benchmark entry with sim/real
    for entry in ground_truth:
        entry["_is_sim"] = _is_sim_entry(entry)

    # ------------------------------------------------------------------
    # Score each entry for similarity (unchanged logic)
    # ------------------------------------------------------------------
    scored: list[tuple[float, dict]] = []
    task_lower = (task_type or "").lower()
    robot_lower = (robot_type or "").lower()
    dataset_lower = dataset_name.lower()

    for entry in ground_truth:
        score = 0.0
        entry_task = entry.get("task_type", "").lower()
        entry_summary = entry.get("features_summary", "").lower()

        # Task type match (strongest signal)
        if task_lower and task_lower == entry_task:
            score += 5.0
        elif task_lower and task_lower in entry_task:
            score += 3.0

        # Robot type / bimanual matching via summary
        if robot_lower:
            if robot_lower in entry_summary:
                score += 3.0
            elif "aloha" in robot_lower and "bimanual" in entry_summary:
                score += 2.0
            elif "aloha" in dataset_lower and "aloha" in entry.get("id", "").lower():
                score += 3.0

        # Dataset name keyword overlap
        name_parts = [
            p for p in dataset_lower.replace("/", " ").replace("_", " ").split()
            if len(p) > 2 and p not in ("sim", "the", "and", "lerobot")
        ]
        for part in name_parts:
            if part in entry.get("id", "").lower() or part in entry_summary:
                score += 1.0

        if score > 0:
            scored.append((score, entry))

    if not scored:
        return None

    # Sort by similarity score
    scored.sort(key=lambda x: -x[0])

    # ------------------------------------------------------------------
    # Domain-aware filtering
    # ------------------------------------------------------------------
    same_domain = [(s, e) for s, e in scored if e["_is_sim"] == user_is_sim]
    cross_domain = [(s, e) for s, e in scored if e["_is_sim"] != user_is_sim]

    if len(same_domain) >= 3:
        # Enough same-domain matches — use only those
        top_entries = [e for _, e in same_domain[:8]]
        includes_cross = False
        cross_excluded = len(cross_domain)
    else:
        # Pad with cross-domain entries, up to 8 total
        n_cross_needed = min(8 - len(same_domain), len(cross_domain))
        top_entries = [e for _, e in same_domain]
        top_entries += [e for _, e in cross_domain[:n_cross_needed]]
        includes_cross = n_cross_needed > 0
        cross_excluded = len(cross_domain) - n_cross_needed

    same_domain_count = sum(1 for e in top_entries if e["_is_sim"] == user_is_sim)
    cross_domain_in_results = sum(1 for e in top_entries if e["_is_sim"] != user_is_sim)

    if not top_entries:
        return None

    # ------------------------------------------------------------------
    # Build peer entries with domain + difficulty labels
    # ------------------------------------------------------------------
    all_peers = [
        PeerEntry(
            id=e["id"],
            dataset=e["dataset"],
            num_episodes=e.get("num_episodes", 0),
            success_rate=e.get("success_rate", 0.0),
            policy=e.get("policy", "unknown"),
            task_type=e.get("task_type", "unknown"),
            is_sim=e["_is_sim"],
            difficulty_tier=_difficulty_tier(e.get("success_rate", 0), e["_is_sim"]),
        )
        for e in top_entries
    ]

    successful_peers = [p for p in all_peers if p.success_rate >= 0.70]

    # ------------------------------------------------------------------
    # Stats — only from same-domain peers for fair ranking
    # ------------------------------------------------------------------
    same_domain_peers = [p for p in all_peers if p.is_sim == user_is_sim]
    peer_episodes_same = [p.num_episodes for p in same_domain_peers if p.num_episodes > 0]
    peer_avg_ep = (
        sum(peer_episodes_same) / len(peer_episodes_same) if peer_episodes_same else 0
    )

    # Percentile — computed ONLY against same-domain peers
    if peer_episodes_same:
        below = sum(1 for ep in peer_episodes_same if num_episodes < ep)
        ratio = below / len(peer_episodes_same)
        if ratio >= 0.75:
            percentile = "bottom 25%"
        elif ratio >= 0.50:
            percentile = "below average"
        elif ratio >= 0.25:
            percentile = "above average"
        else:
            percentile = "top 25%"
    else:
        percentile = "unknown"

    # ------------------------------------------------------------------
    # Filtering note
    # ------------------------------------------------------------------
    domain_label = "simulation" if user_is_sim else "real-world"
    other_label = "real-world" if user_is_sim else "simulation"
    task_label = task_type.replace("_", " ") if task_type else "similar"

    if includes_cross:
        filtering_note = (
            f"Compared against {same_domain_count} {domain_label} {task_label} benchmarks "
            f"+ {cross_domain_in_results} {other_label} (labelled). "
            f"Note: includes {other_label} benchmarks — "
            + ("real-world success rates are typically 20-40% lower."
               if user_is_sim
               else "simulation success rates are typically 20-40% higher.")
        )
    else:
        total_cross = cross_excluded + cross_domain_in_results
        filtering_note = (
            f"Compared against {same_domain_count} {domain_label} {task_label} benchmarks"
            + (f" ({total_cross} {other_label} benchmarks excluded)"
               if total_cross > 0
               else "")
        )

    # Actionable tip
    tip = _generate_tip(num_episodes, peer_avg_ep, successful_peers, coverage_score)

    return CommunityComparison(
        similar_successful=successful_peers,
        similar_all=all_peers,
        your_episodes=num_episodes,
        peer_avg_episodes=peer_avg_ep,
        your_episode_percentile=percentile,
        your_coverage=coverage_score,
        peer_avg_coverage=None,  # benchmark doesn't store coverage
        actionable_tip=tip,
        your_domain=user_domain,
        same_domain_count=same_domain_count,
        cross_domain_count=cross_domain_in_results,
        includes_cross_domain=includes_cross,
        filtering_note=filtering_note,
    )


def _generate_tip(
    num_episodes: int,
    peer_avg_ep: float,
    successful_peers: list[PeerEntry],
    coverage_score: float | None,
) -> str:
    """Generate one actionable recommendation from the comparison."""
    if not successful_peers:
        return "No similar datasets with >70% success found \u2014 you're pioneering this task type."

    min_success_eps = min(p.num_episodes for p in successful_peers)
    avg_success_eps = sum(p.num_episodes for p in successful_peers) / len(successful_peers)

    if num_episodes < min_success_eps:
        gap = int(min_success_eps - num_episodes)
        return f"Collect ~{gap} more episodes to match the smallest successful peer ({min_success_eps} eps)."

    if num_episodes < avg_success_eps * 0.8:
        gap = int(avg_success_eps - num_episodes)
        return f"Collect ~{gap} more episodes to match successful peers (avg {int(avg_success_eps)} eps)."

    if coverage_score is not None and coverage_score < 0.70:
        return "Episode count matches peers \u2014 focus on diversity (varied start positions, speeds, approaches)."

    return "Your dataset size matches successful peers \u2014 focus on data quality and policy tuning."


def comparison_to_dict(comp: CommunityComparison) -> dict:
    """Convert to a JSON-serializable dict."""
    return {
        "similar_successful": [
            {
                "id": p.id,
                "dataset": p.dataset,
                "num_episodes": p.num_episodes,
                "success_rate": p.success_rate,
                "policy": p.policy,
                "is_sim": p.is_sim,
                "difficulty_tier": p.difficulty_tier,
            }
            for p in comp.similar_successful
        ],
        "similar_all": [
            {
                "id": p.id,
                "dataset": p.dataset,
                "num_episodes": p.num_episodes,
                "success_rate": p.success_rate,
                "policy": p.policy,
                "task_type": p.task_type,
                "is_sim": p.is_sim,
                "difficulty_tier": p.difficulty_tier,
            }
            for p in comp.similar_all
        ],
        "your_episodes": comp.your_episodes,
        "peer_avg_episodes": round(comp.peer_avg_episodes, 1),
        "your_episode_percentile": comp.your_episode_percentile,
        "actionable_tip": comp.actionable_tip,
        "your_domain": comp.your_domain,
        "same_domain_count": comp.same_domain_count,
        "cross_domain_count": comp.cross_domain_count,
        "includes_cross_domain": comp.includes_cross_domain,
        "filtering_note": comp.filtering_note,
    }
