########################################################################################################
# The EmbeddingRWKV Model - https://github.com/add_later
########################################################################################################

from types import SimpleNamespace
from typing import List, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn

from .reference.rwkv7 import RWKV_x070


Tokens = List[int]
BatchTokens = Sequence[Sequence[int]]


def generate_eos_mask(tokens: Union[Tokens, BatchTokens], eos_token_id: int) -> torch.Tensor:
    if isinstance(tokens, Sequence) and len(tokens) > 0 and isinstance(tokens[0], Sequence):
        # batch of sequences
        batch_size = len(tokens)
        max_len = max(len(t) for t in tokens)
        eos_mask = torch.zeros((batch_size, max_len), dtype=torch.bool)
        for i, t in enumerate(tokens):
            for j, token_id in enumerate(t):
                if token_id == eos_token_id:
                    eos_mask[i, j] = True
        return eos_mask
    else:
        # single sequence
        seq_len = len(tokens)
        eos_mask = torch.zeros((1, seq_len), dtype=torch.bool)
        for j, token_id in enumerate(tokens):
            if token_id == eos_token_id:
                eos_mask[0, j] = True
        return eos_mask

def extract_module_state_dict(state_dict, module_name: str):
    extracted = {}
    for key in state_dict:
        if key.startswith(module_name + "."):
            new_key = key[len(module_name) + 1 :]
            extracted[new_key] = state_dict[key]
    return extracted


class NonlinearHead(nn.Module):
    def __init__(self, dim_in, dim_hidden, dim_out=None):
        super().__init__()
        dim_out = dim_out or dim_in
        self.fc1 = nn.Linear(dim_in, dim_hidden, bias=False)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(dim_hidden, dim_out, bias=False)
        self.norm = nn.LayerNorm(dim_out)
        # zero initialization to make the head an identity function at the beginning
        nn.init.zeros_(self.fc2.weight)

    def forward(self, x):
        # Residual connection: stabilize training, avoid destroying embedding geometry
        h = self.fc2(self.act(self.fc1(x)))
        return self.norm(h + x)
   

class MultiEOSPooling(nn.Module):
    def __init__(self):
        super().__init__()

    @staticmethod
    def _masked_mean(x: torch.Tensor, mask: torch.Tensor):
        # x: [B,L,D], mask: [B,L] (bool)
        # to handle variable number of eos tokens
        # mask the value to 0 where eos_mask is False
        x = x.masked_fill(~mask.unsqueeze(-1), 0)
        return x.sum(1) / mask.sum(1, keepdim=True)    # [B,D]

    def forward(self, x: torch.Tensor, eos_mask: torch.Tensor):
        # x: [B,L,D], eos_mask: [B,L] (bool)
        return self._masked_mean(x, eos_mask)
    

class MultiTaskHead(nn.Module):
    def __init__(self, dim_in: int, dim_hidden: int):
        super().__init__()
        self.cls_head = NonlinearHead(dim_in, dim_hidden)
        self.sts_head = NonlinearHead(dim_in, dim_hidden)
        self.retr_head = NonlinearHead(dim_in, dim_hidden)
        self.task_keys = ["[CLS]", "[STS]", "[RETR]"]
        self.num_tasks = len(self.task_keys)

    def forward(self, x: torch.Tensor, task_type: str):
        # x: [B,D], task: str
        if task_type == "[CLS]":
            x = self.cls_head(x)
        elif task_type == "[STS]":
            x = self.sts_head(x)
        elif task_type == "[RETR]":
            x = self.retr_head(x)
        else:
            raise ValueError(f"Unknown task type: {task_type}")
        return x


class EmbeddingRWKV(nn.Module):
    def __init__(self, model_path: str):
        super().__init__()
        state_dict = torch.load(model_path, map_location='cpu', mmap=True, weight_only=True)
        rwkv_dict = extract_module_state_dict(state_dict, "rwkv")
        head_dict = extract_module_state_dict(state_dict, "head")
        # prepare rwkv
        model_prefix = model_path[:-4] if model_path.endswith(".pth") else model_path
        args = SimpleNamespace(MODEL_NAME=model_prefix)
        self.rwkv = RWKV_x070(args, rwkv_dict)
        # prepare head
        self.n_head, self.head_size = rwkv_dict['blocks.0.att.r_k'].shape
        self.n_embd = self.n_head * self.head_size
        self.head = MultiTaskHead(self.n_embd, self.n_embd)
        self.head.load_state_dict(head_dict)
        self.eos_pool = MultiEOSPooling()
        self.eos_token_id = 65535

    def generate_zero_state(self, bsz):
        return self.rwkv.generate_zero_state(bsz)

    @torch.inference_mode()
    def forward(
        self,
        tokens: Union[Tokens, BatchTokens],
        state: Optional[List[torch.Tensor]] = None,
        full_output: bool = True,
        task_type: str = None,
    ):
        is_batch = isinstance(tokens, Sequence) and len(tokens) > 0 and isinstance(tokens[0], Sequence)
        if state is None:
            B = len(tokens) if is_batch else 0
            state = self.generate_zero_state(B)
        if is_batch:
            output = self.rwkv.forward_batch(tokens, state, full_output)
        else:
            output = self.rwkv.forward(tokens, state, full_output)
        # prepare eos mask and pooling
        eos_mask = generate_eos_mask(tokens, self.eos_token_id).to(output.device)
        feature = self.eos_pool(output, eos_mask)
        # apply head
        if task_type is None:
            task_type = "[RETR]"
            embedding = self.head(feature, task_type)
        else:
            embedding = self.head(feature, task_type)

        return embedding, state
