import math

import numpy as np
import pytest

import tiled_roofline as tr


@pytest.mark.parametrize("shape, b", [((8, 8, 8), 2), ((37, 23, 41), 8), ((1, 33, 17), 8), ((7, 5, 3), 1), ((9, 11, 13), 64)])
def test_tiled_matmul_is_correct_when_sizes_do_not_divide(shape, b):
    """Guards an off-by-one on edge tiles: a wrong slice end gives a wrong C, not an error."""
    m, k, n = shape
    rng = np.random.default_rng(0)
    A, B = rng.standard_normal((m, k)), rng.standard_normal((k, n))
    C, _ = tr.tiled_matmul(A, B, b)
    assert np.allclose(C, A @ B, atol=1e-9)
    assert np.allclose(C, tr.naive_matmul(A.tolist(), B.tolist()), atol=1e-9)


@pytest.mark.parametrize("shape, b", [((64, 64, 64), 8), ((37, 23, 41), 8), ((1, 33, 17), 8), ((9, 11, 13), 64), ((5, 3, 7), 1)])
def test_counted_loop_matches_closed_form(shape, b):
    """The formula in the lesson and the loop must agree, ragged edges included.
    Guards a closed form that drifts from the loop it claims to describe."""
    m, k, n = shape
    A, B = np.ones((m, k)), np.ones((k, n))
    _, counted = tr.tiled_matmul(A, B, b)
    assert counted == tr.tiled_traffic(m, n, k, b)
    assert counted.flops == tr.matmul_flops(m, n, k)


def test_square_formula_and_tile_one_is_the_naive_count():
    """Loads = 2 n^3 / b for square sizes; tile 1 is exactly the no-reuse count."""
    n, b = 64, 16
    assert tr.tiled_traffic(n, n, n, b).loads == 2 * n**3 // b
    assert tr.tiled_traffic(n, n, n, 1) == tr.naive_traffic(n, n, n)
    assert tr.tiled_traffic(n, n, n, n) == tr.min_traffic(n, n, n)    # one tile: every element moves once


def test_worked_example_four_by_four_with_two_by_two_tiles():
    """The numbers the lesson prints by hand: 8 steps of 8 loads, 16 stores, 128 FLOPs."""
    C, t = tr.tiled_matmul(tr.EXAMPLE_A, tr.EXAMPLE_B, 2)
    assert np.array_equal(C, tr.EXAMPLE_A @ tr.EXAMPLE_B)
    assert C.tolist() == [[2, 6, 5, 1], [7, 9, 7, 5], [4, 6, 3, 1], [5, 7, 5, 3]]    # the matrix printed in the lesson
    steps = tr.tile_steps(4, 4, 4, 2)
    assert len(steps) == 8 and {s[3] for s in steps} == {8}
    assert sum(s[3] for s in tr.tile_steps(5, 3, 7, 2)) == tr.tiled_traffic(5, 3, 7, 2).loads    # ragged edges: guards a wrong edge-tile size in the per-step loads
    assert (t.loads, t.stores, t.flops) == (64, 16, 128)
    assert tr.naive_traffic(4, 4, 4).loads == 128
    assert tr.naive_traffic(4, 4, 4).intensity(1) == pytest.approx(0.89, abs=0.005)    # per number moved: 128 / 144
    assert t.intensity(1) == pytest.approx(1.6)                                          # 128 / 80
    # top-left output tile, the two partial products the lesson tabulates
    a, b = tr.EXAMPLE_A, tr.EXAMPLE_B
    step0 = a[0:2, 0:2] @ b[0:2, 0:2]
    step1 = a[0:2, 2:4] @ b[2:4, 0:2]
    assert step0.tolist() == [[1, 4], [3, 7]]
    assert step1.tolist() == [[1, 2], [4, 2]]
    assert (step0 + step1).tolist() == [[2, 6], [7, 9]]


def test_a100_ridge_from_the_datasheet_numbers():
    """312 TFLOP/s over 2,039 GB/s. Guards a unit slip (TFLOP vs GFLOP, GB vs GiB)."""
    assert tr.ridge(tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH) == pytest.approx(153.0, abs=0.05)


def test_roofline_min_and_classification_tie():
    peak, bw = 100.0, 10.0               # ridge = 10
    assert tr.attainable(2, peak, bw) == 20 and tr.attainable(50, peak, bw) == 100
    assert tr.classify(9.99, peak, bw) == "memory-bound"
    assert tr.classify(10.0, peak, bw) == "compute-bound"    # a tie sits on the roof


def test_intensity_rises_with_tile_and_is_about_b_over_element_size():
    n, s = 4096, tr.FP16_BYTES
    values = [tr.tiled_intensity(n, b, s) for b in (1, 4, 16, 64, 256)]
    assert values == sorted(values) and len(set(values)) == len(values)
    assert tr.tiled_intensity(n, 1, s) == pytest.approx(0.5, rel=1e-3)
    assert tr.tiled_intensity(n, 64, s) == pytest.approx(64 / s, rel=0.02)
    # one tile covering everything moves each element once: the best case
    assert tr.tiled_intensity(n, n, s) == pytest.approx(tr.gemm_intensity(n, n, n, s))


@pytest.mark.parametrize("kib, expected_b, expected_intensity", [(16, 52, 25.8), (64, 104, 50.6), (256, 209, 99.9), (1024, 418, 195.0)])
def test_tile_capacity_bound(kib, expected_b, expected_intensity):
    """Three b x b tiles must fit. Guards forgetting the C tile (2 b^2 gives 128 at 64 KiB) and a rounded-up
    root (105). The rows are the capacity table in the lesson."""
    s = tr.FP16_BYTES
    b = tr.max_tile(kib * 1024, s)
    assert tr.fast_bytes_needed(b, s) <= kib * 1024 < tr.fast_bytes_needed(b + 1, s)
    assert b == expected_b
    assert tr.tiled_intensity(4096, b, s) == pytest.approx(expected_intensity, abs=0.05)


def test_check_yourself_capacity_answer():
    """Check-yourself question 2: 48 KiB, 4-byte numbers -> b = 64, 15.9 FLOP/byte at n = 4096."""
    b = tr.max_tile(48 * 1024, 4)
    assert b == 64
    assert tr.tiled_intensity(4096, b, 4) == pytest.approx(15.9, abs=0.05)


@pytest.mark.parametrize("n, expected", [(2, 1 / 3), (16, 2.667), (256, 42.667), (4096, 682.667)])
def test_best_case_intensity_table_for_four_byte_floats(n, expected):
    """The lesson's best-case table, 2n / (3s) with s = 4. Guards a slip in the byte count or the 3n^2 numbers."""
    assert tr.gemm_intensity(n, n, n, 4) == pytest.approx(expected, abs=0.0005)


def test_smallest_tile_reaching_the_ridge_is_above_the_large_n_rule_for_two_causes():
    """Lesson quotes 342, not the rule of thumb 306. Guards quoting s * ridge as exact. Two causes:
    writing C alone moves 306 to 318; the whole-tile count (13 tiles per side give 151.7, 12 give 163.8)
    moves it to 342."""
    s, peak, bw = tr.FP16_BYTES, tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    ridge = tr.ridge(peak, bw)
    assert round(s * ridge) == 306
    assert math.ceil(tr.ridge_tile_if_tile_count_were_fractional(4096, s, peak, bw)) == 318
    b = tr.smallest_tile_for_ridge(4096, s, peak, bw)
    assert b == 342 and b > 318
    assert tr.tiled_intensity(4096, b, s) >= ridge > tr.tiled_intensity(4096, b - 1, s)
    assert tr.tiled_intensity(4096, 342, s) == pytest.approx(163.8, abs=0.05)    # 12 tiles per side
    assert tr.tiled_intensity(4096, 341, s) == pytest.approx(151.7, abs=0.05)    # 13 tiles per side
    assert tr.fast_bytes_needed(b, s) == 701_784


def test_decode_is_memory_bound_and_batching_crosses_the_ridge():
    s, peak, bw = tr.FP16_BYTES, tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    assert tr.gemm_intensity(1, 4096, 4096, s) == pytest.approx(1.0, abs=0.001)
    assert tr.classify(tr.gemm_intensity(1, 4096, 4096, s), peak, bw) == "memory-bound"
    assert tr.classify(tr.gemm_intensity(4096, 4096, 4096, s), peak, bw) == "compute-bound"
    m = tr.batch_to_reach_ridge(4096, 4096, s, peak, bw)
    assert m == 166
    assert tr.gemm_intensity(m - 1, 4096, 4096, s) < tr.ridge(peak, bw) <= tr.gemm_intensity(m, 4096, 4096, s)


@pytest.mark.parametrize("m, intensity, tflops, share", [
    (1, 1.00, 2.04, 0.7), (8, 7.97, 16.25, 5.2), (32, 31.51, 64.24, 20.6), (128, 120.47, 245.64, 78.7),
    (165, 152.70, 311.35, 99.8), (166, 153.55, 312.00, 100.0), (1024, 682.67, 312.00, 100.0)])
def test_decode_table_rows_match_the_lesson(m, intensity, tflops, share):
    """Guards the decode/prefill table drifting from the code that prints it."""
    peak, bw = tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    i = tr.gemm_intensity(m, 4096, 4096, tr.FP16_BYTES)
    att = tr.attainable(i, peak, bw)
    assert i == pytest.approx(intensity, abs=0.005)
    assert att / 1e12 == pytest.approx(tflops, abs=0.005)
    assert 100 * att / peak == pytest.approx(share, abs=0.05)


def test_decode_ceiling_for_a_seven_billion_parameter_model():
    """6.87 ms to read 14 GB of FP16 weights at 2,039 GB/s -> at most 146 tokens/s. Guards a bytes/params slip."""
    s, bw = tr.FP16_BYTES, tr.A100_BANDWIDTH
    assert tr.weight_read_seconds(7e9, s, bw) * 1e3 == pytest.approx(6.866, abs=0.001)
    assert tr.decode_tokens_per_s(7e9, s, bw) == pytest.approx(145.6, abs=0.1)


def test_naive_multiply_floors_on_the_a100():
    """The naive 4096 table: 1.019 TFLOP/s (0.327% of peak), floors 0.4405 ms and 134.83 ms (306 times),
    0.0494 ms at the minimum traffic. Guards time_floors and the bytes-per-element slip."""
    n, s, peak, bw = 4096, tr.FP16_BYTES, tr.A100_PEAK_FLOPS, tr.A100_BANDWIDTH
    naive, best = tr.naive_traffic(n, n, n), tr.min_traffic(n, n, n)
    assert tr.attainable(naive.intensity(s), peak, bw) == pytest.approx(1.019e12, abs=5e8)
    assert tr.attainable(naive.intensity(s), peak, bw) / peak * 100 == pytest.approx(0.327, abs=0.0005)
    t_math, t_naive = tr.time_floors(naive.flops, naive.nbytes(s), peak, bw)
    _, t_min = tr.time_floors(best.flops, best.nbytes(s), peak, bw)
    assert t_math * 1e3 == pytest.approx(0.4405, abs=0.00005)
    assert t_naive * 1e3 == pytest.approx(134.83, abs=0.005)
    assert t_naive / t_math == pytest.approx(306.07, abs=0.005)
    assert t_min * 1e3 == pytest.approx(0.0494, abs=0.00005) and t_min < t_math
    assert best.nbytes(s) == 100_663_296


def test_timing_helper_runs_and_reports_positive_rates():
    """Smoke test only: the reader runs this on their own machine, so no timing is asserted."""
    rates = tr.time_on_this_machine(n_small=8, n_mm=16, n_mv=32)
    assert len(rates) == 5 and all(v > 0 for v in rates.values())