import importlib

import torch
from kernels.benchmark import Benchmark


class SolAttnBenchmark(Benchmark):
    seed: int = 42

    def _setup(self, tokens, heads):
        self.q = torch.randn(
            1,
            tokens,
            heads,
            128,
            device=self.device,
            dtype=torch.bfloat16,
        )
        self.k = torch.randn_like(self.q)
        self.v = torch.randn_like(self.q)
        module = importlib.import_module(f"{self.kernel.__name__}.triton_ref")
        self.triton_sol_attn = module.sol_attn

    def _run(self):
        self.out = self.kernel.sol_attn(
            self.q,
            self.k,
            self.v,
            tau=1.0,
            thresh_type="exact",
        )

    def _reference(self):
        return self.triton_sol_attn(
            self.q,
            self.k,
            self.v,
            tau=1.0,
            thresh_type="exact",
        )

    def setup_base(self):
        self._setup(tokens=4096, heads=8)

    def benchmark_base(self):
        self._run()

    def verify_base(self):
        return self._reference()

    def setup_video(self):
        self._setup(tokens=16384, heads=16)

    def benchmark_video(self):
        self._run()

    def verify_video(self):
        return self._reference()