import json
import math
import os
import random
import unittest

from prain.config import Miss, Policy
from prain.policy import SlotCache, write_stats
from tests.helpers import TmpCase


def lru(**kw):
    return Policy(eviction="lru", **kw)


class Eviction(unittest.TestCase):
    def test_pinned_experts_stay_and_lru_evicts_the_oldest(self):
        cache = SlotCache(pinned=[3, 5], num_experts=8, cache_slots=2, policy=lru())
        self.assertEqual(cache.plan({0: 1, 3: 1}), [({0: 2, 3: 0}, [(0, 2)])])
        self.assertEqual(cache.plan({0: 1, 1: 1}), [({0: 2, 1: 3}, [(1, 3)])])
        self.assertEqual(cache.plan({1: 1, 2: 1}), [({1: 3, 2: 2}, [(2, 2)])])  # 0 was oldest
        self.assertEqual((cache.hits, cache.misses), (3, 3))
        self.assertEqual(cache.uses, [2, 2, 1, 1, 0, 0, 0, 0])

    def test_lfu_keeps_the_frequent_expert_lru_the_recent_one(self):
        for policy, survivor in ((lru(), 1), (Policy(eviction="lfu"), 0)):
            cache = SlotCache([], 3, 2, policy)
            for e in (0, 0, 0, 1, 2):
                cache.plan({e: 1})
            self.assertIn(survivor, cache.slots, policy.eviction)

    def test_gdsf_keeps_the_expensive_disk_expert(self):
        cost = [1.0, 4.0, 1.0]  # expert 1 lives on disk
        for eviction, survivor in (("lfu", 0), ("gdsf", 1)):
            cache = SlotCache([], 3, 2, Policy(eviction=eviction), cost)
            for e in (0, 0, 1, 2):
                cache.plan({e: 1})
            self.assertIn(survivor, cache.slots, eviction)


class Rounds(unittest.TestCase):
    def test_hottest_experts_are_copied_last_and_stay(self):
        cache = SlotCache([], 10, 2, lru())
        rounds = cache.plan({0: 1, 1: 9, 2: 3, 3: 8, 4: 2, 5: 1})
        self.assertEqual([sorted(slot_of) for slot_of, _ in rounds], [[0, 5], [2, 4], [1, 3]])
        self.assertEqual(set(cache.slots), {1, 3})

    def test_scan_keeps_the_protected_expert_and_promotes_the_hottest(self):
        cache = SlotCache([], 6, 3, lru(scan_tokens=4, scan_keep=1, promote=1))
        cache.plan({0: 1})
        cache.plan({1: 1})
        rounds = cache.plan({2: 1, 3: 5, 4: 2, 5: 3}, tokens=8)
        self.assertEqual(rounds, [({2: 2, 4: 0}, [(2, 2), (4, 0)]),
                                  ({5: 2, 3: 0}, [(5, 2), (3, 0)])])
        self.assertEqual(set(cache.slots), {1, 3, 5})  # 1 was protected
        self.assertEqual(cache.evictor.last[3], cache.now)  # promoted
        self.assertEqual(cache.evictor.last[5], -math.inf)  # admitted at low priority

    def test_cpu_experts_get_no_slot(self):
        cache = SlotCache([0], 6, 2, lru())
        cache.plan({1: 1})
        counts = {0: 1, 1: 1, 2: 1, 3: 4}
        self.assertEqual(cache.choose_cpu(counts, Miss()), set())
        self.assertEqual(cache.choose_cpu(counts, Miss(mode="cpu")), {2, 3})
        cpu = cache.choose_cpu(counts, Miss(mode="auto", cpu_max_tokens=2))
        self.assertEqual(cpu, {2})
        self.assertEqual(cache.plan(counts, 4, cpu), [({0: 0, 1: 1, 3: 2}, [(3, 2)])])
        self.assertEqual((cache.cpu, cache.misses), (1, 2))

    def test_random_plans_always_compute_every_expert_from_the_right_slot(self):
        rng = random.Random(1)
        for _ in range(300):
            n = rng.randint(2, 12)
            pinned = rng.sample(range(n), rng.randint(0, n // 2))
            slots = rng.randint(1, 4)
            policy = Policy(eviction=rng.choice(["lru", "lfu", "gdsf"]),
                            scan_tokens=rng.randint(1, 6), scan_keep=rng.randint(0, slots - 1),
                            promote=rng.randint(0, 3))
            cache = SlotCache(pinned, n, slots, policy, [rng.choice([1.0, 4.0]) for _ in range(n)])
            content = {slot: e for e, slot in cache.pinned.items()}  # what each slot really holds
            for _ in range(12):
                if rng.random() < 0.3:
                    guesses = [(e, 0.5) for e in rng.sample(range(n), rng.randint(1, n))]
                    recent = set(cache.recent)
                    for e, slot in cache.prefetch(guesses, rng.randint(0, 3)):
                        self.assertNotIn(content.get(slot), recent - {None})
                        content[slot] = e
                routed = rng.sample(range(n), rng.randint(1, n))
                heat = {e: rng.randint(1, 9) for e in routed}
                cpu = cache.choose_cpu(heat, Miss(mode=rng.choice(["copy", "auto"]), cpu_max_tokens=3))
                computed = []
                for slot_of, loads in cache.plan(heat, rng.randint(1, 8), cpu):
                    self.assertEqual(len(set(slot_of.values())), len(slot_of))
                    for e, slot in loads:
                        self.assertEqual(slot_of[e], slot)
                        content[slot] = e
                    for e, slot in slot_of.items():
                        self.assertEqual(content[slot], e)
                    computed += slot_of
                self.assertEqual(sorted(computed + sorted(cpu)), sorted(routed))
                held = list(cache.slots.values())
                self.assertEqual(len(set(held)), len(held))
                self.assertEqual(len(held) + len(cache.free), slots)


class Prefetch(TmpCase):
    def test_prefetch_is_low_priority_and_never_evicts_the_last_forward(self):
        cache = SlotCache([], 4, 2, Policy(eviction="lfu"))
        cache.plan({0: 1})
        self.assertEqual(cache.prefetch([(1, 0.5), (2, 0.4)], budget=2), [(1, 1)])  # 0 is recent
        cache.plan({1: 1})  # the prefetch paid off
        self.assertEqual(cache.prefetch([(3, 0.9)], budget=1), [(3, 0)])  # evicts 0, not 1
        cache.plan({2: 1})  # 3 was never used: it goes first
        self.assertEqual((cache.prefetches, cache.prefetch_hits, cache.prefetch_wasted), (2, 1, 1))
        self.assertEqual(set(cache.slots), {1, 2})
        path = os.path.join(self.dir, "stats.json")
        write_stats(path, {7: cache})
        with open(path) as f:
            stats = json.load(f)["layers"]["7"]
        self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]),
                         ([1, 1, 1, 0], 1, 2))
        self.assertEqual(stats["vram"], [1, 2])  # the experts in VRAM now, for the web monitor
        write_stats(path, {7: SlotCache([5, 3], 8, 2, Policy(eviction="lfu"))})
        with open(path) as f:
            self.assertEqual(json.load(f)["layers"]["7"]["vram"], [3, 5])  # pinned ones too