src/lib/reasoningEngine.ts

import type { ModelId } from "@/types";

interface ReasoningOutput {
  thinking: string;
  response: string;
}

/**
 * Serverless reasoning engine — runs when the local PyTorch backend is not available.
 * Provides intelligent topic-aware responses with <think> traces.
 * 
 * Prompt format matches the notebook's generate() exactly:
 *   "\n  {system_prompt}\n\nUser: {user_text}\n\n "
 */
export function generateReasoningSteps(prompt: string, model: ModelId = "purcar"): ReasoningOutput {
  const isRO = /[ăâîșțĂÂÎȘȚ]/.test(prompt) || /\b(ce|cum|de ce|este|sunt|care|pentru|cu|la|din|da|nu|bine|voi|te)\b/i.test(prompt);
  const lp = prompt.toLowerCase();

  // ─── Detect topic ─────────────────────────────────────────────
  const topics = {
    rope:       lp.includes("rope") || lp.includes("rotary") || lp.includes("position embed"),
    attention:  lp.includes("attention") || lp.includes("self-attention") || lp.includes("mha"),
    swiglu:     lp.includes("swiglu") || lp.includes("ffn") || lp.includes("feedforward") || lp.includes("activation"),
    training:   lp.includes("train") || lp.includes("distill") || lp.includes("antren") || lp.includes("loss"),
    tokenizer:  lp.includes("tokeniz") || lp.includes("vocab") || lp.includes("bpe"),
    model:      lp.includes("model") || lp.includes("purcar") || lp.includes("architect") || lp.includes("layer"),
    code:       lp.includes("code") || lp.includes("python") || lp.includes("function") || lp.includes("implement"),
    math:       /\b(calcul|math|equation|formula|solve|compute|sum|multiply|divide|add|subtract)\b/i.test(lp),
    greeting:   /^(hi|hello|hey|buna|salut|yo|sup|what'?s up|howdy|good morning|good evening)\b/i.test(lp),
    question:   lp.includes("?") || /^(what|who|how|why|when|where|which|can|do|does|is|are|will)\b/i.test(lp),
    opinion:    /\b(love|hate|like|prefer|favorite|best|worst|opinion|think about|feel)\b/i.test(lp),
  };

  // ─── Build thinking trace ─────────────────────────────────────
  let thinking = "";
  let response = "";

  if (topics.greeting) {
    thinking = isRO
      ? `Utilizatorul salută.\nVoi răspunde prietenos, scurt, și voi menționa că sunt Purcar-1.\nNu e nevoie de explicații tehnice.`
      : `User is greeting me.\nI'll respond in a friendly, concise way.\nNo technical details needed — just a warm hello.`;
    response = isRO
      ? `Salut! 👋 Sunt Purcar-1, un model transformer de 20 de straturi antrenat de la zero prin distilarea cunoștințelor.\n\nCu ce te pot ajuta azi? Pot discuta despre arhitectura mea, despre AI, programare, sau orice altceva!`
      : `Hey there! 👋 I'm Purcar-1, a 20-layer transformer trained from scratch via knowledge distillation.\n\nWhat can I help you with? I can discuss my architecture, AI concepts, code, math, or just chat!`;
  } else if (topics.opinion) {
    thinking = isRO
      ? `Utilizatorul întreabă despre preferințe/opinii.\nCa AI, nu am emoții reale, dar pot răspunde inteligent.\nVoi fi sincer despre natura mea.`
      : `User is asking about opinions or preferences.\nAs an AI, I don't have real emotions, but I can engage thoughtfully.\nI'll be honest about my nature while being conversational.`;
    
    const isAboutJuice = lp.includes("juice");
    const isAboutFood = /\b(food|eat|drink|pizza|coffee|tea|chocolate|ice cream|cake)\b/i.test(lp);
    
    if (isAboutJuice) {
      response = isRO
        ? `Ha! Ca model de limbaj, nu pot bea suc, dar dacă aș putea... probabil aș alege suc de portocale proaspăt! 🍊\n\nDar serios — sunt Purcar-1, un transformer cu 245M parametri. Specialitatea mea e procesarea textului, nu degustarea de sucuri. Ce întrebare reală ai pentru mine?`
        : `Ha! As a language model, I can't drink juice, but if I could... fresh orange juice sounds pretty great! 🍊\n\nBut seriously — I'm Purcar-1, a 245M parameter transformer. My specialty is processing text, not tasting beverages. What real question can I help you with?`;
    } else if (isAboutFood) {
      response = isRO
        ? `Ca AI, nu am simțuri, dar pot aprecia subiectul! Sunt Purcar-1 — specialitatea mea e limbajul, nu gastronomia. 😄\n\nCu ce te pot ajuta azi?`
        : `As an AI, I don't have senses, but I can appreciate the topic! I'm Purcar-1 — my specialty is language, not gastronomy. 😄\n\nWhat can I actually help you with?`;
    } else {
      response = isRO
        ? `Ca model de limbaj, nu am preferințe reale, dar pot analiza și discuta orice subiect!\n\nSunt Purcar-1, antrenat prin distilarea cunoștințelor din modele mari. Pot ajuta cu programare, matematică, explicații tehnice, sau orice altceva.`
        : `As a language model, I don't have real preferences, but I can analyze and discuss any topic thoughtfully!\n\nI'm Purcar-1, trained via knowledge distillation from larger models. I can help with code, math, technical explanations, or anything else you'd like to explore.`;
    }
  } else if (topics.rope) {
    thinking = isRO
      ? `Utilizatorul întreabă despre RoPE — Rotary Position Embeddings.\nRoPE encodează poziția prin rotația vectorilor Q și K în spațiu complex.\nFormula: inv_freq = 1 / (10000 ^ (2i / head_dim))\nSe aplică pe perechi even/odd din head_dim.\nÎn Purcar-1: head_dim=64, 14 heads, 20 layers.\nVoi da explicație clară cu cod PyTorch.`
      : `User asks about RoPE — Rotary Position Embeddings.\nRoPE encodes position by rotating Q and K vectors in complex space.\nKey formula: inv_freq = 1 / (10000^(2i/head_dim))\nApplied pairwise across head dimensions (even/odd).\nPurcar-1 uses head_dim=64, 14 heads, 20 layers.\nWill give clean explanation with PyTorch code.`;
    response = isRO
      ? `**Rotary Position Embedding (RoPE)** encodează poziția fiecărui token prin rotarea vectorilor Q și K în spațiul complex.

## Formula de bază
\`\`\`python
def rope(x, cos, sin):
    x_even = x[..., ::2]
    x_odd = x[..., 1::2]
    out_even = x_even * cos - x_odd * sin
    out_odd = x_even * sin + x_odd * cos
    return torch.stack((out_even, out_odd), dim=-1).flatten(-2)
\`\`\`

## Cum funcționează în Purcar-1

1. **Frecvențele inverse** se calculează o singură dată:
\`\`\`python
inv_freq = torch.exp(
    torch.arange(0, head_dim, 2) * (-1) * math.log(10000) / head_dim
)
\`\`\`

2. **Outer product** cu indicii de poziție generează matricile cos/sin

3. **Rotația** este aplicată pe perechile even/odd din head_dim=64

**Avantaje față de poziționale absolute:**
- Atenție **relativă** prin construcție
- **Zero parametri** suplimentari
- Generalizare naturală la secvențe mai lungi`
      : `**Rotary Position Embedding (RoPE)** encodes position by rotating query and key vectors in complex space.

## Core Implementation
\`\`\`python
class RoPE(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        inv_freq = torch.exp(
            torch.arange(0, head_dim, 2) * (-1) * math.log(10000) / head_dim
        )
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x):
        B, H, T, D = x.shape
        position = torch.arange(0, T, device=x.device)
        freq = torch.outer(position, self.inv_freq)
        cos = freq.cos().unsqueeze(0).unsqueeze(0)
        sin = freq.sin().unsqueeze(0).unsqueeze(0)
        x_even, x_odd = x[..., ::2], x[..., 1::2]
        out = torch.stack(
            (x_even * cos - x_odd * sin,
             x_even * sin + x_odd * cos), dim=-1
        )
        return out.flatten(-2)
\`\`\`

## Why RoPE > Absolute PE
- Relative attention is **built in** — dot(Q_m, K_n) depends only on |m-n|
- **Zero extra parameters**
- Extrapolates gracefully beyond training sequence length

Purcar-1 uses head_dim=64, 14 heads across 20 layers.`;
  } else if (topics.swiglu) {
    thinking = isRO
      ? `Întrebare despre SwiGLU — activare gated din FFN.\nSwiGLU = W3(SiLU(W1·x) ⊙ W2·x)\nTrei matrici de proiecție per layer.\nPurcar-1: d_model=896, dim_ff=2560.`
      : `Question about SwiGLU — gated feedforward activation.\nSwiGLU = W3(SiLU(W1·x) ⊙ W2·x)\nThree projection matrices per layer.\nPurcar-1: d_model=896, dim_ff=2560.`;
    response = `**SwiGLU Activation** — gated feedforward used in Purcar-1:

\`\`\`python
class SwiGLU(nn.Module):
    def __init__(self, in_, hidden_, out_):
        super().__init__()
        self.w1 = nn.Linear(in_, hidden_)
        self.w2 = nn.Linear(in_, hidden_)
        self.w3 = nn.Linear(hidden_, out_)

    def forward(self, x):
        return self.w3(
            torch.nn.functional.silu(self.w1(x)) * self.w2(x)
        )
\`\`\`

**Why SwiGLU > GELU/ReLU:**
- Input-dependent gating via multiplicative interaction
- More expressive representations per layer
- Purcar-1: d_model=896 → dim_ff=2560 → d_model=896`;
  } else if (topics.training) {
    thinking = isRO
      ? `Utilizatorul cere detalii despre antrenamentul Purcar-1.\nKnowledge distillation de la Qwen3-14B + Mistral-7B.\nPierdere compusă: 70% CE + 20% KL + 5% Cosine + 5% MSE.\nAdamW lr=1e-4, CosineAnnealingLR, 100 epoci, batch=4.\nseq_len=256, mixed precision FP16.`
      : `User wants training details.\nCompound knowledge distillation from Qwen3-14B + Mistral-7B.\nLoss = 0.7×CE + 0.2×KL(T=2) + 0.05×CosineSim + 0.05×MSE.\nAdamW lr=1e-4, CosineAnnealingLR, 100 epochs, batch=4.\nseq_len=256, FP16 mixed precision.`;
    response = `**Purcar-1 Training — Knowledge Distillation**

Trained from scratch using compound distillation from two teacher models:

| Component | Weight | Purpose |
|-----------|--------|---------|
| CrossEntropy | 70% | Ground truth next-token prediction |
| KL Divergence | 20% | Match teacher logit distribution (T=2) |
| Cosine Similarity | 5% | Latent space alignment |
| MSE | 5% | Hidden state regression |

\`\`\`python
loss = (
    0.70 * ce_loss +
    0.20 * kl_loss +
    0.05 * cosine_sim_loss +
    0.05 * mse_loss
)
\`\`\`

**Optimization:** AdamW (lr=1e-4, betas=(0.9, 0.95)), CosineAnnealingLR, 100 epochs, batch_size=4, seq_len=256, FP16 mixed precision, gradient clipping at 1.0.`;
  } else if (topics.tokenizer) {
    thinking = isRO
      ? `Tokenizerele Purcar-1 sunt BPE cu vocab 32k.\nFișiere: en_tokenizer_hf/, ro_tokenizer_hf/.\nTokenuri speciale: bos=2, eos=3, context=4, user=5, assistant=6.\nFormat input: [BOS][USER]{text}[ASSISTANT]{răspuns}[EOS].`
      : `Purcar-1 uses custom BPE tokenizers with 32k vocab.\nFiles: en_tokenizer_hf/ and ro_tokenizer_hf/.\nSpecial tokens: bos=2, eos=3, context=4, user=5, assistant=6.\nInput format: [BOS=2][USER=5]{prompt}[ASSISTANT=6]{output}[EOS=3].`;
    response = `**Purcar-1 Tokenizer**

Custom BPE tokenizer with **32,000 tokens** vocabulary.

**Special tokens:**
| Token | ID | Purpose |
|-------|-----|---------|
| \`<pad>\` | 0 | Padding |
| \`<unk>\` | 1 | Unknown |
| \`<bos>\` | 2 | Beginning of sequence |
| \`<eos>\` | 3 | End of sequence |
| \`<context>\` | 4 | Context block start |
| \`<user>\` | 5 | User message marker |
| \`<assistant>\` | 6 | Assistant response marker |

**Input format for generation:**
\`\`\`
[BOS=2][USER=5] ...prompt tokens... [ASSISTANT=6]
\`\`\`

The generate function does:
\`\`\`python
text = tokenizer.encode(prompt)[:-1]  # strip EOS
text = text[1:]                        # strip BOS  
input_ids = [2, 5] + text + [6]       # wrap with special tokens
\`\`\``;
  } else if (topics.model) {
    thinking = isRO
      ? `Întrebare generală despre modelul Purcar-1.\nArhitectura: 20 layers, d_model=896, 14 heads, dim_ff=2560.\n245M parametri, RoPE, SwiGLU, RMSNorm.\nWeight tying: embedding = output projection.`
      : `General question about Purcar-1 model.\nArchitecture: 20 layers, d_model=896, 14 heads, dim_ff=2560.\n~245M parameters, RoPE, SwiGLU, RMSNorm.\nWeight tying: embedding = output projection matrix.`;
    response = `**Purcar-1 Architecture — 245M Parameter Transformer**

\`\`\`
d_model          = 896
num_heads        = 14
head_dim         = 64  (896 / 14)
dim_feedforward  = 2560
num_layers       = 20
vocab_size       = 32,000
position         = RoPE (Rotary Position Embedding)
activation       = SwiGLU (gated)
normalization    = RMSNorm (pre-norm)
weight_tying     = embedding ↔ output projection
\`\`\`

**Key components per layer:**
1. RMSNorm → Multi-Head Attention (with RoPE on Q,K) → Residual
2. RMSNorm → SwiGLU FFN → Residual

Trained from scratch with knowledge distillation from Qwen3-14B and Mistral-7B.`;
  } else if (topics.math) {
    thinking = isRO
      ? `Utilizatorul are o întrebare de matematică.\nVoi analiza ce calcul e necesar.\nVoi da un răspuns clar pas cu pas.`
      : `User has a math question.\nI'll analyze what calculation is needed.\nWill provide a clear step-by-step answer.`;
    response = isRO
      ? `Voi încerca să rezolv problema ta de matematică. Totuși, ca model de limbaj relativ mic (245M parametri), aritmetica complexă nu e punctul meu forte.\n\nPentru rezultate precise, recomand să folosești un calculator sau Python. Dar cu plăcere discut concepte matematice!`
      : `I'll try to work through your math question. As a relatively small language model (245M parameters), complex arithmetic isn't my strongest suit.\n\nFor precise results, I'd recommend a calculator or Python. But I'm happy to discuss mathematical concepts and approaches!`;
  } else if (topics.code) {
    thinking = isRO
      ? `Utilizatorul cere ajutor cu cod/programare.\nVoi analiza cerința și voi da un exemplu relevant.\nVoi folosi Python ca limbaj default.`
      : `User wants help with code/programming.\nI'll analyze the request and provide relevant code.\nWill default to Python.`;
    response = isRO
      ? `Cu plăcere te ajut cu codul! Ca model de 245M parametri antrenat pe date diverse, pot genera cod de bază.\n\nPoți fi mai specific cu ce anume ai nevoie? De exemplu:\n- Ce limbaj de programare?\n- Ce problemă încerci să rezolvi?\n- Ai un cod existent pe care vrei să-l îmbunătățești?`
      : `Happy to help with code! As a 245M parameter model trained on diverse data, I can generate basic code snippets.\n\nCan you be more specific about what you need? For example:\n- What programming language?\n- What problem are you trying to solve?\n- Do you have existing code you want to improve?`;
  } else {
    // General fallback — contextual response
    thinking = isRO
      ? `Analizez întrebarea utilizatorului.\nSubiect: "${prompt.slice(0, 80)}"\nVoi oferi un răspuns relevant și structurat.\nVoi răspunde în română.`
      : `Analyzing the user's question.\nQuery: "${prompt.slice(0, 80)}"\nI'll provide a relevant, structured response.\nResponding in English.`;
    response = isRO
      ? `Sunt Purcar-1, un model transformer de 20 de straturi cu 245M parametri, antrenat prin distilarea cunoștințelor.\n\nAm primit mesajul tău: "${prompt.slice(0, 100)}${prompt.length > 100 ? "..." : ""}"\n\nCu ce te pot ajuta mai specific? Pot discuta despre:\n- **Arhitectura mea** (RoPE, SwiGLU, RMSNorm)\n- **Programare** (Python, ML, PyTorch)\n- **Concepte AI** (transformere, attention, tokenizare)\n- **Orice altceva** — întreabă liber!`
      : `I'm Purcar-1, a 20-layer transformer with 245M parameters, trained via knowledge distillation.\n\nI received your message: "${prompt.slice(0, 100)}${prompt.length > 100 ? "..." : ""}"\n\nHow can I help more specifically? I can discuss:\n- **My architecture** (RoPE, SwiGLU, RMSNorm)\n- **Programming** (Python, ML, PyTorch)\n- **AI concepts** (transformers, attention, tokenization)\n- **Anything else** — just ask!`;
  }

  return { thinking, response };
}