"""OptiQ CLI: convert, benchmark, and evaluate models."""

import os
import sys

import click


@click.group()
@click.version_option(version="0.1.0")
def cli():
    """OptiQ: Optimizing compiler for PyTorch → MLX with data-driven mixed-precision quantization."""
    pass


@cli.command()
@click.argument("model")
@click.option("--target", default="mlx", help="Target backend (mlx)")
@click.option("--model-type", default="auto", type=click.Choice(["auto", "llm"]),
              help="Model type (auto-detected if not specified)")
@click.option("--target-bpw", default=4.5, type=float, help="Target bits per weight (default: 4.5)")
@click.option("--candidate-bits", default="4,8", help="Candidate bit-widths (comma-separated, e.g. 3,4,8)")
@click.option("--group-size", default=64, type=int, help="Quantization group size")
@click.option("--output", "-o", default=None, help="Output directory")
@click.option("--n-calibration", default=8, type=int, help="Number of calibration samples")
@click.option("--skip-baselines", is_flag=True, help="Skip baseline conversions")
def convert(model, target, model_type, target_bpw, candidate_bits,
            group_size, output, n_calibration, skip_baselines):
    """Convert a model to optimized MLX format.

    Uses per-layer KL divergence sensitivity analysis with multi-tier bit
    allocation (2/3/4/8-bit) to minimize size while preserving quality.

    Examples:

        optiq convert Qwen/Qwen3-0.6B-base --target-bpw 4.0

        optiq convert Qwen/Qwen3-0.6B-base --target-bpw 3.5 --candidate-bits 2,3,4,8
    """
    # Parse candidate bits
    bits_list = [int(b.strip()) for b in candidate_bits.split(",")]

    if output is None:
        safe_name = model.replace("/", "_").replace(".", "_")
        output = f"optiq_output/{safe_name}"

    click.echo(f"OptiQ: Converting {model} (type={model_type}) → {target}")
    click.echo(f"  Target BPW: {target_bpw}, candidate bits: {bits_list}, group_size: {group_size}")

    from .models.llm import run_llm_pipeline

    results = run_llm_pipeline(
        model_name=model,
        output_dir=output,
        target_bpw=target_bpw,
        n_calibration=n_calibration,
        candidate_bits=bits_list,
        group_size=group_size,
        skip_baselines=skip_baselines,
    )
    click.echo(f"\nDone! OptiQ model saved to: {results.get('optiq_path', output)}")


@cli.command()
@click.argument("model_path")
@click.option("--baseline", default=None, help="Baseline model path for comparison")
@click.option("--n-samples", default=50, type=int, help="Number of evaluation samples")
def benchmark(model_path, baseline, n_samples):
    """Benchmark a converted model.

    Examples:

        optiq benchmark ./optiq_model --baseline ./uniform_4bit
    """
    from .utils.benchmark import (
        BenchmarkResult,
        measure_llm_perplexity,
        measure_llm_throughput,
        get_model_size,
        measure_model_bpw,
        print_comparison_table,
    )

    results = []

    click.echo(f"Benchmarking LLM: {model_path}")
    size = get_model_size(model_path)
    bpw = measure_model_bpw(model_path)

    click.echo("  Measuring perplexity...")
    ppl = measure_llm_perplexity(model_path, n_samples=n_samples)

    click.echo("  Measuring throughput...")
    tps = measure_llm_throughput(model_path)

    results.append(BenchmarkResult(
        name=os.path.basename(model_path),
        model_size_mb=size,
        bpw=bpw,
        perplexity=ppl,
        tokens_per_sec=tps,
    ))

    if baseline:
        click.echo(f"\nBenchmarking baseline: {baseline}")
        b_size = get_model_size(baseline)
        b_bpw = measure_model_bpw(baseline)

        click.echo("  Measuring perplexity...")
        b_ppl = measure_llm_perplexity(baseline, n_samples=n_samples)

        click.echo("  Measuring throughput...")
        b_tps = measure_llm_throughput(baseline)

        results.append(BenchmarkResult(
            name=os.path.basename(baseline),
            model_size_mb=b_size,
            bpw=b_bpw,
            perplexity=b_ppl,
            tokens_per_sec=b_tps,
        ))

    print_comparison_table(results)


@cli.command(name="eval")
@click.argument("model_path")
@click.option("--task", default="gsm8k", type=click.Choice(["gsm8k", "ai2d", "wer"]),
              help="Evaluation task")
@click.option("--baseline", default=None, help="Baseline model path for comparison")
@click.option("--n-samples", default=200, type=int, help="Number of evaluation samples")
def eval_cmd(model_path, task, baseline, n_samples):
    """Evaluate a model on standard benchmarks.

    Examples:

        optiq eval ./optiq_model --task gsm8k --baseline ./uniform_4bit

        optiq eval ./vlm_model --task ai2d --n-samples 100
    """
    if task == "gsm8k":
        from .eval.gsm8k import evaluate_gsm8k, print_gsm8k_report

        models = [(model_path, os.path.basename(model_path))]
        if baseline:
            models.append((baseline, os.path.basename(baseline)))

        all_results = []
        for path, name in models:
            click.echo(f"\nEvaluating {name} on GSM8K ({n_samples} questions)...")
            result = evaluate_gsm8k(path, n_samples=n_samples)
            print_gsm8k_report(result)
            all_results.append((name, result))

        if len(all_results) > 1:
            click.echo(f"\n  GSM8K Comparison")
            click.echo(f"  {'=' * 50}")
            click.echo(f"  {'Model':<30} {'Accuracy':>10}")
            click.echo(f"  {'-' * 50}")
            for name, r in all_results:
                click.echo(f"  {name:<30} {r.accuracy:>9.1%}")
            click.echo(f"  {'=' * 50}")

    elif task == "ai2d":
        from .eval.ai2d import evaluate_ai2d, print_ai2d_report
        from .core.importer import load_vlm_model
        import torch

        click.echo(f"\nLoading VLM: {model_path}")
        model_info = load_vlm_model(model_path, dtype=torch.float16)

        click.echo(f"\nEvaluating on AI2D ({n_samples} questions)...")
        result = evaluate_ai2d(
            model_info.model, model_info.tokenizer,
            n_samples=n_samples,
        )
        print_ai2d_report(result)

    elif task == "wer":
        from .eval.wer import evaluate_wer_mlx, print_wer_report

        click.echo(f"\nEvaluating WER: {model_path} ({n_samples} samples)...")
        result = evaluate_wer_mlx(model_path, n_samples=n_samples)
        print_wer_report(result)


@cli.command()
@click.argument("model_path")
@click.option("--calibrate", is_flag=True, help="Calibrate by measuring actual throughput")
def latency(model_path, calibrate):
    """Predict decode throughput for a quantized model.

    Uses the Apple Silicon latency model to estimate throughput without
    running inference. With --calibrate, measures actual throughput once
    and adjusts predictions for remaining models.

    Examples:

        optiq latency ./optiq_model

        optiq latency ./optiq_model --calibrate
    """
    import json
    from .core.latency import detect_hardware, CalibratedLatencyModel
    from .utils.benchmark import get_model_size

    hw = detect_hardware()
    click.echo(f"Hardware: {hw.name} ({hw.memory_bandwidth_gbps:.0f} GB/s)")

    size_mb = get_model_size(model_path)
    size_bytes = size_mb * 1024 * 1024

    # Read config to get layer count
    config_path = os.path.join(model_path, "config.json")
    n_layers = 197
    if os.path.exists(config_path):
        with open(config_path) as f:
            config = json.load(f)
        n_layers = config.get("num_hidden_layers", 28) * 7 + 1

    # Raw prediction (linear layers only)
    linear_us = (size_bytes / hw.memory_bandwidth_bps) * 1e6 * 1.15 + 5.0 * n_layers
    raw_tps = 1e6 / linear_us

    click.echo(f"Model: {model_path}")
    click.echo(f"  Size: {size_mb:.1f} MB")
    click.echo(f"  Linear-layer latency: {linear_us:.0f} us")

    if calibrate:
        import time
        import mlx.core as mx
        from mlx_lm import load, generate
        import numpy as np

        click.echo("  Calibrating (measuring actual throughput)...")
        model, tokenizer = load(model_path)
        prompt = "The quick brown fox jumps over the lazy dog."

        for _ in range(8):
            generate(model, tokenizer, prompt=prompt, max_tokens=10, verbose=False)

        times = []
        for _ in range(15):
            mx.synchronize()
            t0 = time.perf_counter()
            generate(model, tokenizer, prompt=prompt, max_tokens=50, verbose=False)
            mx.synchronize()
            times.append(50 / (time.perf_counter() - t0))

        actual_tps = float(np.median(times))
        del model

        cal = CalibratedLatencyModel(hw=hw)
        cal.calibrate(linear_us, actual_tps)

        click.echo(f"  Actual throughput: {actual_tps:.1f} tok/s")
        click.echo(f"  Overhead: {cal.overhead_us:.0f} us ({cal.overhead_us/(1e6/actual_tps)*100:.0f}%)")
        click.echo(f"  Calibrated prediction: {actual_tps:.1f} tok/s")
    else:
        click.echo(f"  Raw prediction: {raw_tps:.0f} tok/s (linear layers only)")
        click.echo(f"  Note: actual throughput will be lower due to attention/norm/framework overhead.")
        click.echo(f"  Use --calibrate for accurate prediction.")


@cli.command(name="kv-cache")
@click.argument("model_path")
@click.option("--target-bits", default=5.0, type=float, help="Target average KV cache bits (default: 5.0)")
@click.option("--candidate-bits", default="4,8", help="Candidate bit-widths (comma-separated)")
@click.option("--n-samples", default=5, type=int, help="Number of calibration samples")
@click.option("--seq-len", default=512, type=int, help="Calibration sequence length")
@click.option("--group-size", default=64, type=int, help="Quantization group size")
@click.option("--output", "-o", default=None, help="Output directory for results")
def kv_cache_cmd(model_path, target_bits, candidate_bits, n_samples, seq_len,
                 group_size, output):
    """Analyze and optimize KV cache quantization.

    Measures per-layer KV cache sensitivity and generates a mixed-precision
    configuration that keeps sensitive layers at higher precision.

    Examples:

        optiq kv-cache Qwen/Qwen3-0.6B-base --target-bits 5.0

        optiq kv-cache ./my_model --target-bits 6.0 --n-samples 10
    """
    import json
    from .core.kv_cache import (
        measure_kv_sensitivity, optimize_kv_cache, KVCacheConfig,
    )

    bits_list = [int(b.strip()) for b in candidate_bits.split(",")]

    if output is None:
        output = "optiq_output/kv_cache"
    os.makedirs(output, exist_ok=True)

    click.echo(f"OptiQ KV Cache Analysis: {model_path}")
    click.echo(f"  Target bits: {target_bits}, candidates: {bits_list}")

    # Step 1: Sensitivity analysis
    click.echo("\nStep 1: Measuring per-layer KV cache sensitivity...")
    results = measure_kv_sensitivity(
        model_path,
        candidate_bits=bits_list,
        n_samples=n_samples,
        seq_len=seq_len,
        group_size=group_size,
    )

    # Save sensitivity results
    sens_path = os.path.join(output, "kv_sensitivity.json")
    data = [
        {
            "layer_idx": r.layer_idx,
            "layer_name": r.layer_name,
            "sensitivities": {str(k): v for k, v in r.sensitivities.items()},
            "n_kv_heads": r.n_kv_heads,
            "head_dim": r.head_dim,
        }
        for r in results
    ]
    with open(sens_path, "w") as f:
        json.dump(data, f, indent=2)
    click.echo(f"  Saved sensitivity data to {sens_path}")

    # Step 2: Optimize
    click.echo(f"\nStep 2: Optimizing KV cache allocation (target: {target_bits} bits)...")
    configs = optimize_kv_cache(
        results,
        target_kv_bits=target_bits,
        candidate_bits=bits_list,
        group_size=group_size,
    )

    # Save config
    config_path = os.path.join(output, "kv_config.json")
    config_data = [{"layer_idx": c.layer_idx, "bits": c.bits, "group_size": c.group_size}
                   for c in configs]
    with open(config_path, "w") as f:
        json.dump(config_data, f, indent=2)
    click.echo(f"  Saved KV config to {config_path}")

    # Summary
    min_bits = min(bits_list)
    max_bits = max(bits_list)
    n_high = sum(1 for c in configs if c.bits == max_bits)
    n_low = sum(1 for c in configs if c.bits == min_bits)
    achieved = sum(c.bits for c in configs) / len(configs)
    click.echo(f"\n  Result: {n_low} layers @ {min_bits}-bit, {n_high} layers @ {max_bits}-bit")
    click.echo(f"  Achieved: {achieved:.2f} average bits")


if __name__ == "__main__":
    cli()
