"""Validate Orbit's success predictions against benchmark ground truth.

Loads benchmark entries linked to datasets with Orbit analysis, runs the
predictor, and computes accuracy metrics. Supports before/after comparison
with the new advanced analysis adjustment factors.

Usage:
    cd /Users/rahillasne/orbit-robotics/backend
    python -m orbit.analyzer.validation.validate_predictions

Or from the repo root:
    python src/orbit/analyzer/validation/validate_predictions.py
"""

from __future__ import annotations

import json
import logging
import sqlite3
import sys
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


def _find_db() -> str:
    """Find orbit.db in expected locations."""
    candidates = [
        Path(__file__).parent.parent.parent.parent.parent / "backend" / "orbit.db",
        Path.cwd() / "orbit.db",
        Path.cwd() / "backend" / "orbit.db",
    ]
    for p in candidates:
        if p.exists():
            return str(p)
    raise FileNotFoundError("Could not find orbit.db")


def load_validation_pairs(db_path: str | None = None) -> list[dict]:
    """Load benchmark entries that have linked datasets with Orbit analysis.

    Returns list of dicts with:
        - actual_success_rate: from benchmark
        - fingerprint: dataset fingerprint JSON
        - policy: policy used
        - task_name: task description
        - robot: robot type
        - num_episodes: training episodes
        - quality_score: Orbit quality score
    """
    if db_path is None:
        db_path = _find_db()

    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row

    rows = conn.execute("""
        SELECT
            br.id as benchmark_id,
            br.success_rate,
            br.policy,
            br.task_name,
            br.robot,
            br.num_training_episodes,
            br.dataset_name,
            di.orbit_quality_score,
            di.dataset_fingerprint,
            di.hf_repo_id,
            di.total_episodes,
            di.robot_type
        FROM benchmark_results br
        JOIN dataset_index di ON br.dataset_index_id = di.id
        WHERE di.has_orbit_analysis = 1
          AND di.dataset_fingerprint IS NOT NULL
          AND br.success_rate >= 0
          AND br.success_rate <= 1
    """).fetchall()
    conn.close()

    pairs = []
    for r in rows:
        fp = r["dataset_fingerprint"]
        if isinstance(fp, str):
            try:
                fp = json.loads(fp)
            except json.JSONDecodeError:
                continue

        pairs.append({
            "benchmark_id": r["benchmark_id"],
            "actual_success_rate": r["success_rate"],
            "policy": r["policy"],
            "task_name": r["task_name"],
            "robot": r["robot"] or r["robot_type"],
            "num_episodes": r["num_training_episodes"] or r["total_episodes"],
            "quality_score": r["orbit_quality_score"],
            "fingerprint": fp,
            "hf_repo_id": r["hf_repo_id"],
        })

    return pairs


def predict_from_fingerprint(
    pair: dict,
    use_advanced: bool = False,
) -> float | None:
    """Run the success predictor using fingerprint data directly.

    Constructs TaskContext and PolicyFit from fingerprint metrics
    rather than generating synthetic episodes, giving the predictor
    real dataset characteristics.

    Args:
        pair: Validation pair dict from load_validation_pairs.
        use_advanced: Whether to include advanced analysis adjustments
            (simulated from fingerprint data).

    Returns:
        Predicted success rate or None if prediction fails.
    """
    try:
        from orbit.analyzer.gp_features import _infer_episode_sufficiency
        from orbit.analyzer.policy_fit import (
            POLICY_PROFILES,
            PolicyFit,
            normalize_policy_name,
        )
        from orbit.analyzer.success_predictor import CalibratedSuccessPredictor
        from orbit.analyzer.task_context import TaskContext

        fp = pair["fingerprint"]
        quality = pair["quality_score"]
        if quality is None:
            quality = 0.5

        # Normalize quality from 0-100 to 0-1 if needed
        if quality is not None and quality > 1.0:
            quality = quality / 100.0

        num_episodes = pair["num_episodes"] or 50

        # Extract fingerprint sub-dicts
        cov = fp.get("coverage", {})
        sig = fp.get("signal_health", {})
        act_stats = fp.get("action_stats", {})
        ep_stats = fp.get("episode_stats", {})

        # Infer dimensions from fingerprint
        action_dim = act_stats.get("action_dim", 6)
        is_bimanual = action_dim >= 12
        mean_ep_length = ep_stats.get("mean_length", 100)
        consistency = cov.get("consistency", 0.5)

        # Build TaskContext from fingerprint metrics
        dim_score = min(1.0, action_dim / 14.0)
        length_score = min(1.0, mean_ep_length / 500.0)
        consistency_penalty = max(0, 0.2 * (1.0 - consistency))

        task_complexity = (
            0.30 * dim_score
            + 0.25 * length_score
            + 0.25 * consistency_penalty
            + 0.20 * 0.5  # default coordination (unknown without real episodes)
        )
        if is_bimanual:
            task_complexity = min(1.0, task_complexity + 0.15)

        task_context = TaskContext(
            task_complexity_score=task_complexity,
            dimension_complexity=dim_score,
            temporal_complexity=length_score,
            coordination_complexity=0.5,
            object_interaction_complexity=0.5,
            trajectory_multimodality=0.3,
            is_bimanual=is_bimanual,
            action_dim=action_dim,
            mean_episode_length=mean_ep_length,
            estimated_task_type="moderate",
        )

        # Build PolicyFit from fingerprint metrics
        policy_family = normalize_policy_name(pair["policy"] or "act")
        profile = POLICY_PROFILES.get(
            policy_family, POLICY_PROFILES["diffusion_policy"]
        )

        ep_suff = _infer_episode_sufficiency(
            num_episodes, task_complexity, policy_family
        )

        # Action space compatibility
        dim_low, dim_high = profile.get("action_dim_range", (2, 14))
        if dim_low <= action_dim <= dim_high:
            action_compat = 1.0
        elif action_dim < dim_low:
            action_compat = max(0.3, action_dim / dim_low)
        else:
            action_compat = max(0.3, dim_high / action_dim)

        # Consistency compatibility
        cons_sensitivity = profile.get("consistency_sensitivity", 0.5)
        consistency_compat = (
            consistency * cons_sensitivity + (1.0 - cons_sensitivity)
        )

        fit_score = float(
            np.clip(
                ep_suff * 0.30
                + action_compat * 0.20
                + consistency_compat * 0.50,
                0.0,
                1.0,
            )
        )

        policy_fit = PolicyFit(
            fit_score=fit_score,
            episode_sufficiency=ep_suff,
            action_space_compatibility=action_compat,
            consistency_compatibility=consistency_compat,
            policy_type=policy_family,
        )

        # Signal diagnostics from fingerprint
        alive_ratio = sig.get("joints_alive_ratio", 1.0)
        dead_count = max(0, int((1.0 - alive_ratio) * action_dim))
        dead_dims = list(range(dead_count))
        clipping = sig.get("clipping_rate", 0.0)

        blockers = []
        warnings = []
        if dead_count > 0:
            blockers.append(f"{dead_count} dead joints")
        if clipping > 0.2:
            warnings.append("Clipping detected")
        if consistency < 0.5:
            warnings.append("Low consistency")

        sig_dict = {
            "blockers": blockers,
            "warnings": warnings,
            "dead_dimensions": dead_dims,
        }

        predictor = CalibratedSuccessPredictor()

        # Build exclusion set to prevent data leakage: exclude this
        # benchmark entry and near-duplicates (same repo + policy) from GT
        exclude_ids: set[str] = set()
        benchmark_id = pair.get("benchmark_id")
        if benchmark_id is not None:
            exclude_ids.add(f"db_{benchmark_id}")

        hf_repo = (pair.get("hf_repo_id") or "").lower()
        policy_norm = normalize_policy_name(pair.get("policy") or "")
        for e in predictor._ground_truth:
            e_dataset = (e.get("dataset", "") or "").lower()
            e_policy_norm = normalize_policy_name(e.get("policy") or "")
            if hf_repo and hf_repo in e_dataset and e_policy_norm == policy_norm:
                exclude_ids.add(e.get("id", ""))

        kwargs = dict(
            data_quality=quality,
            task_context=task_context,
            policy_fit=policy_fit,
            signal_diagnostics=sig_dict,
            task_type=pair["task_name"],
            episode_count=num_episodes,
            exclude_ids=exclude_ids,
        )

        if use_advanced:
            jerk_level = act_stats.get("mean_jerk", 0.5)
            if jerk_level > 1.0:
                penalty = min(0.15, jerk_level * 0.02)
                kwargs["action_divergence"] = {
                    "estimated_success_penalty": penalty,
                }

            std_len = ep_stats.get("std_length", 0)
            mean_len = ep_stats.get("mean_length", 100)
            if mean_len > 0 and std_len / max(1, mean_len) > 0.3:
                remove_ratio = 0.25
                kwargs["episode_ranking"] = {
                    "remove_count": int(remove_ratio * num_episodes),
                    "total": num_episodes,
                }

        pred = predictor.predict(**kwargs)
        return pred.predicted_success_rate

    except Exception as e:
        logger.debug("Prediction failed for %s: %s", pair.get("hf_repo_id"), e)
        return None


def compute_metrics(
    actuals: list[float],
    predicteds: list[float],
) -> dict:
    """Compute validation metrics.

    Returns dict with: mae, spearman_rho, within_10pct, within_20pct, n
    """
    actuals_arr = np.array(actuals)
    preds_arr = np.array(predicteds)
    n = len(actuals_arr)

    if n < 2:
        return {"mae": float("nan"), "spearman_rho": float("nan"),
                "within_10pct": 0.0, "within_20pct": 0.0, "n": n}

    # Mean absolute error
    mae = float(np.mean(np.abs(actuals_arr - preds_arr)))

    # Spearman rank correlation
    from scipy.stats import spearmanr
    rho, _ = spearmanr(actuals_arr, preds_arr)

    # Within ±10% and ±20%
    abs_diff = np.abs(actuals_arr - preds_arr)
    within_10 = float(np.mean(abs_diff <= 0.10))
    within_20 = float(np.mean(abs_diff <= 0.20))

    return {
        "mae": round(mae, 4),
        "spearman_rho": round(float(rho), 4),
        "within_10pct": round(within_10, 4),
        "within_20pct": round(within_20, 4),
        "n": n,
    }


def run_validation(db_path: str | None = None, max_pairs: int = 200) -> dict:
    """Run full validation: before/after comparison.

    Returns dict with baseline_metrics, advanced_metrics, and improvement.
    """
    pairs = load_validation_pairs(db_path)
    print(f"Loaded {len(pairs)} validation pairs (benchmark + Orbit analysis)")

    if len(pairs) < 3:
        print("WARNING: Not enough validation pairs. Need at least 3.")
        print("Run the benchmark linkage script first:")
        print("  cd backend && python -m app.scripts.link_benchmarks")
        return {"error": "insufficient_pairs", "n": len(pairs)}

    # Limit to max_pairs for speed
    if len(pairs) > max_pairs:
        rng = np.random.RandomState(42)
        indices = rng.choice(len(pairs), size=max_pairs, replace=False)
        pairs = [pairs[i] for i in indices]
        print(f"Sampled {max_pairs} pairs for validation")

    # Baseline (without advanced factors)
    print("\nRunning baseline predictions...")
    baseline_actuals = []
    baseline_preds = []
    for pair in pairs:
        pred = predict_from_fingerprint(pair, use_advanced=False)
        if pred is not None:
            baseline_actuals.append(pair["actual_success_rate"])
            baseline_preds.append(pred)

    baseline_metrics = compute_metrics(baseline_actuals, baseline_preds)
    print(f"Baseline: {baseline_metrics['n']} predictions")
    print(f"  MAE: {baseline_metrics['mae']:.4f}")
    print(f"  Spearman rho: {baseline_metrics['spearman_rho']:.4f}")
    print(f"  Within ±10%: {baseline_metrics['within_10pct']:.1%}")
    print(f"  Within ±20%: {baseline_metrics['within_20pct']:.1%}")

    # Advanced (with new adjustment factors)
    print("\nRunning advanced predictions...")
    advanced_actuals = []
    advanced_preds = []
    for pair in pairs:
        pred = predict_from_fingerprint(pair, use_advanced=True)
        if pred is not None:
            advanced_actuals.append(pair["actual_success_rate"])
            advanced_preds.append(pred)

    advanced_metrics = compute_metrics(advanced_actuals, advanced_preds)
    print(f"Advanced: {advanced_metrics['n']} predictions")
    print(f"  MAE: {advanced_metrics['mae']:.4f}")
    print(f"  Spearman rho: {advanced_metrics['spearman_rho']:.4f}")
    print(f"  Within ±10%: {advanced_metrics['within_10pct']:.1%}")
    print(f"  Within ±20%: {advanced_metrics['within_20pct']:.1%}")

    # Comparison
    mae_improvement = baseline_metrics["mae"] - advanced_metrics["mae"]
    rho_improvement = advanced_metrics["spearman_rho"] - baseline_metrics["spearman_rho"]
    print(f"\nImprovement:")
    print(f"  MAE: {mae_improvement:+.4f} ({'better' if mae_improvement > 0 else 'worse'})")
    print(f"  Spearman rho: {rho_improvement:+.4f} ({'better' if rho_improvement > 0 else 'worse'})")

    return {
        "baseline_metrics": baseline_metrics,
        "advanced_metrics": advanced_metrics,
        "mae_improvement": round(mae_improvement, 4),
        "rho_improvement": round(rho_improvement, 4),
    }


if __name__ == "__main__":
    logging.basicConfig(level=logging.WARNING)

    # Try to find DB
    db_path = None
    for p in [
        Path("orbit.db"),
        Path("backend/orbit.db"),
        Path(__file__).parent.parent.parent.parent.parent / "backend" / "orbit.db",
    ]:
        if p.exists():
            db_path = str(p)
            break

    if db_path is None:
        print("ERROR: Could not find orbit.db")
        sys.exit(1)

    result = run_validation(db_path)
    if "error" not in result:
        print("\nValidation complete!")
