import json
import os
import threading
import time

import torch

from prain import profile
from tests.helpers import TmpCase


class Inner(torch.nn.Module):
    def forward(self, x):
        time.sleep(0.002)
        return x


class Outer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.ModuleList([Inner(), Inner()])
        self.head = Inner()

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        with profile.span("copy"):
            time.sleep(0.001)
        return self.head(x)


class TestProfile(TmpCase):
    def setUp(self):
        super().setUp()
        self.syncs = []
        saved = profile.DIR, profile.SYNC[0]
        profile.DIR, profile.SYNC[0] = self.dir, lambda: self.syncs.append(1)
        self.addCleanup(lambda: (setattr(profile, "DIR", saved[0]),
                                 profile.SYNC.__setitem__(0, saved[1])))

    def steps(self, model, tokens):
        for n in tokens:
            profile.step_begin()
            with profile.span("execute_model"):
                profile.note_step(model, n)
                model(torch.zeros(1))
            profile.step_end()

    def folded(self, kind):
        with open(os.path.join(self.dir, f"{kind}.folded")) as f:
            return {stack: int(us) for stack, us in (l.rsplit(" ", 1) for l in f.read().splitlines())}

    def test_records_the_requested_steps_as_folded_stacks(self):
        model = Outer()
        self.steps(model, [1])  # no start file: nothing recorded, nothing synced
        self.assertEqual(self.syncs, [])
        self.write("start", "3")
        self.steps(model, [1, 1, 300, 1])  # records 2 decode steps and 1 prefill, not the 4th
        self.assertFalse(os.path.exists(os.path.join(self.dir, "start")))
        with open(os.path.join(self.dir, "steps.json")) as f:
            self.assertEqual(json.load(f)["steps"], {"decode": 2, "prefill": 1})
        decode = self.folded("decode")
        self.assertEqual(set(decode), {"step", "step;execute_model", "step;execute_model;Outer",
                                       "step;execute_model;Outer;Inner",  # both layers, merged
                                       "step;execute_model;Outer;copy",
                                       "step;execute_model;Outer;head:Inner"})
        for stack, ms in (("Outer;Inner", 8), ("Outer;head:Inner", 4), ("Outer;copy", 2)):
            self.assertGreaterEqual(decode[f"step;execute_model;{stack}"], ms * 1000)
            self.assertLess(decode[f"step;execute_model;{stack}"], ms * 1000 * 3)
        self.assertGreater(self.folded("prefill")["step;execute_model;Outer;Inner"], 4000)
        with open(os.path.join(self.dir, "decode.calls")) as f:
            calls = dict(line.rsplit(" ", 1) for line in f.read().splitlines())
        self.assertEqual(calls["step;execute_model;Outer;Inner"], "4")  # 2 layers x 2 steps
        self.assertEqual(calls["step"], "2")
        recorded = len(self.syncs)
        self.steps(model, [1])  # finished: the module hooks are gone
        self.assertEqual(len(self.syncs), recorded)
        self.assertFalse(model.head._forward_hooks or model.head._forward_pre_hooks)

    def test_hooks_leave_other_threads_inputs_and_outputs_alone(self):
        model = Outer()
        self.write("start", "1")
        profile.step_begin()
        profile.note_step(model, 1)  # installs the hooks
        out = []
        thread = threading.Thread(target=lambda: out.append(model(torch.ones(1))))
        thread.start()
        thread.join()
        self.assertTrue(torch.equal(out[0], torch.ones(1)))
        profile.step_end()

    def test_svg_shows_time_per_step(self):
        lines = ["step 500", "step;attn 3000", "step;attn;copy 1000"]
        svg = profile.svg(lines, steps=2, title="decode")
        self.assertTrue(svg.startswith("<svg"))
        self.assertIn("attn: 2.0 ms per step (88.9%)", svg)  # 4000 us of 4500, over 2 steps
        self.assertIn("copy: 0.5 ms per step", svg)
        self.assertEqual(profile.top(lines, steps=2)[0], ("attn", 1.5, 66.7, None))  # self time
        calls = ["step;attn 10", "step;attn;copy 4", "step 2"]
        self.assertEqual(profile.top(lines, steps=2, calls=calls)[0], ("attn", 1.5, 66.7, 5.0))