import torch
import torch.nn as nn
import torch.nn.functional as F

_TRITON_AVAILABLE = False
try:
    import triton
    import triton.language as tl

    @triton.jit
    def _wkv7_fwd_kernel(
        R, K, V, DECAY, A, O,
        STATE_OUT, STATE_IN,
        sab_scale, T,
        stride_b, stride_t, stride_h,
        H: tl.constexpr, D: tl.constexpr, BLOCK_D: tl.constexpr,
        RETURN_STATE: tl.constexpr, HAS_INIT_STATE: tl.constexpr,
    ):
        pid = tl.program_id(0)
        b_idx = pid // H
        h_idx = pid % H
        base = b_idx * stride_b + h_idx * stride_h

        di = tl.arange(0, BLOCK_D)
        dj = tl.arange(0, BLOCK_D)
        mask_i = di < D
        mask_j = dj < D

        if HAS_INIT_STATE:
            s_off = b_idx * (H * D * D) + h_idx * (D * D)
            state_ptrs = STATE_IN + s_off + di[:, None] * D + dj[None, :]
            state_mask = mask_i[:, None] & mask_j[None, :]
            state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32)
        else:
            state = tl.zeros((BLOCK_D, BLOCK_D), dtype=tl.float32)

        for t in range(T):
            t_off = base + t * stride_t
            kt = tl.load(K + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
            vt = tl.load(V + t_off + di, mask=mask_i, other=0.0).to(tl.float32)
            rt = tl.load(R + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
            dt = tl.load(DECAY + t_off + dj, mask=mask_j, other=1.0).to(tl.float32)
            at = tl.load(A + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)

            sa = tl.sum(state * (-kt)[None, :], axis=1)
            ka = kt * at
            sab = sa[:, None] * ka[None, :]
            state = state * dt[None, :] + sab_scale * sab + vt[:, None] * kt[None, :]
            state = tl.minimum(tl.maximum(state, -10.0), 10.0)

            out_t = tl.sum(state * rt[None, :], axis=1)
            tl.store(O + t_off + di, out_t, mask=mask_i)

        if RETURN_STATE:
            s_off = b_idx * (H * D * D) + h_idx * (D * D)
            state_ptrs = STATE_OUT + s_off + di[:, None] * D + dj[None, :]
            state_mask = mask_i[:, None] & mask_j[None, :]
            tl.store(state_ptrs, state, mask=state_mask)

    def _wkv7_scan_triton(r, decay, k, v, a, sab_scale):
        B, T, H, D = r.shape
        r, k, v, decay, a = [x.contiguous() for x in (r, k, v, decay, a)]
        o = torch.empty_like(r)
        stride_b, stride_t, stride_h = T * H * D, H * D, D
        BLOCK_D = triton.next_power_of_2(D)
        _wkv7_fwd_kernel[(B * H,)](
            r, k, v, decay, a, o,
            None, None,
            float(sab_scale), T,
            stride_b, stride_t, stride_h,
            H=H, D=D, BLOCK_D=BLOCK_D,
            RETURN_STATE=False, HAS_INIT_STATE=False,
        )
        return o

    if torch.cuda.is_available():
        _TRITON_AVAILABLE = True
except Exception:
    pass

_FLA_AVAILABLE = False
try:
    import torch.distributed.tensor as _tdt
    if not hasattr(_tdt, 'Replicate'):
        try:
            from torch.distributed._tensor import Replicate as _R, Shard as _S
            _tdt.Replicate = _R; _tdt.Shard = _S
        except ImportError:
            pass
    if not hasattr(_tdt, 'Placement'):
        try:
            from torch.distributed._tensor.placement_types import Placement as _P
            _tdt.Placement = _P
        except ImportError:
            pass
    if not hasattr(_tdt, 'distribute_module'):
        _tdt.distribute_module = lambda *a, **kw: None
    from fla.ops.rwkv7 import chunk_rwkv7 as _fla_chunk_rwkv7
    if torch.cuda.is_available():
        _test_r = torch.randn(1, 1, 2, 64, device='cuda', dtype=torch.bfloat16, requires_grad=True)
        _test_w = -torch.ones(1, 1, 2, 64, device='cuda', dtype=torch.bfloat16)
        _test_o, _ = _fla_chunk_rwkv7(_test_r, _test_w, _test_r, _test_r, _test_r, _test_r,
                                       head_first=False)
        _test_o.sum().backward()
        if not _test_r.grad.isnan().any():
            _FLA_AVAILABLE = True
        del _test_r, _test_w, _test_o
        torch.cuda.empty_cache()
    else:
        _FLA_AVAILABLE = True
except Exception:
    pass


class BiRWKV7Layer(nn.Module):

    def __init__(self, hidden_size, num_heads):
        super().__init__()
        assert hidden_size % num_heads == 0
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads

        self.mu_r = nn.Parameter(torch.zeros(hidden_size))
        self.mu_w = nn.Parameter(torch.zeros(hidden_size))
        self.mu_k = nn.Parameter(torch.zeros(hidden_size))
        self.mu_v = nn.Parameter(torch.zeros(hidden_size))
        self.mu_a = nn.Parameter(torch.zeros(hidden_size))
        self.mu_g = nn.Parameter(torch.zeros(hidden_size))

        self.W_r = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_k = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_v = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_w = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_a = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_g = nn.Linear(hidden_size, hidden_size, bias=False)

        self.sab_gate = nn.Parameter(torch.tensor(-5.0))

        self.group_norm = nn.GroupNorm(num_heads, hidden_size)
        self.W_o = nn.Linear(hidden_size, hidden_size, bias=False)

        nn.init.normal_(self.W_w.weight, std=0.01)
        nn.init.normal_(self.W_a.weight, std=0.01)
        nn.init.normal_(self.W_g.weight, std=0.02)

    def _token_shift(self, x):
        x_prev = F.pad(x[:, :-1], (0, 0, 1, 0))

        def mix(mu):
            return x + (x_prev - x) * torch.sigmoid(mu)

        return {
            'r': mix(self.mu_r), 'w': mix(self.mu_w),
            'k': mix(self.mu_k), 'v': mix(self.mu_v),
            'a': mix(self.mu_a), 'g': mix(self.mu_g),
        }

    def _wkv7_scan_fla(self, r, w, k, v, a, sab_scale):
        B, T, H, D = r.shape
        orig_dtype = r.dtype
        r, w, k, v, a = [x.float() for x in (r, w, k, v, a)]
        k_scaled = k * (D ** -0.5)
        w_log = -0.6065306597633104 * torch.sigmoid(w)
        a_sig = torch.sigmoid(a)
        a_fla = -k_scaled
        b_fla = sab_scale * k_scaled * a_sig
        o, _ = _fla_chunk_rwkv7(r, k_scaled, v, a_fla, b_fla, log_w=w_log, scale=1.0, head_first=False)
        return o.to(orig_dtype)

    def _wkv7_scan_python(self, r, w, k, v, a, sab_scale):
        B, T, H, D = r.shape
        orig_dtype = r.dtype

        r, w, k, v, a = [x.float() for x in (r, w, k, v, a)]
        k = k * (D ** -0.5)
        decay = torch.exp(-0.6065306597633104 * torch.sigmoid(w))
        a = torch.sigmoid(a)

        state = torch.zeros(B, H, D, D, device=r.device, dtype=torch.float32)
        outputs = []

        for t in range(T):
            if t > 0 and t % 16 == 0:
                state = state.detach()

            kt, vt, rt, at, dt = k[:, t], v[:, t], r[:, t], a[:, t], decay[:, t]

            sa = torch.einsum('bhij,bhj->bhi', state, -kt)
            sab = torch.einsum('bhi,bhj->bhij', sa, kt * at)
            state = state * dt.unsqueeze(-2) + sab_scale * sab + torch.einsum('bhi,bhj->bhij', vt, kt)
            state = state.clamp(-10.0, 10.0)

            outputs.append(torch.einsum('bhij,bhj->bhi', state, rt))

        return torch.stack(outputs, dim=1).to(orig_dtype)

    def _wkv7_scan(self, r, w, k, v, a, sab_scale):
        if _TRITON_AVAILABLE and r.is_cuda:
            B, T, H, D = r.shape
            orig_dtype = r.dtype
            r, w, k, v, a = [x.float() for x in (r, w, k, v, a)]
            k = k * (D ** -0.5)
            decay = torch.exp(-0.6065306597633104 * torch.sigmoid(w))
            a = torch.sigmoid(a)
            return _wkv7_scan_triton(r, decay, k, v, a, sab_scale).to(orig_dtype)
        if _FLA_AVAILABLE and r.is_cuda:
            return self._wkv7_scan_fla(r, w, k, v, a, sab_scale)
        return self._wkv7_scan_python(r, w, k, v, a, sab_scale)

    def forward(self, x, attention_mask=None, **kwargs):
        B, T, C = x.shape
        H, D = self.num_heads, self.head_size

        mixed = self._token_shift(x)
        r = self.W_r(mixed['r']).view(B, T, H, D)
        w = self.W_w(mixed['w']).view(B, T, H, D)
        k = self.W_k(mixed['k']).view(B, T, H, D)
        v = self.W_v(mixed['v']).view(B, T, H, D)
        a = self.W_a(mixed['a']).view(B, T, H, D)
        g = torch.sigmoid(self.W_g(mixed['g']))

        sab_scale = torch.sigmoid(self.sab_gate)

        out_fwd = self._wkv7_scan(r, w, k, v, a, sab_scale)
        out_bwd = self._wkv7_scan(
            r.flip(1), w.flip(1), k.flip(1), v.flip(1), a.flip(1), sab_scale
        ).flip(1)

        out = (out_fwd + out_bwd).reshape(B, T, C) * 0.5
        out = self.group_norm(out.transpose(1, 2)).transpose(1, 2)
        out = self.W_o(out * g)

        return out, None


def init_from_attention(birwkv, attn_module):
    q_proj = k_proj = v_proj = o_proj = None

    if hasattr(attn_module, 'Wqkv'):
        fused = attn_module.Wqkv.weight.data
        C = fused.shape[1]
        q_proj, k_proj, v_proj = fused[:C], fused[C:2*C], fused[2*C:]
    else:
        for name in ['q_proj', 'query', 'W_q', 'wq']:
            if hasattr(attn_module, name):
                q_proj = getattr(attn_module, name).weight.data
                break
        for name in ['k_proj', 'key', 'W_k', 'wk']:
            if hasattr(attn_module, name):
                k_proj = getattr(attn_module, name).weight.data
                break
        for name in ['v_proj', 'value', 'W_v', 'wv']:
            if hasattr(attn_module, name):
                v_proj = getattr(attn_module, name).weight.data
                break

    for name in ['Wo', 'out_proj', 'o_proj', 'dense', 'W_o', 'wo']:
        if hasattr(attn_module, name):
            o_proj = getattr(attn_module, name).weight.data
            break

    transferred = []
    for src, dst, label in [
        (q_proj, birwkv.W_r, 'Q->R'),
        (k_proj, birwkv.W_k, 'K->K'),
        (v_proj, birwkv.W_v, 'V->V'),
        (o_proj, birwkv.W_o, 'O->O'),
    ]:
        if src is not None:
            dst.weight.data.copy_(src)
            transferred.append(label)

    return transferred