"""Tests for chip_roofline: AI, B*, L*, per-token latency curves."""
from __future__ import annotations

import math

import pytest

from tests.analytical_visualization.chip_roofline import (
    arithmetic_intensity,
    balance_context,
    bound_regime,
    critical_batch,
    good_batch,
    good_context,
    knee_batch,
    kv_bytes_per_token,
    per_token_latency_curve,
    step_latency_curve,
    t_com,
    t_mem_long,
    t_mem_short,
    total_active_params,
    utilization_at,
)
from tests.analytical_visualization.model_config import (
    MachineParams, ModelConfig,
)
from tests.analytical_visualization.model_presets import PRESETS


# ── Arithmetic intensity ────────────────────────────────────────────


def test_arithmetic_intensity_matches_ratio():
    """AI = peak_flops / bw_hbm — pure ratio, no util factor."""
    m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
    assert arithmetic_intensity(m) == pytest.approx(8e12 / 256e9)


def test_ai_scales_with_flops_and_bandwidth():
    m1 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
    m2 = MachineParams(peak_tflops_f16=16.0, bw_hbm_gbs=256.0)
    m3 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=128.0)
    assert arithmetic_intensity(m2) == pytest.approx(2 * arithmetic_intensity(m1))
    assert arithmetic_intensity(m3) == pytest.approx(2 * arithmetic_intensity(m1))


# ── Critical batch B* ───────────────────────────────────────────────


def test_critical_batch_h100_reference():
    """H100-class: 989 TFLOPs BF16, 3.35 TB/s HBM3 → B* ≈ 295."""
    m = MachineParams(peak_tflops_f16=989.0, bw_hbm_gbs=3350.0)
    model = PRESETS["Llama 3 8B"].model  # bf16 (b=2)
    b_star = critical_batch(m, model)
    assert b_star == pytest.approx(295, rel=0.05), b_star


def test_critical_batch_scales_with_sparsity():
    """MoE 8× sparsity gives 8× B*."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    b_dense = critical_batch(m, model, sparsity=1)
    b_moe8 = critical_batch(m, model, sparsity=8)
    assert b_moe8 == pytest.approx(8 * b_dense)


def test_critical_batch_bf16_formula():
    """For BF16 (b=2), B* = AI in flops-per-byte units → numerically."""
    m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
    model = ModelConfig(bytes_per_elem=2)
    ai = arithmetic_intensity(m)
    assert critical_batch(m, model) == pytest.approx(ai), (
        critical_batch(m, model), ai,
    )


# ── Balance context L* ──────────────────────────────────────────────


def test_balance_context_positive_finite():
    """L* for a real model on real machine is a positive finite number."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    l_star = balance_context(m, model)
    assert 0 < l_star < 1e9, l_star


def test_balance_context_scales_with_bandwidth():
    """L* = 2N*W/(C*kv_bpt): doubling HBM BW doubles L*.
    Slower memory hits the KV wall at a shorter context. Machine-only
    change so N_active and kv_bpt stay fixed."""
    fast = MachineParams(bw_hbm_gbs=512.0)
    slow = MachineParams(bw_hbm_gbs=256.0)
    model = PRESETS["Llama 3 8B"].model
    assert (balance_context(fast, model)
            == pytest.approx(2 * balance_context(slow, model)))


# ── Knee ────────────────────────────────────────────────────────────


def test_knee_equals_b_star_at_short_context():
    """S_kv → 0: B_knee → B*."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    b_star = critical_batch(m, model)
    assert knee_batch(m, model, s_kv=1) == pytest.approx(b_star, rel=1e-3)


def test_knee_diverges_at_balance_context():
    """S_kv = L*: knee is infinite; S_kv > L*: no knee (None)."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    l_star = balance_context(m, model)
    assert knee_batch(m, model, s_kv=int(2 * l_star)) is None
    # At 0.9 * L*, knee ~= 10 * B*; assert at least 5x to survive rounding.
    below = knee_batch(m, model, s_kv=int(0.9 * l_star))
    assert below is not None and below > 5 * critical_batch(m, model)


def test_knee_slides_right_with_context():
    """S_kv = L*/2 → B_knee = 2 * B*."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    b_star = critical_batch(m, model)
    l_star = balance_context(m, model)
    got = knee_batch(m, model, s_kv=int(l_star / 2))
    assert got == pytest.approx(2 * b_star, rel=0.01), (got, 2 * b_star)


# ── Per-token latency curve ─────────────────────────────────────────


def test_latency_curve_monotonically_decreasing():
    """total_s must strictly decrease as batch increases (weight/B shrinks)."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    pts = per_token_latency_curve(m, model, [1, 2, 4, 8, 16, 32, 64],
                                    s_kv=1024)
    totals = [p.total_s for p in pts]
    for a, b in zip(totals, totals[1:]):
        assert b < a, (a, b)


def test_latency_curve_asymptotes_to_compute_plus_kv():
    """At very large B, total → compute + KV (weight/B → 0)."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    pts = per_token_latency_curve(m, model, [10_000_000], s_kv=1024)
    p = pts[0]
    assert p.total_s == pytest.approx(p.compute_s + p.kv_s, rel=1e-3)


def test_latency_at_b_star_is_roughly_2x_floor():
    """At B*, weight = compute (dense, S_kv small), so total ≈ 2 * compute."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    b_star = int(round(critical_batch(m, model)))
    pts = per_token_latency_curve(m, model, [b_star], s_kv=1)
    p = pts[0]
    # weight_s ≈ compute_s at B*
    assert p.weight_s == pytest.approx(p.compute_s, rel=0.05), (
        p.weight_s, p.compute_s,
    )


# ── Regime classification ───────────────────────────────────────────


def test_regime_memory_bound_at_b1():
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    assert bound_regime(m, model, batch=1, s_kv=1024) == "memory-bound"


def test_regime_kv_bound_past_balance_context():
    """S_kv > L* with reasonable batch → KV dominates."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    l_star = balance_context(m, model)
    b_star = int(round(critical_batch(m, model)))
    assert bound_regime(m, model, batch=b_star * 4,
                         s_kv=int(3 * l_star)) == "kv-bound"


# ── Component sanity ────────────────────────────────────────────────


def test_total_active_params_llama3_8b_ballpark():
    """Llama 3 8B — attn+FFN alone (no embeddings/LM head) is ~6.9B,
    HF-reported total 8.03B. Guard the ballpark."""
    n = total_active_params(PRESETS["Llama 3 8B"].model)
    assert 6.5e9 < n < 8e9, n


def test_kv_bytes_per_token_llama3_8b():
    """Llama 3 8B: 2*8*128*2 bytes/layer/token × 32 layers = 131072 bytes."""
    kv_bpt = kv_bytes_per_token(PRESETS["Llama 3 8B"].model)
    assert kv_bpt == 2 * 8 * 128 * 2 * 32


# ── Regime terms + good-batch / good-context recommendations ──────


def test_t_mem_short_shrinks_with_batch():
    """t_mem_short = N·b/(W·B) — doubling B halves it."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    assert t_mem_short(m, model, 2) == pytest.approx(
        t_mem_short(m, model, 1) / 2
    )


def test_t_mem_long_equals_t_com_at_l_star():
    """L* is defined by t_mem_long(L*) == t_com. Cross-check."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    l_star = balance_context(m, model)
    assert t_mem_long(m, model, int(round(l_star))) == pytest.approx(
        t_com(m, model), rel=0.01,
    )


def test_t_com_is_batch_independent():
    """t_com = 2·N/C. Constant — doesn't depend on B or S_kv."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    tc = t_com(m, model)
    assert tc == pytest.approx(2 * total_active_params(model) / m.peak_flops)


def test_good_batch_is_two_b_star():
    """Recommendation is 2 × B* (Pope's rule)."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    rec = good_batch(m, model)
    assert rec.effective == pytest.approx(2 * critical_batch(m, model))
    assert rec.b_star == pytest.approx(critical_batch(m, model))


def test_good_batch_moe_scales_with_sparsity():
    """MoE sparsity=8 → recommended B is 8× dense recommendation."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    r_dense = good_batch(m, model, sparsity=1.0)
    r_moe = good_batch(m, model, sparsity=8.0)
    assert r_moe.effective == pytest.approx(8 * r_dense.effective)


def test_good_context_returns_l_star_as_ceiling():
    """Recommendation is L*; utilization drops past it."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    rec = good_context(m, model, s_kv=100)
    assert rec.l_star == pytest.approx(balance_context(m, model))
    assert rec.max_efficient == pytest.approx(rec.l_star)


def test_utilization_at_l_star_is_half():
    """At S_kv = L*, compute and KV read take equal time → 50% util."""
    l_star = 3000.0
    assert utilization_at(int(l_star), l_star) == pytest.approx(0.5)


def test_utilization_at_2_l_star_is_one_third():
    """At S_kv = 2·L*, util = 1/(1+2) = 33.3%."""
    l_star = 3000.0
    assert utilization_at(int(2 * l_star), l_star) == pytest.approx(1 / 3, rel=1e-3)


def test_utilization_at_zero_context_is_100_percent():
    """Zero context, no KV to read → all time is compute."""
    assert utilization_at(0, 3000.0) == pytest.approx(1.0)


def test_utilization_drops_monotonically_with_context():
    """util(S_kv) is strictly decreasing."""
    l_star = 3000.0
    vals = [utilization_at(s, l_star) for s in [100, 1000, 3000, 6000, 30000]]
    for a, b in zip(vals, vals[1:]):
        assert b < a


# ── Step latency (undivided by B) ──────────────────────────────────


def test_step_weight_is_batch_invariant():
    """step_weight = N·b/W — same for every B (loaded once per step)."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    pts = step_latency_curve(m, model, [1, 8, 64, 512], s_kv=1024)
    ws = [p.weight_s for p in pts]
    assert all(w == pytest.approx(ws[0]) for w in ws), ws


def test_step_compute_linear_in_batch():
    """step_compute = 2·N·B/C — doubling B doubles compute."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    pts = step_latency_curve(m, model, [1, 2, 8], s_kv=1024)
    assert pts[1].compute_s == pytest.approx(2 * pts[0].compute_s)
    assert pts[2].compute_s == pytest.approx(8 * pts[0].compute_s)


def test_step_kv_linear_in_batch():
    """step_kv = B · S_kv · kv_bpt / W — linear in B."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    pts = step_latency_curve(m, model, [1, 4], s_kv=1024)
    assert pts[1].kv_s == pytest.approx(4 * pts[0].kv_s)


def test_step_total_equals_per_token_times_batch():
    """The step and per-token views are just ÷ B / × B of each other."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    B_range = [1, 4, 16, 64]
    s_kv = 1024
    step_pts = step_latency_curve(m, model, B_range, s_kv=s_kv)
    tok_pts = per_token_latency_curve(m, model, B_range, s_kv=s_kv)
    for sp, tp in zip(step_pts, tok_pts):
        assert sp.total_s == pytest.approx(tp.total_s * sp.batch)


def test_step_and_per_token_identical_at_b1():
    """At B=1, step latency == per-token cost (dividing by 1)."""
    m = MachineParams()
    model = PRESETS["Llama 3 8B"].model
    sp = step_latency_curve(m, model, [1], s_kv=1024)[0]
    tp = per_token_latency_curve(m, model, [1], s_kv=1024)[0]
    assert sp.total_s == pytest.approx(tp.total_s)
    assert sp.weight_s == pytest.approx(tp.weight_s)
    assert sp.compute_s == pytest.approx(tp.compute_s)
    assert sp.kv_s == pytest.approx(tp.kv_s)


# ── App wiring: tab exists on the Streamlit app ────────────────────


def test_roofline_tab_registered_in_app():
    """The new 'Chip roofline & B*' tab is listed in st.tabs and gets a
    corresponding `with tab_roofline:` block."""
    from pathlib import Path
    src = (Path(__file__).parent / "app.py").resolve().read_text(
        encoding="utf-8"
    )
    assert src.count('"Chip roofline & B*"') == 1
    assert "with tab_roofline:" in src