from types import MethodType
import torch
import torch.nn.functional as F
import triton
import triton.language as tl

base_linear = F.linear
checks = []
checked = set()
writer_owners = {}


@triton.jit
def joint_matvec(X, Q, W, A, Y, NQ: tl.constexpr, NW: tl.constexpr,
                 NA: tl.constexpr, K: tl.constexpr, BK: tl.constexpr):
    start = tl.program_id(0) * 2
    rows = start + tl.arange(0, 2)
    columns = tl.arange(0, BK)
    x = tl.load(X + columns, columns < K, 0).to(tl.float32)
    if start < NQ:
        weights = Q + rows[:, None] * K + columns[None, :]
    elif start < NQ + NW:
        weights = W + (rows[:, None] - NQ) * K + columns[None, :]
    else:
        weights = A + (rows[:, None] - NQ - NW) * K + columns[None, :]
    weight = tl.load(weights,
        (rows[:, None] < NQ + NW + NA) & (columns[None, :] < K), 0).to(tl.float32)
    result = tl.sum(weight * x[None, :], axis=1)
    tl.store(Y + rows, result, rows < NQ + NW + NA)


def writer_linear(values, weight, bias=None):
    owner = writer_owners.get(weight.data_ptr())
    if owner is not None and owner._b1_joint_active and values.numel() == values.shape[-1]:
        begin = owner.q_proj.out_features
        result = owner._b1_joint_output[..., begin:begin + weight.shape[0]].reshape(*values.shape[:-1], weight.shape[0])
        owner._b1_joint_active = False
        owner._b1_joint_output = None
        return result
    return base_linear(values, weight, bias)


def install(model):
    for m in model.modules():
        if not (hasattr(m, '_query_scores') and hasattr(m, '_inference_writer_gate_up')):
            continue
        m._b1_joint_output = None
        m._b1_joint_active = False
        writer_owners[m._inference_writer_gate_up.data_ptr()] = m
        previous_query, previous_admission = m.q_proj.forward, m.admission_down.forward
        assert all(w.is_contiguous() and w.shape[0] % 2 == 0 for w in
            (m.q_proj.weight, m._inference_writer_gate_up, m.admission_down.weight))

        def query(projection, values, *args, owner=m, previous=previous_query, **kwargs):
            owner._b1_joint_active = values.numel() == values.shape[-1]
            if not owner._b1_joint_active:
                return previous(values, *args, **kwargs)
            q, w, a = owner.q_proj.weight, owner._inference_writer_gate_up, owner.admission_down.weight
            n = q.shape[0] + w.shape[0] + a.shape[0]
            output = values.new_empty((*values.shape[:-1], n))
            joint_matvec[triton.cdiv(n, 2),](values, q, w, a, output,
                q.shape[0], w.shape[0], a.shape[0], q.shape[1],
                triton.next_power_of_2(q.shape[1]), num_warps=4, num_stages=1)
            if owner not in checked:
                expected = torch.cat([base_linear(values, weight) for weight in (q, w, a)], -1)
                checks.append(dict(layer=len(checks), shape=list(output.shape),
                    exact=bool(torch.equal(expected, output)), finite=bool(output.isfinite().all()),
                    max_abs=float((expected.float() - output.float()).abs().max())))
                checked.add(owner)
            owner._b1_joint_output = output
            return output[..., :q.shape[0]]

        def admission(projection, values, *args, owner=m, previous=previous_admission, **kwargs):
            if not owner._b1_joint_active:
                return previous(values, *args, **kwargs)
            return owner._b1_joint_output[..., -owner.admission_down.out_features:]

        m.q_proj.forward = MethodType(query, m.q_proj)
        m.admission_down.forward = MethodType(admission, m.admission_down)
    assert len(writer_owners) == 8
    F.linear = writer_linear