from dataclasses import dataclass
from pathlib import Path
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn
from torch.nn import CrossEntropyLoss

from transformers.modeling_utils import PreTrainedModel
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import (
    ModelOutput,
    is_torch_cuda_available,
)


rwkv5_cuda_kernel = None
class Rwkv5Config(PretrainedConfig):

    model_type = "rwkv5"

    def __init__(
        self,
        vocab_size=65536,
        hidden_size=768,
        num_hidden_layers=24,
        attention_hidden_size=None,
        head_size=64,
        intermediate_size=None,
        layer_norm_epsilon=1e-5,
        bos_token_id=0,
        eos_token_id=0,
        rescale_every=0,
        tie_word_embeddings=False,
        use_cache=False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.attention_hidden_size = attention_hidden_size if attention_hidden_size is not None else hidden_size
        self.head_size = head_size
        self.intermediate_size = None
        self.layer_norm_epsilon = layer_norm_epsilon
        self.rescale_every = rescale_every
        self.use_cache = use_cache

        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id

        super().__init__(
            tie_word_embeddings=tie_word_embeddings, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs
        )

@torch.jit.script
def rwkv_linear_attention(
    B: int,
    H: int,
    S: int,
    T: int,
    n_head: int,
    hidden,
    time_decay,
    time_first,
    receptance,
    key,
    value,
    gate,
    lxw,
    lxb,
):
    key = key.to(torch.float32).view(B, T, H, S).transpose(1, 2).transpose(-2, -1)
    value = value.to(torch.float32).view(B, T, H, S).transpose(1, 2)
    receptance = receptance.to(torch.float32).view(B, T, H, S).transpose(1, 2)
    time_decay = torch.exp(-torch.exp(time_decay.float())).reshape(-1, 1, 1).reshape(n_head, -1, 1)
    time_first = time_first.float().reshape(-1, 1, 1).reshape(n_head, -1, 1)
    lxw = lxw.float()
    lxb = lxb.float()
    out = torch.zeros_like(key).reshape(B, T, H, S)
    for t in range(T):
        rt = receptance[:, :, t : t + 1, :]
        kt = key[:, :, :, t : t + 1]
        vt = value[:, :, t : t + 1, :]
        at = kt @ vt
        out[:, t] = (rt @ (time_first * at)).squeeze(2)

    out = out.reshape(B * T, H * S)
    out = F.group_norm(out, num_groups=H, weight=lxw, bias=lxb).reshape(B, T, H * S)
    out = out.to(dtype=hidden.dtype) * gate

    return out

@torch.jit.script
def extract_key_value( hidden, shifted, time_mix_key, time_mix_value, time_mix_receptance, time_mix_gate):
    key = hidden * time_mix_key + shifted * (1 - time_mix_key)
    value = hidden * time_mix_value + shifted * (1 - time_mix_value)
    receptance = hidden * time_mix_receptance + shifted * (1 - time_mix_receptance)
    gate = hidden * time_mix_gate + shifted * (1 - time_mix_gate)

    return receptance, key, value, gate

class RwkvSelfAttention(nn.Module):
    def __init__(self, config, layer_id=0):
        super().__init__()
        self.config = config
        self.layer_id = layer_id
        hidden_size = config.hidden_size
        # https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py#L146
        num_attention_heads = hidden_size // config.head_size
        self.num_attention_heads = num_attention_heads
        attention_hidden_size = (
            config.attention_hidden_size if config.attention_hidden_size is not None else hidden_size
        )
        self.attention_hidden_size = attention_hidden_size

        self.time_decay = nn.Parameter(torch.empty(num_attention_heads, config.head_size))
        self.time_faaaa = nn.Parameter(torch.empty(num_attention_heads, config.head_size))
        self.time_mix_gate = nn.Parameter(torch.empty(1, 1, hidden_size))

        self.time_mix_key = nn.Parameter(torch.empty(1, 1, hidden_size))
        self.time_mix_value = nn.Parameter(torch.empty(1, 1, hidden_size))
        self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, hidden_size))

        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
        self.key = nn.Linear(hidden_size, attention_hidden_size, bias=False)
        self.value = nn.Linear(hidden_size, attention_hidden_size, bias=False)
        self.receptance = nn.Linear(hidden_size, attention_hidden_size, bias=False)
        self.gate = nn.Linear(hidden_size, attention_hidden_size, bias=False)
        self.output = nn.Linear(attention_hidden_size, hidden_size, bias=False)
        # https://github.com/BlinkDL/RWKV-LM/blob/3db37a72356b736966ddd377268f02b80963af3f/RWKV-v4neo/src/model.py#L190C1-L190C1
        self.ln_x = nn.GroupNorm(hidden_size // config.head_size, hidden_size)




    def forward(self, hidden):
        B = hidden.shape[0]
        H = self.time_decay.shape[0]
        S = hidden.shape[-1] // H
        T = hidden.shape[1]


        # Mix hidden with the previous timestep to produce key, value, receptance
        shifted = self.time_shift(hidden)

        receptance, key, value, gate = extract_key_value(hidden, shifted, self.time_mix_key, self.time_mix_value, self.time_mix_receptance, self.time_mix_gate)

        # https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L693
        key = self.key(key)
        value = self.value(value)
        receptance = self.receptance(receptance)
        gate = F.silu(self.gate(gate), inplace=True)

        rwkv = rwkv_linear_attention(
            B,
            H,
            S,
            T,
            self.num_attention_heads,
            hidden,
            self.time_decay,
            self.time_faaaa,
            receptance,
            key,
            value,
            gate,
            self.ln_x.weight,
            self.ln_x.bias,
        )


        return self.output(rwkv)


class RwkvFeedForward(nn.Module):
    def __init__(self, config, layer_id=0):
        super().__init__()
        self.config = config
        self.layer_id = layer_id
        self.hidden_size = config.hidden_size
        # https://github.com/BlinkDL/RWKV-LM/blob/3db37a72356b736966ddd377268f02b80963af3f/RWKV-v4neo/train.py#L168
      
        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
        self.time_maa_k = nn.Parameter(torch.randn(1,1,self.hidden_size))

        self.key = nn.Linear(self.hidden_size, self.hidden_size * 2, bias=False)
        self.value = nn.Linear(self.hidden_size * 2, self.hidden_size, bias=False)
        #self.value.weight.data.zero_()
        
    
    def forward(self, x):
        
        xx = self.time_shift(x) - x        
        k = x + xx * self.time_maa_k
        k = torch.relu(self.key(k)) ** 2
        return self.value(k)

class RwkvBlock(nn.Module):
    def __init__(self, config, layer_id):
        super().__init__()
        self.config = config
        self.layer_id = layer_id

        if layer_id == 0:
            self.pre_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)

        self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)

        self.attention = torch.jit.script(RwkvSelfAttention(config, layer_id))
        self.feed_forward = torch.jit.script(RwkvFeedForward(config, layer_id))

    def forward(self, hidden):
        if self.layer_id == 0:
            hidden = self.pre_ln(hidden)
        attention = self.attention(self.ln1(hidden))
        hidden = hidden + attention

        feed_forward = self.feed_forward(self.ln2(hidden))
        hidden = hidden + feed_forward

        return hidden


class Rwkv5PreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = Rwkv5Config
    base_model_prefix = "rwkv"
    _no_split_modules = ["RwkvBlock"]
    _keep_in_fp32_modules = ["time_decay", "time_first"]
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, RwkvSelfAttention):
            layer_id = module.layer_id
            num_hidden_layers = module.config.num_hidden_layers
            hidden_size = module.config.hidden_size
            attention_hidden_size = module.attention_hidden_size
            num_attention_heads = hidden_size // module.config.head_size

            ratio_0_to_1 = layer_id / (num_hidden_layers - 1)  # 0 to 1
            ratio_1_to_almost0 = 1.0 - (layer_id / num_hidden_layers)  # 1 to ~0

            time_weight = torch.tensor(
                [i / hidden_size for i in range(hidden_size)],
                dtype=module.time_mix_key.dtype,
                device=module.time_mix_key.device,
            )
            time_weight = time_weight[None, None, :]

            # https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py#L398
            decay_speed = [
                -6.0 + 5.0 * (h / (attention_hidden_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
                for h in range(attention_hidden_size)
            ]
            decay_speed = torch.tensor(decay_speed, dtype=module.time_decay.dtype, device=module.time_decay.device)
            tmp = torch.tensor(
                [
                    (1.0 - (i / (attention_hidden_size - 1.0))) * ratio_0_to_1 + 0.1 * ((i + 1) % 3 - 1)
                    for i in range(attention_hidden_size)
                ],
                dtype=module.time_faaaa.dtype,
                device=module.time_faaaa.device,
            )

            with torch.no_grad():
                module.time_decay.data = decay_speed.reshape(num_attention_heads, module.config.head_size)
                module.time_faaaa.data = tmp.reshape(num_attention_heads, module.config.head_size)
                module.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0)

                module.time_mix_value.data = torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1
                module.time_mix_receptance.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
                module.time_mix_gate.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0)
                
        elif isinstance(module, RwkvFeedForward):
             with torch.no_grad():
                ratio_1_to_almost0 = 1.0 - (module.layer_id / module.config.num_hidden_layers)  # 1 to ~0
                ddd = torch.ones(1, 1, module.hidden_size)
                for i in range(module.hidden_size):
                    ddd[0, 0, i] = i / module.hidden_size
                print(module.time_maa_k.data)
                module.time_maa_k.data = ddd
                print(module.time_maa_k.data)

@dataclass
class Rwkv5CausalLMOutput(ModelOutput):

    loss = None
    logits = None
    state = None
    hidden_states = None
    attentions = None

    def __init__(self, loss=None, logits=None, state=None, hidden_states=None, attentions=None):
        super().__init__(
            loss=loss,
            logits=logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )
        self.state = state


class Rwkv5Model(Rwkv5PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([RwkvBlock(config, layer_id=idx) for idx in range(config.num_hidden_layers)])
        self.ln_out = nn.LayerNorm(config.hidden_size)

        self.layers_are_rescaled = False
        self.gradient_checkpointing = False

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings

    def set_input_embeddings(self, new_embeddings):
        self.embeddings = new_embeddings

    def forward(
        self,
        input_ids = None,
        inputs_embeds= None,
    ):

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is None and inputs_embeds is None:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embeddings(input_ids)

        hidden_states = inputs_embeds

        for idx, block in enumerate(self.blocks):
            hidden_states = block(
                hidden_states
            )

        hidden_states = self.ln_out(hidden_states)


        return hidden_states


class Rwkv5ForCausalLM(Rwkv5PreTrainedModel):
    _tied_weights_keys = ["head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.rwkv = Rwkv5Model(config)
        self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, new_embeddings):
        self.head = new_embeddings

    def prepare_inputs_for_generation(self, input_ids, state=None, inputs_embeds=None, **kwargs):
        # only last token for inputs_ids if the state is passed along.
        if state is not None:
            input_ids = input_ids[:, -1].unsqueeze(-1)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and state is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs["state"] = state
        return model_inputs


    def forward(
        self,
        input_ids = None,
        labels = None,
        inputs_embeds = None,
        return_dict = None,
    ):
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        rwkv_outputs = self.rwkv(
            input_ids,
            inputs_embeds=inputs_embeds,
        )

        logits = self.head(rwkv_outputs)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(logits.device)
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        if not return_dict:
            output = (logits,) + rwkv_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return Rwkv5CausalLMOutput(
            loss=loss,
            logits=logits,
        )


Convert this model to tensorflow