"""Post-training verification.

Compares predicted quality (from orbit gate/analyze) against actual
training outcome. Stores calibration data for improving future predictions.
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)

CALIBRATION_DIR = Path.home() / ".orbit" / "calibration"


@dataclass
class VerificationResult:
    """Comparison of prediction vs actual training outcome."""

    predicted_grade: str  # "B"
    predicted_score: int  # 84
    predicted_success_range: tuple[float, float]  # (0.65, 0.80)
    actual_final_loss: float  # 0.12
    actual_loss_converged: bool  # True
    training_status: str  # "converged", "plateau", "diverged", "nan"
    prediction_accuracy: str  # "correct", "optimistic", "pessimistic"
    feedback: str  # Human-readable summary


def verify_training(
    training_dir: str,
    dataset: str | None = None,
    success_rate: float | None = None,
) -> VerificationResult:
    """Compare predicted vs actual training outcome.

    Args:
        training_dir: Path to training output directory.
        dataset: Optional dataset repo ID for loading predictions.
        success_rate: Optional actual success rate from rollout evaluation.

    Returns:
        VerificationResult comparing prediction to reality.
    """
    from orbit.debugger import (
        _analyze_dynamics,
        _find_config,
        _parse_lerobot_logs,
        _parse_wandb_logs,
    )

    training_path = Path(training_dir)

    # Parse training logs
    log_entries = _parse_lerobot_logs(training_path)
    wandb_entries = _parse_wandb_logs(training_path)
    all_entries = sorted(log_entries + wandb_entries, key=lambda x: x.get("step", 0))

    dynamics = _analyze_dynamics(all_entries)
    config = _find_config(training_path)

    # Extract actual training metrics
    actual_final_loss = _extract_final_loss(all_entries)
    training_status = dynamics["status"]
    converged = training_status in ("converged_well", "partial_convergence")

    # Get predictions (try to load from prior analysis)
    predicted_grade, predicted_score, predicted_range = _load_predictions(
        training_path, dataset, config
    )

    # Run prediction if we have the dataset and no cached prediction
    if predicted_grade == "?" and dataset:
        predicted_grade, predicted_score, predicted_range = _run_prediction(dataset)

    # Determine prediction accuracy
    accuracy = _assess_accuracy(
        predicted_grade=predicted_grade,
        predicted_range=predicted_range,
        training_status=training_status,
        actual_final_loss=actual_final_loss,
        success_rate=success_rate,
    )

    # Generate feedback
    feedback = _generate_feedback(
        predicted_grade=predicted_grade,
        predicted_score=predicted_score,
        predicted_range=predicted_range,
        training_status=training_status,
        actual_final_loss=actual_final_loss,
        accuracy=accuracy,
        success_rate=success_rate,
    )

    result = VerificationResult(
        predicted_grade=predicted_grade,
        predicted_score=predicted_score,
        predicted_success_range=predicted_range,
        actual_final_loss=actual_final_loss,
        actual_loss_converged=converged,
        training_status=training_status,
        prediction_accuracy=accuracy,
        feedback=feedback,
    )

    # Store calibration data if success_rate provided
    if success_rate is not None:
        _store_calibration(result, success_rate, dataset)

    return result


def print_verification(result: VerificationResult, dataset: str | None = None) -> None:
    """Print verification result with rich formatting."""
    from orbit.utils.display import console

    console.print("\n[bold]Post-Training Verification[/bold]")
    console.print("\u2500" * 40)

    # Prediction
    if result.predicted_grade != "?":
        console.print(
            f"\n  Predicted: Grade {result.predicted_grade} "
            f"(score {result.predicted_score}), "
            f"success {result.predicted_success_range[0]:.0%}-{result.predicted_success_range[1]:.0%}"
        )
    else:
        console.print("\n  Predicted: No prior prediction available")

    # Actual
    status_colors = {
        "converged_well": "green",
        "partial_convergence": "yellow",
        "converged_high": "yellow",
        "plateau": "red",
        "nan_loss": "red",
        "loss_explosion": "red",
        "gradient_instability": "red",
    }
    color = status_colors.get(result.training_status, "white")
    status_display = result.training_status.replace("_", " ").upper()

    if not math.isnan(result.actual_final_loss):
        console.print(
            f"  Actual:    [{color}]{status_display}[/{color}], "
            f"final loss {result.actual_final_loss:.4f}"
        )
    else:
        console.print(f"  Actual:    [{color}]{status_display}[/{color}]")

    # Verdict
    accuracy_icons = {
        "correct": ("\u2713", "green"),
        "optimistic": ("\u26a0", "yellow"),
        "pessimistic": ("\u2713", "cyan"),
        "unknown": ("?", "dim"),
    }
    icon, acc_color = accuracy_icons.get(result.prediction_accuracy, ("?", "dim"))
    console.print(
        f"\n  Verdict: [{acc_color}]{icon} PREDICTION {result.prediction_accuracy.upper()}[/{acc_color}]"
    )
    console.print(f"  {result.feedback}")

    if result.prediction_accuracy != "unknown" and dataset:
        console.print(
            f"\n  [dim]If you've run eval rollouts, add: "
            f"--success-rate <rate> to calibrate future predictions.[/dim]"
        )
    console.print()


def _extract_final_loss(entries: list[dict]) -> float:
    """Extract the final valid loss value from entries."""
    valid = [e["loss"] for e in entries if not math.isnan(e["loss"]) and not math.isinf(e["loss"])]
    return valid[-1] if valid else float("nan")


def _load_predictions(
    training_path: Path,
    dataset: str | None,
    config: dict | None,
) -> tuple[str, int, tuple[float, float]]:
    """Try to load predictions from prior analysis artifacts."""
    # Check for orbit analysis output in training dir
    for f in training_path.rglob("*orbit*.json"):
        try:
            data = json.loads(f.read_text())
            if "readiness" in data:
                r = data["readiness"]
                return (
                    r.get("grade", "?"),
                    r.get("score", 0),
                    tuple(r.get("success_range", (0, 0))),
                )
        except (OSError, json.JSONDecodeError, KeyError):
            continue

    return "?", 0, (0.0, 0.0)


def _run_prediction(dataset: str) -> tuple[str, int, tuple[float, float]]:
    """Run orbit analysis to get a prediction."""
    try:
        from orbit.analyzer.coverage import compute_coverage, _extract_array
        from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
        from orbit.analyzer.quality import compute_quality_score
        from orbit.analyzer.success_predictor import DatasetReadinessGrader

        info = load_dataset_info(dataset)
        episode_data = load_episode_data(dataset, max_episodes=100)

        if episode_data is None or episode_data.empty:
            return "?", 0, (0.0, 0.0)

        coverage = compute_coverage(info, episode_data)
        quality = compute_quality_score(info, episode_data)

        action_episodes = []
        for _, ep_df in episode_data.groupby("episode_index"):
            if "action" in ep_df.columns:
                arr = _extract_array(ep_df["action"])
                if arr is not None and arr.size > 0:
                    action_episodes.append(arr)

        grader = DatasetReadinessGrader()
        readiness = grader.grade(
            quality_score=quality.overall_score * 100,
            episode_consistency=coverage.episode_consistency,
            signal_health={"dead_dimensions": [], "blockers": [], "warnings": [],
                           "total_joints": 6, "clipping_joints": []},
            policy_fit_score=0.7,
            episode_ranking={"total": len(action_episodes), "remove_count": 0, "review_count": 0},
            action_divergence=0.0,
            episode_count=info.num_episodes,
            policy_type="act",
        )

        return (
            readiness.grade,
            readiness.score,
            getattr(readiness, "success_range", (0.0, 0.0)),
        )
    except Exception:
        return "?", 0, (0.0, 0.0)


def _assess_accuracy(
    predicted_grade: str,
    predicted_range: tuple[float, float],
    training_status: str,
    actual_final_loss: float,
    success_rate: float | None,
) -> str:
    """Assess whether prediction matched reality."""
    if predicted_grade == "?":
        return "unknown"

    # If we have actual success rate, compare directly
    if success_rate is not None and predicted_range != (0.0, 0.0):
        if predicted_range[0] <= success_rate <= predicted_range[1]:
            return "correct"
        elif success_rate < predicted_range[0]:
            return "optimistic"
        else:
            return "pessimistic"

    # Otherwise, compare training status to grade
    good_outcomes = {"converged_well", "partial_convergence"}
    bad_outcomes = {"nan_loss", "loss_explosion", "gradient_instability", "plateau"}
    good_grades = {"A", "B", "C"}

    if training_status in good_outcomes and predicted_grade in good_grades:
        return "correct"
    elif training_status in bad_outcomes and predicted_grade in ("D", "F"):
        return "correct"
    elif training_status in good_outcomes and predicted_grade in ("D", "F"):
        return "pessimistic"
    elif training_status in bad_outcomes and predicted_grade in good_grades:
        return "optimistic"

    return "correct"  # converged_high with C/D is roughly expected


def _generate_feedback(
    predicted_grade: str,
    predicted_score: int,
    predicted_range: tuple[float, float],
    training_status: str,
    actual_final_loss: float,
    accuracy: str,
    success_rate: float | None,
) -> str:
    """Generate human-readable feedback."""
    if accuracy == "correct":
        if training_status in ("converged_well", "partial_convergence"):
            return "Training outcome matches prediction. The dataset quality was sufficient."
        else:
            return "Training failed as predicted. The dataset quality issues flagged by orbit were real."
    elif accuracy == "optimistic":
        return (
            "Prediction was too optimistic — training failed despite a passing grade. "
            "Run orbit analyze --deep to identify hidden data issues."
        )
    elif accuracy == "pessimistic":
        return (
            "Prediction was too pessimistic — training succeeded despite a low grade. "
            "The model was more robust than expected."
        )
    else:
        loss_str = f" Final loss: {actual_final_loss:.4f}." if not math.isnan(actual_final_loss) else ""
        return f"No prior prediction to compare.{loss_str}"


def _store_calibration(
    result: VerificationResult,
    success_rate: float,
    dataset: str | None,
) -> None:
    """Store calibration data for future regression."""
    try:
        CALIBRATION_DIR.mkdir(parents=True, exist_ok=True)
        entry = {
            "predicted_grade": result.predicted_grade,
            "predicted_score": result.predicted_score,
            "predicted_success_range": list(result.predicted_success_range),
            "actual_success_rate": success_rate,
            "actual_final_loss": result.actual_final_loss
            if not math.isnan(result.actual_final_loss)
            else None,
            "training_status": result.training_status,
            "prediction_accuracy": result.prediction_accuracy,
            "dataset": dataset,
        }

        # Append to calibration log
        cal_file = CALIBRATION_DIR / "calibration_log.jsonl"
        with open(cal_file, "a") as f:
            f.write(json.dumps(entry) + "\n")

        logger.info("Calibration data stored in %s", cal_file)
    except OSError as e:
        logger.debug("Failed to store calibration: %s", e)
