import os
import unittest

from prain.config import LayerPlacement, Miss, Policy, Prefetch
from prain.sim import (Cost, read_forwards, simulate, synthetic_trace, tune, write_trace)
from tests.helpers import TmpCase

SMALL = dict(layers=4, experts=32, top_k=4, forwards=120, prompt_every=60, prompt_tokens=32)


def forwards_of(records):
    forwards, last = [], None
    for r in records:
        if last is None or r["l"] <= last:
            forwards.append({})
        forwards[-1][r["l"]] = r
        last = r["l"]
    return forwards


def all_ram(slots, layers=4, experts=32):
    return {layer: LayerPlacement((), tuple(range(experts)), (), slots) for layer in range(layers)}


class Traces(TmpCase):
    def test_round_trip_and_determinism(self):
        records = synthetic_trace(**SMALL)
        self.assertEqual(records, synthetic_trace(**SMALL))
        path = os.path.join(self.dir, "t.jsonl.gz")
        write_trace(path, records[:8])
        write_trace(path, records[8:], mode="at")  # runs append to the same trace
        forwards = read_forwards(path)
        self.assertEqual(len(forwards), 120)
        self.assertEqual(forwards[0][0], records[0])
        self.assertEqual(forwards[0][0]["t"], 32)
        self.assertEqual(sorted(forwards[1][1]["p"]), ["1", "2"])
        self.assertNotIn("p", forwards[1][3])  # nothing after the last layer

    def test_copy_cost(self):
        self.assertAlmostEqual(Cost().copy_ms(False), 0.991, places=3)
        self.assertAlmostEqual(Cost().copy_ms(True), 4.955, places=3)


class Simulate(unittest.TestCase):
    def setUp(self):
        self.forwards = forwards_of(synthetic_trace(**SMALL))

    def run_sim(self, policy=Policy(), prefetch=Prefetch(), miss=Miss(), slots=3, **kw):
        return simulate(self.forwards, all_ram(slots), policy, prefetch, miss, Cost(), **kw)

    def test_belady_bounds_every_policy(self):
        bound = self.run_sim(belady=True).hit_rate
        for eviction in ("lru", "lfu", "gdsf"):
            self.assertLessEqual(self.run_sim(Policy(eviction=eviction)).hit_rate, bound)

    def test_cpu_path_beats_copies_when_the_cpu_is_faster_than_pcie(self):
        copy = self.run_sim().ms_per_decode
        self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy)

    def test_prefetch_helps_when_predictions_are_right(self):
        none = self.run_sim().hit_rate
        good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))
        self.assertGreater(good.hit_rate, none)
        self.assertGreater(good.counters["prefetch_hits"], good.counters["prefetch_wasted"])

    def test_pregate_needs_predictions(self):
        stripped = [{layer: {k: v for k, v in r.items() if k != "p"} for layer, r in f.items()}
                    for f in self.forwards]
        with self.assertRaisesRegex(ValueError, "needs a trace with predictions"):
            simulate(stripped, all_ram(3), Policy(), Prefetch(mode="pregate"), Miss(), Cost())


class Tune(unittest.TestCase):
    def test_tune_beats_the_defaults_within_the_budget(self):
        forwards = forwards_of(synthetic_trace(**SMALL))
        cost = Cost(expert_mib=64.0)  # 1 GiB = 16 slots = 4 per layer
        best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost)
        self.assertLessEqual(sum(per.values()), 16)
        self.assertEqual(sum(len(r) for r in ram.values()), 16)  # 1 GiB of RAM tier
        default = next(r for v, r in rows if v == (Policy(), Prefetch(), Miss()))
        self.assertLessEqual(final.total_ms, default.total_ms)
        self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows))
        self.assertGreaterEqual(bound.hit_rate, 0.0)