"""Deep analysis: LLM interprets orbit metrics like a senior robotics engineer.

Not a summary — a diagnosis. The prompt injects hundreds of hours of LeRobot
training experience so the model can pattern-match failure modes from metrics.
"""

from __future__ import annotations

import json
import logging

from orbit.llm_utils import cache_get, cache_key, cache_set, call_gemini, load_api_key

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Encoded training expertise
# ---------------------------------------------------------------------------

TRAINING_EXPERTISE = """
You have trained hundreds of imitation learning policies with LeRobot. Here is what you know
that most beginners don't:

ACT (Action Chunking with Transformers):
- Loss = L1 loss on action chunks + optional KL divergence (if use_vae=True)
- Healthy ACT training: loss starts at 0.5-1.0, drops to 0.02-0.08 by step 50k-100k
- If loss stalls above 0.15: almost always a data problem (multimodal demos, inconsistent demonstrations)
- batch_size=1 is a silent killer — gradients oscillate, loss looks like it's training but robot doesn't move
- chunk_size should be ~25-50% of episode length. Too large = model tries to plan too far ahead. Too small = no temporal coherence
- ACT uses temporal ensembling at inference (overlapping chunks averaged). If demonstrations have discontinuities, the ensembling produces jerky motion
- 50 episodes is often enough for a single-variation task, but only if demonstrations are consistent
- When ACT produces "frozen robot" at eval: check action normalization. If mean/std are wrong, all predictions collapse to zero

Diffusion Policy:
- Loss = MSE between predicted noise and actual noise (epsilon prediction)
- Healthy training: loss starts at 0.8-1.0, drops to 0.02-0.05 by step 100k-200k
- DP explicitly models multimodal action distributions — handles demonstrations with multiple strategies
- horizon = length of action sequence predicted at once. Default 16. Longer horizon = more temporal coherence but harder to train
- n_action_steps = how many predicted actions to actually execute before re-planning. Default 8
- DP is more data-hungry than ACT — needs 100+ episodes for reliable training
- DP fails gracefully — it doesn't produce dangerous intermediate actions like ACT does with multimodal data

SmolVLA / VLA models:
- Pre-trained vision-language-action models being fine-tuned
- Need 10-50 episodes for fine-tuning (not 200+)
- Very sensitive to image quality and camera angle consistency
- Action scaling is critical — VLA models output in a normalized space
- Use low learning rate (2e-5) and short training to avoid forgetting pre-trained skills

BC / BC-RNN:
- BC is extremely sensitive to data quality — one bad episode can ruin training
- BC-RNN handles temporal dependencies better but needs more data (100+)
- If BC works but ACT/DP doesn't, your data is probably fine and the issue is hyperparameters

General training knowledge:
- #1 cause of "robot doesn't move": action normalization mismatch between training and deployment
- #2 cause: camera processing differs between training and deployment (resolution, crop)
- #3 cause: batch_size too small (must be >=4 for stable training, >=8 preferred)
- Loss going to NaN: usually learning rate too high (try 5e-5), or dead joints causing gradient explosion
- Loss plateau at high value (>0.15): data problem, not hyperparameter problem. Fix the data.
- Overfitting with small dataset: the policy memorizes exact trajectories and fails on any variation
- Directional bias on joints: robot only moves one direction in training → same at eval
"""


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


def run_deep_analysis(
    dataset_id: str,
    episode_count: int,
    action_dim: int,
    fps: float | None,
    robot_type: str,
    policy_type: str,
    is_simulation: bool,
    readiness_grade: str,
    readiness_score: int,
    coverage_results: dict | None,
    signal_results: dict | None,
    episode_ranking: dict | None,
    divergence_results: dict | None,
    temporal_results: dict | None,
    policy_fit_results: dict | None,
    quality_score: float | None,
    recommendations: list | None,
    no_cache: bool = False,
) -> str | None:
    """Run LLM-powered deep analysis on orbit metrics.

    Returns markdown-formatted expert diagnosis, or None if unavailable.
    """
    api_key = load_api_key()
    if not api_key:
        return None

    metrics = _format_metrics(
        dataset_id, episode_count, action_dim, fps, robot_type,
        policy_type, is_simulation, readiness_grade, readiness_score,
        coverage_results, signal_results, episode_ranking,
        divergence_results, temporal_results, policy_fit_results,
        quality_score, recommendations,
    )

    ck = cache_key(metrics)
    if not no_cache:
        cached = cache_get(ck, "deep_analysis")
        if cached:
            return cached

    prompt = f"""{TRAINING_EXPERTISE}

---

A user just ran `orbit analyze` on their dataset. Here are ALL the metrics.
Read them carefully. Then give your expert assessment.

{metrics}

---

Your assessment must follow this EXACT structure:

**What will happen when you train**
2-3 sentences. Be concrete: "The robot will learn to approach but freeze before grasping"
not "performance may be affected." Predict the specific failure mode if there is one.
Reference the policy type — ACT and Diffusion fail differently on the same data.

**The real issue**
1-2 sentences. Every dataset has a weakness. Name it specifically with the number from the metrics.
If the grade is A or B, the issue is minor — say what would push it to the next level.

**Exactly what to do**
3-4 bullet points with SPECIFIC actions. Not "collect more data." Instead:
"Collect 20 more episodes where [specific gap from coverage analysis]."
"Remove episodes [specific indices] — they were flagged for [specific reason]."
Include orbit commands: `orbit clean`, `orbit suggest`, `orbit analyze --full`.

**What to watch during training**
3 bullet points. Include SPECIFIC numbers:
"Loss should drop below [X] by step [Y] for {policy_type}."
"If loss plateaus above 0.15 after step 50k, the issue is data, not hyperparameters."
"At eval, if the robot [specific behavior], it means [specific data issue]."

Keep it under 350 words total. No hedging. No "it depends." Give your best call."""

    result = call_gemini(prompt, api_key, max_tokens=4096)
    if result:
        cache_set(ck, result, "deep_analysis")
    return result


# ---------------------------------------------------------------------------
# Metrics formatting
# ---------------------------------------------------------------------------


def _format_metrics(
    dataset_id, episode_count, action_dim, fps, robot_type,
    policy_type, is_simulation, readiness_grade, readiness_score,
    coverage, signal, ranking, divergence, temporal,
    policy_fit, quality, recommendations,
) -> str:
    """Format all metrics into a structured block for the LLM."""
    sections = [
        f"DATASET: {dataset_id}",
        f"Episodes: {episode_count} | Action dims: {action_dim} | FPS: {fps or 'unknown'}",
        f"Robot: {robot_type} | Policy: {policy_type} | {'Simulation' if is_simulation else 'Real robot'}",
        f"Readiness: {readiness_grade} ({readiness_score}/100) | Quality: {quality or 'N/A'}",
    ]

    if coverage:
        sections.append("\nCOVERAGE:")
        for k, v in coverage.items():
            if not isinstance(v, (dict, list)):
                sections.append(f"  {k}: {v}")

    if signal:
        sections.append("\nSIGNAL HEALTH:")
        # Extract key facts instead of dumping verbose JSON
        dead = signal.get("dead_dimensions", [])
        if dead:
            sections.append(f"  Dead joints: {dead}")
        blockers = signal.get("blockers", [])
        for b in blockers:
            sections.append(f"  BLOCKER: {b}")
        warnings_list = signal.get("warnings", [])
        for w in warnings_list[:8]:
            sections.append(f"  WARNING: {w}")
        passing = signal.get("passing", [])
        for p in passing[:4]:
            sections.append(f"  OK: {p}")
        clipping = signal.get("clipping_joints", [])
        if clipping:
            sections.append(f"  Clipping joints: {clipping}")
        remove_eps = signal.get("episodes_to_remove", [])
        if remove_eps:
            sections.append(f"  Episodes to remove: {remove_eps[:20]}")
        review_eps = signal.get("episodes_to_review", [])
        if review_eps:
            sections.append(f"  Episodes to review: {review_eps[:20]}")

    if ranking:
        sections.append("\nEPISODE RANKING:")
        for k, v in ranking.items():
            if not isinstance(v, (dict, list)):
                sections.append(f"  {k}: {v}")

    if divergence:
        sections.append("\nACTION DIVERGENCE:")
        for k, v in divergence.items():
            if not isinstance(v, (dict, list)):
                sections.append(f"  {k}: {v}")

    if temporal:
        sections.append("\nTEMPORAL ALIGNMENT:")
        sections.append(f"  {json.dumps(temporal, indent=2, default=str)[:1000]}")

    if policy_fit:
        sections.append("\nPOLICY FIT:")
        sections.append(f"  {json.dumps(policy_fit, indent=2, default=str)[:1000]}")

    if recommendations:
        sections.append("\nTOP RECOMMENDATIONS:")
        for r in recommendations[:5]:
            if isinstance(r, dict):
                sections.append(f"  - {r.get('description', r.get('text', str(r)))}")
            elif hasattr(r, "description"):
                sections.append(f"  - [{r.priority}] {r.description}")
            else:
                sections.append(f"  - {r}")

    return "\n".join(sections)
