"""Training suggestion generator — recommends policy, hyperparameters, and a ready-to-run command.

No LLM calls, no API keys. Works entirely offline using dataset analysis
and encoded knowledge about policy architectures from policy_fit.py.

Supports multiple frameworks: LeRobot, OpenVLA, OpenPI/pi0, GR00T N1,
NVIDIA OSMO, and custom training scripts via --framework flag.
"""

from __future__ import annotations

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

import numpy as np

from orbit.analyzer.policy_fit import POLICY_PROFILES, PolicyFitAnalyzer, normalize_policy_name

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


@dataclass
class TrainingSuggestion:
    """A complete training configuration suggestion."""

    dataset_repo: str
    policy_type: str
    policy_full_name: str
    fit_score: float

    # Dataset summary
    episode_count: int
    action_dim: int
    state_dim: int
    fps: float
    camera_count: int
    camera_resolution: str  # e.g. "480p"

    # Hyperparameters
    chunk_size: int
    batch_size: int
    steps: int
    learning_rate: float
    image_resolution: str | None  # None if no cameras

    # The command
    train_command: str

    # Tips and warnings
    tips: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    # Alternative policy if close
    alternative: dict | None = None  # {"policy": "...", "reason": "..."}


# ---------------------------------------------------------------------------
# GPU memory -> batch size lookup
# ---------------------------------------------------------------------------

# batch_size recommendations: (cameras, resolution_bucket) -> {gpu_gb: batch_size}
_BATCH_SIZES = {
    (1, "480p"): {8: 8, 16: 16, 24: 32, 40: 64, 80: 128},
    (1, "720p"): {8: 4, 16: 8, 24: 16, 40: 32, 80: 64},
    (1, "1080p"): {8: 2, 16: 4, 24: 8, 40: 16, 80: 32},
    (2, "480p"): {8: 4, 16: 8, 24: 16, 40: 32, 80: 64},
    (2, "720p"): {8: 2, 16: 4, 24: 8, 40: 16, 80: 32},
    (2, "1080p"): {8: 2, 16: 2, 24: 4, 40: 8, 80: 16},
}


def _pick_batch_size(camera_count: int, resolution: str, gpu_gb: int | None) -> int:
    """Pick batch size based on camera config and GPU memory."""
    if gpu_gb is None:
        gpu_gb = 24  # assume A5000/RTX 4090

    cam_key = min(camera_count, 2)
    res_key = resolution if resolution in ("480p", "720p", "1080p") else "480p"
    lookup = _BATCH_SIZES.get((cam_key, res_key), _BATCH_SIZES[(1, "480p")])

    # Find closest GPU tier
    best_bs = 8
    for tier_gb, bs in sorted(lookup.items()):
        if gpu_gb >= tier_gb:
            best_bs = bs

    # Never return batch_size=1 — causes silent training failures with ACT
    return max(2, best_bs)


def _resolution_bucket(shapes: dict) -> str:
    """Determine the resolution bucket from dataset shapes."""
    for key, shape in shapes.items():
        if "image" in key.lower() or "video" in key.lower() or "observation.image" in key.lower():
            if len(shape) >= 2:
                h = shape[-3] if len(shape) >= 3 else shape[-2]
                if h >= 1080:
                    return "1080p"
                if h >= 720:
                    return "720p"
    return "480p"


# ---------------------------------------------------------------------------
# Main suggestion logic
# ---------------------------------------------------------------------------


def generate_suggestion(
    dataset_repo: str,
    info,
    episode_data=None,
    policy: str | None = None,
    gpu_memory_gb: int | None = None,
) -> TrainingSuggestion:
    """Generate a training suggestion for a dataset.

    Args:
        dataset_repo: HuggingFace repo ID.
        info: DatasetInfo from dataset_loader.
        episode_data: Optional DataFrame for deeper analysis.
        policy: Optional policy name. Auto-selects if None.
        gpu_memory_gb: Available GPU memory for batch size calculation.

    Returns:
        TrainingSuggestion with everything needed to start training.
    """
    from orbit.analyzer.coverage import _extract_array, compute_coverage
    from orbit.analyzer.signal_diagnostics import compute_signal_diagnostics

    # Determine action/state dims from shapes
    action_dim = _dim_from_shapes(info.shapes, "action")
    state_dim = _dim_from_shapes(info.shapes, "state")
    camera_count = len(info.camera_names) if info.camera_names else 0
    resolution = _resolution_bucket(info.shapes)
    fps = info.fps or 30.0

    # Build action episodes for analysis
    action_episodes = []
    state_episodes = []
    consistency_score = 0.7  # default
    multimodality_score = 0.3

    if episode_data is not None and not episode_data.empty:
        ep_groups = episode_data.groupby("episode_index")
        for _, ep_df in ep_groups:
            if "action" in ep_df.columns:
                ep_actions = _extract_array(ep_df["action"])
                if ep_actions is not None and ep_actions.size > 0:
                    action_episodes.append(ep_actions)
                    if ep_actions.ndim == 2:
                        action_dim = ep_actions.shape[1]
            if "state" in ep_df.columns:
                ep_states = _extract_array(ep_df["state"])
                if ep_states is not None and ep_states.size > 0:
                    state_episodes.append(ep_states)
                    if ep_states.ndim == 2:
                        state_dim = ep_states.shape[1]

        if action_episodes:
            coverage = compute_coverage(info, episode_data)
            consistency_score = coverage.episode_consistency

    # Check for dead joints
    dead_joints = []
    if action_episodes:
        try:
            sig_diag = compute_signal_diagnostics(action_episodes)
            dead_joints = sig_diag.dead_dimensions
        except Exception as e:
            import logging as _log
            _log.getLogger(__name__).debug("Signal diagnostics failed in suggester: %s", e)

    # Compute median episode length
    if action_episodes:
        lengths = [len(ep) for ep in action_episodes]
        median_ep_length = int(np.median(lengths))
    else:
        median_ep_length = int(info.num_frames / max(1, info.num_episodes))

    # Policy selection
    tips = []
    warnings = []
    alternative = None

    fit_score = 0.7  # default if we couldn't compute

    if policy:
        policy_type = normalize_policy_name(policy)
    else:
        # Auto-select via PolicyFitAnalyzer
        if action_episodes:
            pf = PolicyFitAnalyzer()
            pf_result = pf.analyze(
                action_episodes,
                policy_type="auto",
                episode_count=info.num_episodes,
                task_complexity=0.5,
                consistency_score=consistency_score,
                multimodality_score=multimodality_score,
            )
            policy_type = pf_result.policy_type
            fit_score = pf_result.fit_score
            if pf_result.recommended_policies and len(pf_result.recommended_policies) > 1:
                top = pf_result.recommended_policies[0]
                second = pf_result.recommended_policies[1]
                if second["fit_score"] > top["fit_score"] - 0.05:
                    if abs(second["fit_score"] - top["fit_score"]) < 0.01:
                        reason = f"Both score {top['fit_score']:.2f} — equally good fit"
                    else:
                        reason = f"Close fit — scores {second['fit_score']:.2f} vs {top['fit_score']:.2f}"
                    alternative = {
                        "policy": second["policy"],
                        "full_name": second["full_name"],
                        "fit_score": second["fit_score"],
                        "reason": reason,
                    }
        else:
            # Fallback: pick based on episode count heuristic
            if info.num_episodes < 50:
                policy_type = "act"
            elif info.num_episodes < 200:
                policy_type = "diffusion_policy"
            else:
                policy_type = "diffusion_policy"

    # Compute fit score for user-specified policy (auto-select already set it above)
    if policy and action_episodes:
        try:
            pf = PolicyFitAnalyzer()
            pf_result = pf.analyze(
                action_episodes,
                policy_type=policy_type,
                episode_count=info.num_episodes,
                task_complexity=0.5,
                consistency_score=consistency_score,
                multimodality_score=multimodality_score,
            )
            fit_score = pf_result.fit_score
        except Exception as e:
            import logging as _log
            _log.getLogger(__name__).debug("Policy fit analysis failed, using default: %s", e)

    profile = POLICY_PROFILES.get(policy_type, POLICY_PROFILES["diffusion_policy"])
    policy_full_name = profile["full_name"]

    # Hyperparameters
    chunk_size = _pick_chunk_size(policy_type, median_ep_length)
    batch_size = _pick_batch_size(camera_count, resolution, gpu_memory_gb)
    steps = _pick_steps(policy_type, info.num_episodes)
    lr = _pick_learning_rate(policy_type)
    image_res = None if camera_count == 0 else resolution

    # Batch size validation — if too small, suggest resizing
    if batch_size < 4 and resolution in ("720p", "1080p"):
        warnings.append(
            f"Your images are {resolution}. Resize to 480p to fit batch_size=8: "
            "add --dataset.image_transforms.enable=true"
        )
        tips.append(f"Tip: 480p is usually sufficient for manipulation tasks")

    # Never batch_size=1
    if batch_size < 2:
        batch_size = 2
        warnings.append("Batch size forced to 2 — batch_size=1 causes silent training failures with ACT")

    # Dead joint warning
    if dead_joints:
        effective_dim = action_dim - len(dead_joints)
        warnings.append(
            f"Dead joints detected: {dead_joints}. "
            f"Consider --policy.action_dim={effective_dim} (reduced from {action_dim})"
        )

    # Episode count warnings
    min_eps = profile.get("min_episodes_simple", 25)
    if info.num_episodes < 30:
        warnings.append(
            f"Very small dataset ({info.num_episodes} episodes). "
            "Consider collecting 20-30 more episodes before training"
        )
    elif info.num_episodes < min_eps:
        warnings.append(
            f"Dataset has {info.num_episodes} episodes — "
            f"{policy_type} works best with {min_eps}+"
        )

    # Consistency warning
    if consistency_score < 0.5:
        warnings.append(
            f"Low demonstration consistency ({consistency_score:.2f}). "
            "Review your teleop technique — are you performing the task the same way each time?"
        )

    # Training tips
    tips.append(f"Loss should drop below 0.1 by step {steps // 5}")
    tips.append("If loss plateaus above 0.2: your demonstrations may be too inconsistent")
    tips.append("If robot doesn't move during eval: check action normalization stats")

    # Build the command
    output_dir = f"outputs/train/{dataset_repo.split('/')[-1]}-{policy_type}"
    cmd_parts = [
        "lerobot-train",
        f"  --dataset.repo_id={dataset_repo}",
        f"  --policy.type={policy_type}",
        f"  --batch_size={batch_size}",
        f"  --steps={steps}",
    ]
    if policy_type == "act":
        cmd_parts.append(f"  --policy.chunk_size={chunk_size}")
        cmd_parts.append(f"  --policy.n_action_steps={chunk_size}")
    elif policy_type == "diffusion_policy":
        cmd_parts.append(f"  --policy.horizon={chunk_size}")
    elif policy_type == "smolvla":
        pass  # uses model defaults

    cmd_parts.extend([
        f"  --log_freq=200",
        f"  --save_freq={max(1000, steps // 5)}",
        f"  --eval_freq={max(1000, steps // 5)}",
        f"  --policy.device=cuda",
        f"  --output_dir={output_dir}",
    ])

    train_command = " \\\n".join(cmd_parts)

    return TrainingSuggestion(
        dataset_repo=dataset_repo,
        policy_type=policy_type,
        policy_full_name=policy_full_name,
        fit_score=fit_score,
        episode_count=info.num_episodes,
        action_dim=action_dim,
        state_dim=state_dim,
        fps=fps,
        camera_count=camera_count,
        camera_resolution=resolution,
        chunk_size=chunk_size,
        batch_size=batch_size,
        steps=steps,
        learning_rate=lr,
        image_resolution=image_res,
        train_command=train_command,
        tips=tips,
        warnings=warnings,
        alternative=alternative,
    )


def _dim_from_shapes(shapes: dict, key: str) -> int:
    """Extract dimensionality from shapes dict."""
    for k, v in shapes.items():
        if key in k.lower():
            if isinstance(v, (list, tuple)) and len(v) >= 1:
                return v[-1]
    return 6  # default


def _pick_chunk_size(policy_type: str, median_ep_length: int) -> int:
    if policy_type == "act":
        return min(100, max(10, median_ep_length // 2))
    elif policy_type in ("diffusion_policy", "dp3"):
        return min(16, max(4, median_ep_length // 10))
    else:
        return 50  # default


def _pick_steps(policy_type: str, episode_count: int) -> int:
    if policy_type == "act":
        if episode_count < 100:
            return 100_000
        elif episode_count < 500:
            return 200_000
        else:
            return 300_000
    elif policy_type in ("diffusion_policy", "dp3"):
        if episode_count < 200:
            return 200_000
        else:
            return 500_000
    elif policy_type == "smolvla":
        return 50_000  # fine-tuning
    else:
        return 200_000


def _pick_learning_rate(policy_type: str) -> float:
    if policy_type == "smolvla":
        return 2e-5
    return 1e-4


# ---------------------------------------------------------------------------
# Rich output
# ---------------------------------------------------------------------------


def print_suggestion(suggestion: TrainingSuggestion) -> None:
    """Print a training suggestion with rich formatting."""
    from rich.panel import Panel
    from rich.table import Table

    from orbit.utils.display import console

    s = suggestion
    console.print()

    # Summary panel
    fit_label = f"fit: {s.fit_score:.2f}" if s.fit_score else "best fit"
    title = f"Recommended: {s.policy_full_name} ({fit_label})"
    table = Table(show_header=False, border_style="cyan", title=title)
    table.add_column("Metric", style="bold")
    table.add_column("Value")
    table.add_row("Episodes", f"{s.episode_count}")
    table.add_row("Cameras", f"{s.camera_count} x {s.camera_resolution}" if s.camera_count else "none")
    table.add_row("Action dims", str(s.action_dim))
    table.add_row("FPS", f"{s.fps:.0f}")
    if s.alternative:
        alt = s.alternative
        alt_fit = f"fit: {alt['fit_score']:.2f}" if alt.get("fit_score") else ""
        table.add_row(
            "Alternative",
            f"{alt['full_name']} ({alt_fit}) — {alt['reason']}" if alt_fit
            else f"{alt['full_name']} — {alt['reason']}",
        )
    console.print(table)

    # Warnings
    if s.warnings:
        console.print()
        for w in s.warnings:
            console.print(f"  [yellow]\u26a0[/yellow] {w}")

    # The command
    console.print("\n[bold]Copy and run:[/bold]\n")
    console.print(f"[green]{s.train_command}[/green]")

    # Tips
    if s.tips:
        console.print("\n[bold]Training tips:[/bold]")
        for tip in s.tips:
            console.print(f"  \u2022 {tip}")

    # Save as script
    script_name = f"train_{s.dataset_repo.split('/')[-1]}_{s.policy_type}.sh"
    console.print(f"\n[dim]Save as script: echo the command above > {script_name}[/dim]")
    console.print()


# ---------------------------------------------------------------------------
# Multi-framework support
# ---------------------------------------------------------------------------


def generate_multi_framework_suggestion(
    dataset_repo: str,
    info,
    episode_data=None,
    policy: str | None = None,
    gpu_memory_gb: int | None = None,
    framework: str = "lerobot",
):
    """Generate training configs for one or all frameworks.

    Args:
        dataset_repo: Dataset path or HuggingFace repo ID.
        info: DatasetInfo from dataset_loader.
        episode_data: Optional DataFrame for deeper analysis.
        policy: Optional policy name.
        gpu_memory_gb: Available GPU memory.
        framework: Framework name or "all" for all frameworks.

    Returns:
        List of (framework_name, TrainingConfig) tuples.
    """
    from orbit.training.base import FRAMEWORKS, get_framework_config

    # First generate the base LeRobot suggestion to get computed hyperparams
    base = generate_suggestion(
        dataset_repo=dataset_repo,
        info=info,
        episode_data=episode_data,
        policy=policy,
        gpu_memory_gb=gpu_memory_gb,
    )

    frameworks = FRAMEWORKS if framework == "all" else [framework]
    results = []

    for fw in frameworks:
        try:
            config = get_framework_config(
                framework=fw,
                dataset_repo=dataset_repo,
                policy_type=base.policy_type,
                info=info,
                episode_count=info.num_episodes,
                action_dim=base.action_dim,
                camera_count=base.camera_count,
                resolution=base.camera_resolution,
                gpu_memory_gb=gpu_memory_gb,
                batch_size=base.batch_size,
                steps=base.steps,
                learning_rate=base.learning_rate,
                chunk_size=base.chunk_size,
                fps=base.fps,
            )
            results.append((fw, config))
        except Exception as e:
            logger.warning("Failed to generate %s config: %s", fw, e)

    return base, results


def print_multi_framework_suggestion(base_suggestion, framework_configs) -> None:
    """Print training configs for multiple frameworks."""
    from rich.panel import Panel
    from rich.table import Table

    from orbit.utils.display import console

    s = base_suggestion

    # Dataset summary
    console.print()
    fit_label = f"fit: {s.fit_score:.2f}" if s.fit_score else "best fit"
    title = f"Dataset Summary — Policy: {s.policy_full_name} ({fit_label})"
    table = Table(show_header=False, border_style="cyan", title=title)
    table.add_column("Metric", style="bold")
    table.add_column("Value")
    table.add_row("Episodes", f"{s.episode_count}")
    table.add_row("Cameras", f"{s.camera_count} x {s.camera_resolution}" if s.camera_count else "none")
    table.add_row("Action dims", str(s.action_dim))
    table.add_row("FPS", f"{s.fps:.0f}")
    console.print(table)

    # Per-framework configs
    for fw_name, config in framework_configs:
        console.print()
        fw_title = fw_name.upper()
        if fw_name == "osmo":
            fw_title = "NVIDIA OSMO"
        elif fw_name == "openpi":
            fw_title = "OpenPI / pi0"
        elif fw_name == "openvla":
            fw_title = "OpenVLA"
        elif fw_name == "groot":
            fw_title = "GR00T N1"

        console.print(f"[bold cyan]{'=' * 60}[/bold cyan]")
        console.print(f"[bold]{fw_title}[/bold]")
        console.print(f"[dim]GPU: {config.gpu_requirements} | Time: {config.estimated_time}[/dim]")
        console.print(f"[dim]Format: {config.dataset_format}[/dim]")

        if config.conversion_needed:
            console.print(f"\n[yellow]Conversion needed:[/yellow]")
            console.print(f"  [green]{config.conversion_command}[/green]")

        if config.pre_checks:
            console.print(f"\n[bold]Pre-checks:[/bold]")
            for check in config.pre_checks:
                console.print(f"  {check}")

        console.print(f"\n[bold]Command:[/bold]")
        if fw_name == "osmo":
            console.print(f"[green]{config.command}[/green]")
        else:
            console.print(f"\n[green]{config.command}[/green]")

        if config.warnings:
            console.print()
            for w in config.warnings:
                if w:
                    console.print(f"  [yellow]\u26a0[/yellow] {w}")

        if config.tips:
            console.print(f"\n[bold]Tips:[/bold]")
            for tip in config.tips:
                console.print(f"  \u2022 {tip}")

        # Generated config files
        if config.generated_files:
            console.print(f"\n[bold]Generated config files:[/bold]")
            for filepath, content in config.generated_files.items():
                console.print(f"\n  [cyan]{filepath}:[/cyan]")
                for line in content.splitlines():
                    console.print(f"  [dim]{line}[/dim]")

    console.print()
