import logging
import math
from dataclasses import dataclass, field
from typing import Dict, Optional, Tuple, Union

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

from ..base import (
    BaseSequenceModelConfig,
    BaseSequenceModelTrain,
    ResettableParametersModule,
)
from .wkv_kernel import WKV, WKVConfig, WKVTorch

LOGGER = logging.getLogger(__name__)


class L2Wrap(torch.autograd.Function):
    """L2 regularization for the logits."""

    @staticmethod
    def forward(ctx, loss, y):
        ctx.save_for_backward(y)
        return loss

    @staticmethod
    def backward(ctx, grad_output):
        y = ctx.saved_tensors[0]
        # to encourage the logits to be close to 0
        factor = 1e-4 / (y.shape[0] * y.shape[1])
        maxx, ids = torch.max(y, -1, keepdim=True)
        gy = torch.zeros_like(y)
        gy.scatter_(-1, ids, maxx * factor)
        return (grad_output, gy)


def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    elif activation == "silu":
        return F.silu
    elif activation == "relu_squared":
        return lambda x: torch.square(torch.relu(x))
    elif activation == "selu":
        return F.selu
    elif activation == "elu":
        return F.elu
    else:
        raise ValueError(f"Unknown activation function {activation}")


@dataclass
class RWKVConfig(BaseSequenceModelConfig):
    embedding_dim: int = 640
    ffn_dim: int = 2048
    num_layers: int = 12
    attention_dim: int = -1
    wkv_config: Optional[Union[WKVConfig, Dict]] = field(
        default_factory=lambda: WKVConfig()
    )
    l2_logit_reg: bool = False
    use_timemix_timemix: bool = True
    use_timemix_channelmix: bool = True
    channelmix_act_fn: str = "relu_squared"
    init: str = "reproduce_init"  # options "paper_init", "reproduce_init" (reproduce init is different from paper init)
    # the reproduce_init was figured out when only the code was available, not the paper
    # the init as it is described in the paper slightly differs from the reproduce_init
    # the reproduce_init actually performs much better than the paper_init (ca. 0.5 better in train loss after 300 steps on WikiText-103)
    # so we use this as default. reproduce_init is also (likely to be) used in the original code.

    def __post_init__(self):
        # TODO: Check if this was actually needed
        if self.wkv_config is not None:
            self.wkv_config.T_max = max(
                self.wkv_config.T_max, self.context_length
            )
        if self.attention_dim <= 0:
            self.attention_dim = self.embedding_dim


class RWKV(BaseSequenceModelTrain):
    config_class = RWKVConfig

    def __init__(self, config: RWKVConfig, **kwargs):
        super().__init__()
        self.config = config
        self.cfg = config

        # self.embedding = nn.Embedding(
        #     num_embeddings=self.cfg.vocab_size,
        #     embedding_dim=self.cfg.embedding_dim,
        # )

        self.blocks = nn.ModuleList(
            [
                RWKVBlock(config=self.cfg, block_idx=i)
                for i in range(self.cfg.num_layers)
            ]
        )
        if self.cfg.wkv_config is not None:
            LOGGER.info("Using WKV cuda kernel.")
        else:
            LOGGER.info("Using WKV torch kernel.")

        # self.ln_out = nn.LayerNorm(self.cfg.embedding_dim)
        # self.head = nn.Linear(
        #     self.cfg.embedding_dim, self.cfg.vocab_size, bias=False
        # )
        self.reset_parameters()

    @property
    def context_length(self) -> int:
        return self.config.context_length

    # @property
    # def vocab_size(self) -> int:
    #     return self.config.vocab_size

    def reset_parameters(self) -> None:
        # init embedding
        # default init is zero # TODO try this
        # we use a narrow uniform init, in the original code they use the initial learning rate
        # we just set it to a small value
        # emb_init_range = 0.0008  # 1e-3
        # nn.init.uniform_(
        #     self.embedding.weight, a=-emb_init_range, b=emb_init_range
        # )
        # init blocks
        for b in self.blocks:
            b.reset_parameters()
        # init head and layer norm
        # self.head.reset_parameters()
        # self.ln_out.reset_parameters()

    def _create_optim_groups(self, config: RWKVConfig):
        optim_groups = [
            {"params": [p for p in self.parameters()], "weight_decay": 0.0}
        ]
        return optim_groups

    def forward(self, x):
        # no embedding
        # # input shape: (B, T), T <= context_len, T are token ids
        # B, T = x.size()
        # assert T <= self.cfg.context_length, (
        #     f"input sequence length {T} exceeds model "
        #     f"context length {self.cfg.context_length}"
        # )

        # x = self.embedding(x)  # (B, T, C), C = embedding_dim

        for i, block in enumerate(self.blocks):
            x = block(x)

        # x = self.ln_out(x)
        # no head
        # x = self.head(x)
        return x

    def get_loss_func(self):
        def loss_fn(y_hat, y):
            loss = F.cross_entropy(
                y_hat.view(-1, y_hat.size(-1)), y.view(-1), ignore_index=-1
            )
            if self.cfg.l2_logit_reg:
                loss = L2Wrap.apply(loss, y_hat)
            return loss

        return loss_fn


def _calc_gain(weight: torch.Tensor) -> float:
    """Calculate the gain value of the given weight tensor."""
    gain = 1.0
    fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(weight)
    if fan_out > fan_in:
        gain = math.sqrt(fan_out / fan_in)
    return gain


class RWKVBlock(ResettableParametersModule):
    def __init__(self, config: RWKVConfig, block_idx: int):
        super().__init__()
        self.config = config
        self.block_idx = block_idx

        self.ln0 = None
        if self.block_idx == 0:
            self.ln0 = nn.LayerNorm(self.config.embedding_dim)
            # TODO 1) maybe additional positional embedding here (only in block 0)

        self.ln1 = nn.LayerNorm(self.config.embedding_dim)
        self.ln2 = nn.LayerNorm(self.config.embedding_dim)

        # TODO 2) maybe pre feedforward here (channel mix) see line 325f in RWKV-v4neo/model.py
        self.attention_timemix = RWKVTimeMix(
            config=self.config, block_id=self.block_idx
        )
        self.ffn_channelmix = RWKVChannelMix(
            config=self.config, block_id=self.block_idx
        )

    def reset_parameters(self) -> None:
        if self.ln0 is not None:
            self.ln0.reset_parameters()
        self.ln1.reset_parameters()
        self.ln2.reset_parameters()
        self.attention_timemix.reset_parameters()
        self.ffn_channelmix.reset_parameters()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.block_idx == 0 and self.ln0 is not None:
            x = self.ln0(x)
            # TODO 1) maybe positional embedding here (only in block 0)
            # x = x+pos_emb

        # TODO 2) maybe pre feedforward here (channel mix) see line 325f in RWKV-v4neo/model.py
        # residual connection 1
        x = x + self.attention_timemix(self.ln1(x))
        # residual connection 2
        x = x + self.ffn_channelmix(self.ln2(x))
        return x


class RWKVTimeMix(ResettableParametersModule):
    def __init__(self, config: RWKVConfig, block_id: int):
        super().__init__()
        self.config = config
        self.block_id = block_id

        embedding_dim = self.config.embedding_dim
        attention_dim = self.config.attention_dim
        # init time mix constants
        req_grad = True  # TODO make this configurable
        self.time_mix_k = nn.Parameter(
            torch.empty((1, 1, embedding_dim)), requires_grad=req_grad
        )
        self.time_mix_v = nn.Parameter(
            torch.empty((1, 1, embedding_dim)), requires_grad=req_grad
        )
        self.time_mix_r = nn.Parameter(
            torch.empty((1, 1, embedding_dim)), requires_grad=req_grad
        )

        # init time decay
        self.time_decay = nn.Parameter(
            torch.empty((attention_dim,)), requires_grad=req_grad
        )
        self.time_first = nn.Parameter(
            torch.empty((attention_dim,)), requires_grad=req_grad
        )

        # init layers / parameters
        # this shifts the time dimension by 1 forward, pad 0 at first time step, remove last time step:
        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
        self.key = nn.Linear(embedding_dim, attention_dim, bias=False)
        self.value = nn.Linear(embedding_dim, attention_dim, bias=False)
        self.receptance = nn.Linear(embedding_dim, attention_dim, bias=False)
        self.output = nn.Linear(attention_dim, embedding_dim, bias=False)

        if self.config.wkv_config is not None:
            # use CUDA implementation
            self.wkv = WKV(config=self.config.wkv_config)
        else:
            # use pure PyTorch implementation
            self.wkv = WKVTorch()

        self.reset_parameters()

    def reset_parameters(self) -> None:
        # init time mix constants
        time_mix_k, time_mix_v, time_mix_r = self._init_time_mix_constants()
        req_grad = True
        self.time_mix_k = nn.Parameter(time_mix_k, requires_grad=req_grad)
        self.time_mix_v = nn.Parameter(time_mix_v, requires_grad=req_grad)
        self.time_mix_r = nn.Parameter(time_mix_r, requires_grad=req_grad)
        # init time decay
        time_decay, time_first = self._init_time_decay_constants()
        self.time_decay = nn.Parameter(time_decay, requires_grad=req_grad)
        self.time_first = nn.Parameter(time_first, requires_grad=req_grad)
        # init layers / parameters
        if self.config.init == "paper_init":
            # ZERO INIT
            nn.init.zeros_(self.receptance.weight)
            nn.init.zeros_(self.key.weight)
            nn.init.zeros_(self.value.weight)
            # NORMAL INIT
            nn.init.normal_(
                self.output.weight,
                std=math.sqrt(self.config.ffn_dim / self.config.embedding_dim),
            )
        elif self.config.init == "reproduce_init":
            # ZERO INIT
            nn.init.zeros_(self.key.weight)
            nn.init.zeros_(self.receptance.weight)
            nn.init.zeros_(self.output.weight)
            # ORTHOGONAL INIT
            nn.init.orthogonal_(
                self.value.weight, gain=_calc_gain(self.value.weight)
            )
        else:
            raise ValueError(f"Unknown init method {self.config.init}")

    def _compute_rkv(self, x):
        if self.config.use_timemix_timemix:
            xx = self.time_shift(
                x
            )  # Mix x with the previous timestep to produce xk, xv, xr
            xk = x * self.time_mix_k + xx * (1 - self.time_mix_k)
            xv = x * self.time_mix_v + xx * (1 - self.time_mix_v)
            xr = x * self.time_mix_r + xx * (1 - self.time_mix_r)
        else:
            xk = x
            xv = x
            xr = x
        k = self.key(xk)
        v = self.value(xv)
        r = self.receptance(xr)
        sr = torch.sigmoid(r)
        return sr, k, v

    def forward(self, x):
        B, T, C = x.size()  # x = (batch_size, seq_len, embedding_dim)
        attention_dim = self.config.attention_dim
        sr, k, v = self._compute_rkv(
            x
        )  # sr, k, v = (batch_size, seq_len, attention_dim)
        # wkv cuda/torch kernel
        rwkv = sr * self.wkv(
            B, T, attention_dim, self.time_decay, self.time_first, k, v
        )
        return self.output(rwkv)

    def _init_time_mix_constants(
        self,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        num_blocks = self.config.num_layers
        embedding_dim = self.config.embedding_dim

        ratio_0_to_1 = self.block_id / max(1, num_blocks - 1)  # 0 to 1
        ratio_1_to_almost0 = 1.0 - (self.block_id / num_blocks)  # 1 to ~0

        # TODO does this make sense?
        # different time mix constants for each block and each embedding dim
        embed_dim_val = torch.ones(1, 1, embedding_dim)
        for i in range(embedding_dim):
            embed_dim_val[0, 0, i] = i / embedding_dim

        # TODO check constants 0.3 and 0.5
        time_mix_k = torch.pow(embed_dim_val, ratio_1_to_almost0)
        time_mix_v = (
            torch.pow(embed_dim_val, ratio_1_to_almost0) + 0.3 * ratio_0_to_1
        )
        time_mix_r = torch.pow(embed_dim_val, 0.5 * ratio_1_to_almost0)

        return time_mix_k, time_mix_v, time_mix_r

    def _init_time_decay_constants(self) -> Tuple[torch.Tensor, torch.Tensor]:
        num_blocks = self.config.num_layers
        attention_dim = self.config.attention_dim
        ratio_0_to_1 = self.block_id / max(1, num_blocks - 1)  # 0 to 1

        # time decay
        # this encourages the model to decay the information in different memory cells (channel dimensions)
        # at different speeds
        decay_speed = torch.ones(attention_dim)
        for h in range(attention_dim):
            decay_speed[h] = -5 + 8 * (h / (attention_dim - 1)) ** (
                0.7 + 1.3 * ratio_0_to_1
            )
        time_decay = decay_speed

        # time first
        # The alternating zigzag pattern initially creates subtle variations in the tensor elements,
        # which are intended to help the model treat different dimensions of the embedding differently
        zigzag = (
            torch.tensor([(i + 1) % 3 - 1 for i in range(attention_dim)]) * 0.5
        )
        time_first = (
            torch.ones(attention_dim) * torch.log(torch.tensor(0.3)) + zigzag
        )

        return time_decay, time_first


class RWKVChannelMix(ResettableParametersModule):
    def __init__(self, config: RWKVConfig, block_id: int):
        super().__init__()
        self.config = config
        self.block_id = block_id

        self._act_fn = _get_activation_fn(self.config.channelmix_act_fn)

        embedding_dim = self.config.embedding_dim
        ffn_dim = self.config.ffn_dim
        # init time mix constants
        req_grad = True
        self.time_mix_k = nn.Parameter(
            torch.empty((1, 1, embedding_dim)), requires_grad=req_grad
        )
        self.time_mix_r = nn.Parameter(
            torch.empty((1, 1, embedding_dim)), requires_grad=req_grad
        )

        # init layers / parameters
        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
        self.key = nn.Linear(embedding_dim, ffn_dim, bias=False)
        self.receptance = nn.Linear(embedding_dim, embedding_dim, bias=False)
        self.value = nn.Linear(ffn_dim, embedding_dim, bias=False)
        self.reset_parameters()

    def reset_parameters(self):
        # init time mix constants
        time_mix_k, time_mix_r = self._init_time_mix_constants()
        req_grad = True
        self.time_mix_k = nn.Parameter(time_mix_k, requires_grad=req_grad)
        self.time_mix_r = nn.Parameter(time_mix_r, requires_grad=req_grad)
        # init layers / parameters
        if self.config.init == "paper_init":
            # ZERO INIT
            nn.init.zeros_(self.receptance.weight)
            nn.init.zeros_(self.key.weight)
            # NORMAL INIT
            nn.init.normal_(
                self.value.weight,
                std=math.sqrt(self.config.ffn_dim / self.config.embedding_dim),
            )
        elif self.config.init == "reproduce_init":
            # ZERO INIT
            nn.init.zeros_(self.receptance.weight)
            nn.init.zeros_(self.value.weight)
            # ORTHOGONAL INIT
            nn.init.orthogonal_(
                self.key.weight, gain=_calc_gain(self.key.weight)
            )
        else:
            raise ValueError(f"Unknown init method {self.config.init}")

    def _init_time_mix_constants(self) -> Tuple[torch.Tensor, torch.Tensor]:
        num_blocks = self.config.num_layers
        embedding_dim = self.config.embedding_dim

        ratio_1_to_almost0 = 1.0 - (self.block_id / num_blocks)  # 1 to ~0
        embed_dim_val = torch.ones(1, 1, embedding_dim)
        for i in range(embedding_dim):
            embed_dim_val[0, 0, i] = i / embedding_dim

        time_mix_k = torch.pow(embed_dim_val, ratio_1_to_almost0)
        time_mix_r = torch.pow(embed_dim_val, ratio_1_to_almost0)

        return time_mix_k, time_mix_r

    def forward(self, x):
        if self.config.use_timemix_channelmix:
            xx = self.time_shift(x)
            xk = x * self.time_mix_k + xx * (1 - self.time_mix_k)
            xr = x * self.time_mix_r + xx * (1 - self.time_mix_r)
        else:
            xk = x
            xr = x
        k = self.key(xk)
        k = self._act_fn(k)
        kv = self.value(k)
        y = torch.sigmoid(self.receptance(xr)) * kv
        return y