profile.py

"""soup profile — estimate memory, speed, and GPU requirements before training."""

from __future__ import annotations

import json
from pathlib import Path

import typer
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

from soup_cli.utils.gpu import model_size_from_name
from soup_cli.utils.profiler import (
    GPU_MEMORY,
    estimate_speed,
    estimate_total,
    recommend_batch_size,
    recommend_gpu,
)

console = Console()


def profile(
    config: str = typer.Option(
        "soup.yaml", "--config", "-c", help="Path to soup.yaml config file"
    ),
    gpu: str = typer.Option(
        None, "--gpu", "-g",
        help=(
            "Target GPU for recommendations "
            "(e.g., rtx3090, rtx4090, a100, h100). Auto-detects if not set."
        ),
    ),
    json_output: bool = typer.Option(
        False, "--json", help="Output as JSON for scripting"
    ),
):
    """Estimate memory, speed, and GPU requirements BEFORE training."""
    from soup_cli.config.loader import load_config

    config_path = Path(config)
    if not config_path.exists():
        console.print(f"[red]Config file not found:[/] {config}")
        raise typer.Exit(1)

    cfg = load_config(config_path)

    # Determine model size
    model_params_b = model_size_from_name(cfg.base)

    # Determine batch size (use 4 as default estimate for "auto")
    batch_size = cfg.training.batch_size
    if batch_size == "auto":
        batch_size = 4
    else:
        batch_size = int(batch_size)

    # Resolve GPU memory
    gpu_memory_gb = _resolve_gpu_memory(gpu)

    # Compute profile
    result = estimate_total(
        model_name=cfg.base,
        model_params_b=model_params_b,
        quantization=cfg.training.quantization,
        lora_r=cfg.training.lora.r,
        lora_alpha=cfg.training.lora.alpha,
        batch_size=batch_size,
        seq_len=cfg.data.max_length,
        optimizer=cfg.training.optimizer,
        gradient_checkpointing=cfg.training.gradient_checkpointing,
    )

    # Speed estimates
    tokens_per_sec = estimate_speed(
        model_params_b, cfg.training.quantization, batch_size
    )
    samples_per_sec = tokens_per_sec / max(cfg.data.max_length, 1)

    # Batch size recommendation
    recommended_bs = recommend_batch_size(result["total_memory_gb"], gpu_memory_gb)

    # GPU recommendations
    compatible_gpus = recommend_gpu(result["total_memory_gb"])

    # Add speed/time to result
    result["tokens_per_sec"] = round(tokens_per_sec, 1)
    result["samples_per_sec"] = round(samples_per_sec, 2)
    result["recommended_batch_size"] = recommended_bs
    result["compatible_gpus"] = compatible_gpus
    result["gpu_memory_gb"] = gpu_memory_gb

    if json_output:
        console.print(json.dumps(result, indent=2))
        return

    _render_profile(result, cfg, gpu_memory_gb)


def _resolve_gpu_memory(gpu: str | None) -> float:
    """Resolve GPU memory in GB from flag or auto-detection."""
    if gpu is not None:
        gpu_key = gpu.lower().replace(" ", "").replace("-", "")
        if gpu_key not in GPU_MEMORY:
            valid = ", ".join(sorted(GPU_MEMORY.keys()))
            console.print(
                f"[red]Unknown GPU:[/] {gpu}\n"
                f"[dim]Valid options: {valid}[/]"
            )
            raise typer.Exit(1)
        return float(GPU_MEMORY[gpu_key])

    # Auto-detect
    try:
        from soup_cli.utils.gpu import get_gpu_info

        info = get_gpu_info()
        mem_bytes = info.get("memory_total_bytes", 0)
        if mem_bytes > 0:
            return mem_bytes / (1024**3)
    except (ImportError, RuntimeError, OSError):
        pass

    # Default to 24 GB (common consumer GPU)
    return 24.0


def _render_profile(result: dict, cfg, gpu_memory_gb: float) -> None:
    """Render Rich profile output."""
    # Model info
    model_info = (
        f"Model:     [bold]{cfg.base}[/]\n"
        f"Params:    [bold]{result['model_params_b']:.1f}B[/] "
        f"(trainable: {result['trainable_params']:,.0f} with LoRA r={cfg.training.lora.r})\n"
        f"Quantization: [bold]{result['quantization']}[/]"
    )
    if result["gradient_checkpointing"]:
        model_info += "\nGradient checkpointing: [bold green]enabled[/]"

    # Memory breakdown table
    mem_table = Table(show_header=False, box=None, padding=(0, 2))
    mem_table.add_column("Component", style="bold")
    mem_table.add_column("Memory", justify="right")
    mem_table.add_row("Model", f"~{result['model_memory_gb']:.1f} GB")
    mem_table.add_row("LoRA", f"~{result['lora_memory_gb']:.1f} GB")
    mem_table.add_row("Optimizer", f"~{result['optimizer_memory_gb']:.1f} GB")
    mem_table.add_row(
        f"Activations (bs={result['batch_size']}, seq={result['seq_len']})",
        f"~{result['activation_memory_gb']:.1f} GB",
    )
    mem_table.add_row("Overhead", f"~{result['overhead_gb']:.1f} GB")
    mem_table.add_row("-" * 20, "-" * 10)
    mem_table.add_row("[bold]Total[/]", f"[bold]~{result['total_memory_gb']:.1f} GB[/]")

    # Speed info
    speed_info = (
        f"Tokens/sec: ~{result['tokens_per_sec']:,.0f}\n"
        f"Samples/sec: ~{result['samples_per_sec']:.1f}"
    )

    # Recommendations
    recs = []
    fits = result["total_memory_gb"] <= gpu_memory_gb
    if fits:
        recs.append(
            f"[green]OK[/] Fits in {gpu_memory_gb:.0f} GB VRAM"
        )
    else:
        recs.append(
            f"[red]X[/] Does NOT fit in {gpu_memory_gb:.0f} GB VRAM "
            f"(need ~{result['total_memory_gb']:.0f} GB)"
        )

    recs.append(
        f"[green]OK[/] Recommended batch_size: {result['recommended_batch_size']}"
    )

    if result["total_memory_gb"] > 24 and not result["gradient_checkpointing"]:
        recs.append(
            "[yellow]![/] Consider gradient_checkpointing: true for memory savings"
        )

    if result["total_memory_gb"] > 40:
        recs.append(
            "[yellow]![/] Consider DeepSpeed ZeRO-3 or FSDP for distributed training"
        )

    # Compatible GPUs (show top 5)
    gpu_list = result["compatible_gpus"][:5]

    console.print(Panel(model_info, title="[bold]Training Profile[/]"))
    console.print()
    console.print("[bold]GPU Memory Estimate:[/]")
    console.print(mem_table)
    console.print()
    console.print(Panel(speed_info, title="Speed Estimate"))
    console.print()
    console.print(Panel("\n".join(recs), title="Recommendations"))
    console.print()

    if gpu_list:
        console.print("[bold]Compatible GPUs:[/]")
        for gpu_name in gpu_list:
            console.print(f"  - {gpu_name}")