import math
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import spaces
from transformers import (
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedModel,
    GenerationMixin,
)
from transformers.modeling_outputs import CausalLMOutputWithPast

class FWKVConfig(PretrainedConfig):
    """Configuration class for FWKV-ROSA model architecture."""
    model_type = "fwkv"

    def __init__(
        self,
        d_model: int = 512,
        d_emb: int = 128,
        n_layers: int = 14,
        ffn_mult: int = 4,
        vocab_size: int = 50257,
        seq_len: int = 1024,        # trained with 1024
        wkv_floor: float = 0.1,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
        self.d_model = d_model
        self.d_emb = d_emb
        self.n_layers = n_layers
        self.ffn_mult = ffn_mult
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.wkv_floor = wkv_floor

def rosa(x: list[int]) -> list[int]:
    """Causal copy‑signal predictor; returns y[i] = token after longest
    repeating suffix ending at i, or -1 if none."""
    n = len(x)
    if n == 0:
        return []
    y = [-1] * n
    s = 2 * n + 2
    trans = [None] * s
    link = [-1] * s
    length = [0] * s
    last_end = [-1] * s
    trans[0] = {}
    last = 0
    size = 1

    for i, t in enumerate(x):
        cur = size; size += 1
        trans[cur] = {}
        length[cur] = length[last] + 1
        p = last
        while p != -1 and t not in trans[p]:
            trans[p][t] = cur
            p = link[p]
        if p == -1:
            link[cur] = 0
        else:
            q = trans[p][t]
            if length[p] + 1 == length[q]:
                link[cur] = q
            else:
                clone = size; size += 1
                trans[clone] = trans[q].copy()
                length[clone] = length[p] + 1
                link[clone] = link[q]
                last_end[clone] = last_end[q]
                while p != -1 and trans[p][t] == q:
                    trans[p][t] = clone
                    p = link[p]
                link[q] = clone
                link[cur] = clone
        last = cur

        v = cur
        pred = -1
        while v != -1:
            if length[v] > 0 and last_end[v] >= 0:
                pred = x[last_end[v] + 1]
                break
            v = link[v]
        y[i] = pred

        v = last
        while v != -1 and last_end[v] < i:
            last_end[v] = i
            v = link[v]
    return y

def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
    """Hillis–Steele inclusive scan with constant per‑channel decay."""
    W = W.to(dtype=a.dtype)          # keep precision
    val = a
    T = a.shape[1]
    d = 1
    while d < T:
        shifted = F.pad(val[:, :-d, :], (0, 0, d, 0))
        val = val + (W ** d) * shifted
        d *= 2
    return val

class FactorizedTiedHead(nn.Module):
    """Factorized embedding projection and tied output head."""
    def __init__(self, vocab_size: int, d_model: int, d_emb: int):
        super().__init__()
        self.d_model = d_model
        self.d_emb = d_emb
        self.weight = nn.Parameter(torch.empty(vocab_size, d_emb))
        self.proj = nn.Linear(d_emb, d_model, bias=False)

    def embed(self, input_ids):
        return self.proj(F.embedding(input_ids, self.weight))

    def to_emb_space(self, x):
        return F.linear(x, self.proj.weight.t())

    def logits(self, x_emb):
        return F.linear(x_emb, self.weight)

class FWKVBlock(nn.Module):
    """FWKV layer with linear time attention-style recurrent mechanism."""
    def __init__(self, d: int, ffn_mult: int = 4, floor: float = 0.1):
        super().__init__()
        self.floor = floor
        self.proj_k = nn.Linear(d, d, bias=False)
        self.proj_v = nn.Linear(d, d, bias=False)
        self.proj_r = nn.Linear(d, d, bias=False)
        self.proj_out = nn.Linear(d, d, bias=False)
        self.w = nn.Parameter(torch.ones(d) * 2.0)
        self.ffn = nn.Sequential(
            nn.Linear(d, ffn_mult * d, bias=False),
            nn.GELU(),
            nn.Linear(ffn_mult * d, d, bias=False),
        )
        self.norm_wkv = nn.LayerNorm(d)
        self.norm_ffn = nn.LayerNorm(d)

    @property
    def W(self):
        return torch.clamp(torch.sigmoid(self.w), min=self.floor)

    def forward(self, x, state=None):
        B, T, d = x.shape
        W = self.W
        k = self.proj_k(x)
        v = self.proj_v(x)
        r = torch.sigmoid(self.proj_r(x))

        a = k * v
        if state is not None:
            a = a.clone()
            a[:, 0] = a[:, 0] + W * state

        wkv_out = parallel_scan_decay(a, W)
        new_state = wkv_out[:, -1].detach()

        x = self.norm_wkv(x + self.proj_out(r * wkv_out))
        x = self.norm_ffn(x + self.ffn(x))
        return x, new_state

class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
    """Full causal language model utilizing FWKV recurrent layers and ROSA embeddings."""
    config_class = FWKVConfig

    def __init__(self, config):
        super().__init__(config)
        self.shared = FactorizedTiedHead(config.vocab_size, config.d_model, config.d_emb)
        self.rosa_emb = nn.Embedding(config.vocab_size + 1, config.d_emb, padding_idx=0)
        self.blocks = nn.ModuleList([
            FWKVBlock(config.d_model, config.ffn_mult, config.wkv_floor)
            for _ in range(config.n_layers)
        ])
        self.norm = nn.LayerNorm(config.d_model)
        self.post_init()

    def get_input_embeddings(self):
        return self.shared.weight

    def forward(
        self,
        input_ids,
        rosa_ids=None,
        past_key_values=None,
        labels=None,
        use_cache=True,
        **kwargs,
    ):
        if rosa_ids is None:
            rows = [rosa(row.tolist()) for row in input_ids.detach().cpu()]
            rosa_ids = torch.tensor(rows, device=input_ids.device, dtype=torch.long)

        x = self.shared.embed(input_ids)
        rosa_idx = (rosa_ids + 1).clamp(min=0)
        x = x + self.shared.proj(self.rosa_emb(rosa_idx))

        states_in = past_key_values or [None] * len(self.blocks)
        states_out = []
        for block, state in zip(self.blocks, states_in):
            x, new_state = block(x, state)
            states_out.append(new_state)

        x = self.norm(x)
        x_emb = self.shared.to_emb_space(x)
        logits = self.shared.logits(x_emb)

        return CausalLMOutputWithPast(
            loss=None,
            logits=logits,
            past_key_values=states_out if use_cache else None,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                      rosa_ids=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
            if rosa_ids is not None:
                rosa_ids = rosa_ids[:, -1:]
        return {"input_ids": input_ids, "rosa_ids": rosa_ids,
                "past_key_values": past_key_values, "use_cache": True}

USER_TOKEN = "<|user|>"
ASSISTANT_TOKEN = "<|assistant|>"

def load_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Loading FWKV-ROSA from Hub on {device} ...")
    try:
        model = FWKVLanguageModel.from_pretrained("FWKV/FWKV-ROSA")
        model = model.to(device)
        model.eval()
        tokenizer = AutoTokenizer.from_pretrained("FWKV/FWKV-ROSA")
        status = "FWKV-ROSA chat model ready!"
    except Exception as e:
        model, tokenizer = None, None
        status = f"Error loading model: {e}"
        print(status)
    return model, tokenizer, status


model, tokenizer, load_status = load_model()

@spaces.GPU
@torch.no_grad()
def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50):
    """Autoregressive generation with ROSA updates, yielding token lists and comprehensive throughput metrics."""
    device = next(model.parameters()).device
    eos_id = tokenizer.eos_token_id

    # Initial forward pass over the prompt
    inp = torch.tensor([ids], device=device)
    rosa_ids = torch.tensor([rosa(ids)], device=device)
    out = model(input_ids=inp, rosa_ids=rosa_ids, use_cache=True)
    states = out.past_key_values
    logits = out.logits[0, -1]
    generated = list(ids)
    reply_tokens = []
    
    start_time = time.perf_counter()
    prev_step_time = start_time
    instant_tps_list = []

    for _ in range(max_new_tokens):
        scaled = logits / max(temperature, 1e-5)
        if top_k and top_k < scaled.size(-1):
            kth = torch.topk(scaled, top_k).values[-1]
            scaled[scaled < kth] = float('-inf')
        probs = torch.softmax(scaled, dim=-1)
        next_token = torch.multinomial(probs, 1).item()
        generated.append(next_token)
        if next_token == eos_id:
            break

        reply_tokens.append(next_token)
        now = time.perf_counter()

        # Calculate per-step instant duration and speed
        step_duration = now - prev_step_time
        prev_step_time = now

        if step_duration > 0:
            instant_tps = 1.0 / step_duration
            instant_tps_list.append(instant_tps)

        # Compute aggregate throughput metrics
        total_elapsed = now - start_time
        avg_tps = len(reply_tokens) / total_elapsed if total_elapsed > 0 else 0.0
        current_tps = instant_tps_list[-1] if instant_tps_list else avg_tps
        min_tps = min(instant_tps_list) if instant_tps_list else avg_tps
        max_tps = max(instant_tps_list) if instant_tps_list else avg_tps

        stats = {
            "current": current_tps,
            "avg": avg_tps,
            "min": min_tps,
            "max": max_tps,
        }

        yield reply_tokens, stats

        # ROSA prediction for the next step
        next_rosa = rosa(generated)[-1]
        step_inp = torch.tensor([[next_token]], device=device)
        step_rosa = torch.tensor([[next_rosa]], device=device)
        out = model(input_ids=step_inp, rosa_ids=step_rosa,
                    past_key_values=states, use_cache=True)
        states = out.past_key_values
        logits = out.logits[0, -1]

def extract_text_content(content) -> str:
    """Safely extract plain text from string, list, or dict content structures returned by Gradio."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict):
                if "text" in item:
                    parts.append(str(item["text"]))
                elif "content" in item:
                    parts.append(extract_text_content(item["content"]))
            else:
                parts.append(str(item))
        return " ".join(parts)
    if isinstance(content, dict):
        if "text" in content:
            return str(content["text"])
        return str(content)
    return str(content) if content is not None else ""

def chat_function(message, history):
    """Gradio ChatInterface streaming handler formatted with speed stats (Live, Avg, Min, Max)."""
    messages = []
    for turn in history:
        if isinstance(turn, (list, tuple)):
            user_msg, asst_msg = turn
            messages.append({"role": "user", "content": extract_text_content(user_msg)})
            if asst_msg:
                messages.append({"role": "assistant", "content": extract_text_content(asst_msg)})
        elif isinstance(turn, dict):
            messages.append({
                "role": turn.get("role", "user"),
                "content": extract_text_content(turn.get("content", ""))
            })
    messages.append({"role": "user", "content": extract_text_content(message)})

    # Encode token sequence according to model chat template
    user_id = tokenizer.convert_tokens_to_ids(USER_TOKEN)
    asst_id = tokenizer.convert_tokens_to_ids(ASSISTANT_TOKEN)
    eos_id = tokenizer.eos_token_id
    ids = []
    for turn in messages:
        role = turn["role"]
        content = turn["content"]
        if not content.strip():
            continue
        content_ids = tokenizer.encode(" " + content)
        if role == "user":
            ids += [user_id] + content_ids
        elif role == "assistant":
            ids += [asst_id] + content_ids + [eos_id]

    # Truncate left if context exceeds model max sequence length
    max_len = model.config.seq_len if model else 1024
    if len(ids) > max_len:
        ids = ids[-max_len:]

    # Prompt assistant response
    ids.append(asst_id)

    # Stream generated output with full throughput statistics
    for reply_tokens, stats in generate_reply_stream(ids, max_new_tokens=150, temperature=0.8, top_k=50):
        reply = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip()
        metrics_bar = (
            f"⚡ **{stats['current']:.1f} tok/s** "
            f"*(Avg: **{stats['avg']:.1f}** | Min: **{stats['min']:.1f}** | Max: **{stats['max']:.1f}** tok/s)*"
        )
        yield f"{reply}\n\n{metrics_bar}"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"""
    # ⚡ FWKV-ROSA Chat
    **Model:** [FWKV/FWKV-ROSA](https://huggingface.co/FWKV/FWKV-ROSA)
    *{load_status}*

    This is a 56M‑parameter recurrent LM trained with the RWKV‑8 ROSA
    copy‑signal mechanism. It uses the chat template:

        `<|user|> message <|assistant|> reply <eos>`

    You can chat naturally; the model will remember recent context up to
    {model.config.seq_len if model else 1024} tokens.
    """)

    chatbot = gr.ChatInterface(
        fn=chat_function,
        title="",
        description="",
        examples=[
            "Explain how a linear recurrent network can still copy long‑range patterns.",
            "Write a short poem about a fox discovering a hidden library.",
        ],
    )

if __name__ == "__main__":
    demo.launch()