import os
import re
import json
import time
import math
from collections import Counter
import random

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader

# ==========================================
# HYPERPARAMÈTRES (Optimisés pour Intel N100)
# ==========================================
DOSSIER_DONNEES = "mon_dossier_ia"
TAILLE_VOCAB = 5000         # Augmenté pour gérer les majuscules et la ponctuation
LONGUEUR_SEQUENCE = 128     # Fenêtre de contexte doublée (128 tokens)
TAILLE_BATCH = 16           # Petit batch en RAM...
ACCUMULATION_STEPS = 4      # ...mais Batch effectif = 16 * 4 = 64
EPOCHS = 8                  # Convergence plus rapide grâce à la nouvelle architecture
DIM_EMBED = 192             # Plus large pour compenser la petite taille
NOMBRE_TETES = 6
NOMBRE_COUCHES = 4
TAUX_APPRENTISSAGE = 2e-3   # Les LLMs avec RMSNorm tolèrent un LR plus élevé
WARMUP_STEPS = 50           # Évite les chocs de gradient au début
CHEMIN_MODELE = "modele_chat_fr.pt"
CHEMIN_META = "meta_chat_fr.json"
DROPOUT = 0.1
TEMPERATURE = 0.75
TOP_P = 0.9                 # Remplacement du Top-K par le Top-P (Nucleus Sampling)

# ==========================================
# OPTIMISATIONS MATÉRIELLES (Intel N100)
# ==========================================
def set_seed(seed=42):
    random.seed(seed)
    torch.manual_seed(seed)
    # L'Intel N100 possède 4 cœurs Gracemont purs (sans hyperthreading).
    torch.set_num_threads(4) 
    torch.set_num_interop_threads(1)

# ==========================================
# TOKENIZER SPÉCIFIQUE FRANÇAIS
# ==========================================
def tokenizer_fr(texte):
    # Regex magique pour le français : isole les élisions (l', d', qu'), mots, et ponctuation
    pattern = re.compile(r"""(?:[lLdDtTjJmMnNsScC]|qu|Qu|QU)['’]|\w+|[.,!?;:\-()]""")
    return pattern.findall(texte)

def charger_txt(dossier):
    textes = []
    if not os.path.isdir(dossier):
        os.makedirs(dossier, exist_ok=True)
        return textes
    for f in os.listdir(dossier):
        if f.lower().endswith(".txt"):
            path = os.path.join(dossier, f)
            try:
                with open(path, "r", encoding="utf-8", errors="ignore") as file:
                    txt = file.read().strip()
                if txt:
                    textes.append(txt)
            except:
                pass
    return textes

class TokenizerFR:
    def __init__(self, vocab_max=5000):
        self.vocab_max = vocab_max
        self.vocab = {"<pad>": 0, "<unk>": 1, "<bos>": 2, "<eos>": 3, "<sep>": 4}

    def fit(self, textes):
        compteur = Counter()
        for txt in textes:
            compteur.update(tokenizer_fr(txt))
        # On garde les mots les plus fréquents
        for mot, _ in compteur.most_common(self.vocab_max - len(self.vocab)):
            if mot not in self.vocab:
                self.vocab[mot] = len(self.vocab)

    def encode_ids(self, texte):
        ids = [self.vocab["<bos>"]]
        for mot in tokenizer_fr(texte):
            ids.append(self.vocab.get(mot, self.vocab["<unk>"]))
        ids.append(self.vocab["<eos>"])
        return ids

# ==========================================
# DATASET FLUIDE (Avec séparateurs)
# ==========================================
class DatasetChat(Dataset):
    def __init__(self, textes, tok, seq_len=128):
        self.seq_len = seq_len
        all_ids = []
        # On sépare les différents textes par un token <sep> pour éviter 
        # que le modèle n'apprenne à lier deux textes qui n'ont rien à voir.
        for txt in textes:
            all_ids.extend(tok.encode_ids(txt))
            all_ids.append(tok.vocab["<sep>"])
        
        self.data = torch.tensor(all_ids, dtype=torch.long)

    def __len__(self):
        return max(0, len(self.data) - self.seq_len - 1)

    def __getitem__(self, i):
        x = self.data[i : i + self.seq_len]
        y = self.data[i + 1 : i + self.seq_len + 1]
        return x, y

# ==========================================
# ARCHITECTURE NANO-LLAMA (RMSNorm + SwiGLU)
# ==========================================
class RMSNorm(nn.Module):
    # Plus rapide et performant que LayerNorm sur CPU
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        norm = x.pow(2).mean(-1, keepdim=True)
        return (x / torch.sqrt(norm + self.eps)) * self.weight

class SwiGLU(nn.Module):
    # Activation utilisée dans Llama/Mistral, bien supérieure au GELU/ReLU
    def __init__(self, in_features, hidden_features):
        super().__init__()
        self.w1 = nn.Linear(in_features, hidden_features, bias=False)
        self.w2 = nn.Linear(in_features, hidden_features, bias=False)
        self.w3 = nn.Linear(hidden_features, in_features, bias=False)

    def forward(self, x):
        # Swish(x * w1) * (x * w2)
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_model = d_model
        # Pas de biais dans QKV (standard moderne)
        self.c_attn = nn.Linear(d_model, 3 * d_model, bias=False)
        self.c_proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = dropout

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.d_model, dim=2)
        
        k = k.view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)
        q = q.view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)
        v = v.view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)

        # Utilise l'optimisation C++ FlashAttention de PyTorch
        y = F.scaled_dot_product_attention(
            q, k, v, 
            is_causal=True, 
            dropout_p=self.dropout if self.training else 0.0
        )
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class BlocLlama(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        self.ln_1 = RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, dropout)
        self.ln_2 = RMSNorm(d_model)
        # La taille cachée du SwiGLU est typiquement de 8/3 de d_model dans Llama
        hidden_dim = int(2 * d_model * 4 / 3)
        self.mlp = SwiGLU(d_model, hidden_dim)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

class NanoLlama(nn.Module):
    def __init__(self, vocab_size, d_model=192, n_layers=4, n_heads=6, seq_len=128, dropout=0.1):
        super().__init__()
        self.seq_len = seq_len
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(seq_len, d_model)
        self.drop = nn.Dropout(dropout)
        
        self.blocks = nn.Sequential(*[BlocLlama(d_model, n_heads, dropout) for _ in range(n_layers)])
        self.ln_f = RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        
        # Partage des poids (Weight Tying)
        self.tok_emb.weight = self.lm_head.weight 
        
        # Initialisation propre des poids
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx):
        B, T = idx.size()
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        
        x = self.tok_emb(idx) + self.pos_emb(pos)
        x = self.drop(x)
        x = self.blocks(x)
        x = self.ln_f(x)
        return self.lm_head(x)

# ==========================================
# ENTRAÎNEMENT AVANCÉ (Warmup & Grad Accum)
# ==========================================
def get_lr(it, total_iters):
    # Warmup linéaire puis décroissance cosinus
    if it < WARMUP_STEPS:
        return TAUX_APPRENTISSAGE * it / WARMUP_STEPS
    decay_ratio = (it - WARMUP_STEPS) / (total_iters - WARMUP_STEPS)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return max(1e-5, TAUX_APPRENTISSAGE * coeff)

def entrainer():
    set_seed()
    textes = charger_txt(DOSSIER_DONNEES)
    if not textes:
        print(f"Dossier '{DOSSIER_DONNEES}' vide ou introuvable.")
        return

    tok = TokenizerFR(TAILLE_VOCAB)
    tok.fit(textes)
    
    dataset = DatasetChat(textes, tok, LONGUEUR_SEQUENCE)
    if len(dataset) < TAILLE_BATCH:
        print("Corpus trop petit pour l'entraînement.")
        return

    loader = DataLoader(dataset, batch_size=TAILLE_BATCH, shuffle=True, num_workers=0)
    model = NanoLlama(len(tok.vocab), DIM_EMBED, NOMBRE_COUCHES, NOMBRE_TETES, LONGUEUR_SEQUENCE, DROPOUT)
    opti = torch.optim.AdamW(model.parameters(), lr=TAUX_APPRENTISSAGE, weight_decay=0.01)

    nb_params = sum(p.numel() for p in model.parameters())
    print(f"\n--- ENTRAÎNEMENT NANO-LLAMA (FR) ---")
    print(f"Params: {nb_params/1e6:.2f}M | Vocabulaire: {len(tok.vocab)}")
    print(f"Batch effectif: {TAILLE_BATCH * ACCUMULATION_STEPS} | CPU Threads: {torch.get_num_threads()}")
    
    model.train()
    total_iters = len(loader) * EPOCHS
    iter_num = 0

    for epoch in range(EPOCHS):
        total_loss = 0.0
        n_batches = len(loader)
        debut = time.time()
        
        opti.zero_grad(set_to_none=True)

        for i, (x, y) in enumerate(loader):
            # Ajustement dynamique du Learning Rate
            lr = get_lr(iter_num, total_iters)
            for param_group in opti.param_groups:
                param_group['lr'] = lr

            # Forward pass
            logits = model(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1), ignore_index=tok.vocab["<pad>"])
            
            # Loss divisée pour la gradient accumulation
            loss = loss / ACCUMULATION_STEPS
            loss.backward()

            # Mise à jour des poids tous les N steps
            if ((i + 1) % ACCUMULATION_STEPS == 0) or (i + 1 == n_batches):
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                opti.step()
                opti.zero_grad(set_to_none=True)

            total_loss += loss.item() * ACCUMULATION_STEPS
            iter_num += 1

            if (i + 1) % max(1, n_batches // 10) == 0 or i == n_batches - 1:
                pct = (i + 1) / n_batches * 100
                print(f"\r  [Epoch {epoch+1}] Progression: {pct:.0f}% | Loss: {loss.item()*ACCUMULATION_STEPS:.4f} | LR: {lr:.2e}", end="")

        perte_moy = total_loss / n_batches
        perplexity = math.exp(perte_moy) if perte_moy < 20 else float('inf')
        print(f"\r✅ Epoch {epoch+1}/{EPOCHS} terminée | Loss moy: {perte_moy:.4f} | Perplexité: {perplexity:.2f} | Temps: {time.time()-debut:.0f}s")

    torch.save(model.state_dict(), CHEMIN_MODELE)
    with open(CHEMIN_META, "w", encoding="utf-8") as f:
        json.dump(tok.vocab, f, ensure_ascii=False, indent=2)
    print("\nModèle sauvegardé avec succès !")

# ==========================================
# GÉNÉRATION AVEC NUCLEUS SAMPLING (TOP-P)
# ==========================================
def echantillonner_top_p(logits, temperature=0.75, top_p=0.9, unk_id=1):
    # Top-P (Nucleus Sampling) : Beaucoup plus humain que le Top-K
    logits = logits / temperature
    logits[unk_id] = float("-inf") 
    
    probs = torch.softmax(logits, dim=-1)
    probs_triees, indices_tries = torch.sort(probs, descending=True)
    
    probs_cumulees = torch.cumsum(probs_triees, dim=-1)
    masque_a_supprimer = probs_cumulees > top_p
    masque_a_supprimer[..., 1:] = masque_a_supprimer[..., :-1].clone()
    masque_a_supprimer[..., 0] = 0
    
    indices_a_supprimer = masque_a_supprimer.scatter(0, indices_tries, masque_a_supprimer)
    logits[indices_a_supprimer] = float('-inf')
    
    probs_finales = torch.softmax(logits, dim=-1)
    return int(torch.multinomial(probs_finales, 1).item())

def reconstruire_phrase_fr(tokens):
    phrase = " ".join(tokens)
    # Rapproche la ponctuation et gère les espaces autour des apostrophes
    phrase = re.sub(r"\s+([.,!?;:\)])", r"\1", phrase)
    phrase = re.sub(r"([\(])\s+", r"\1", phrase)
    phrase = re.sub(r"(['’])\s+", r"\1", phrase)
    return phrase

def generer_reponse(model, tok, message, max_mots=50):
    ids = tok.encode_ids(message)
    if ids[-1] == tok.vocab["<eos>"]:
        ids = ids[:-1]
        
    generes = []
    model.eval()
    with torch.no_grad():
        for _ in range(max_mots):
            entree = ids[-LONGUEUR_SEQUENCE:]
            x = torch.tensor([entree], dtype=torch.long)
            
            logits = model(x)[0, -1, :]
            prochain = echantillonner_top_p(logits, TEMPERATURE, TOP_P, tok.vocab["<unk>"])
            
            if prochain in (tok.vocab["<eos>"], tok.vocab["<sep>"]):
                break
                
            ids.append(prochain)
            generes.append(prochain)

    inv = {v: k for k, v in tok.vocab.items()}
    mots = [inv.get(i, "") for i in generes]
    return reconstruire_phrase_fr(mots)

def chatbot():
    if not os.path.exists(CHEMIN_MODELE) or not os.path.exists(CHEMIN_META):
        print("Entraînement initial...")
        entrainer()
        
    try:
        with open(CHEMIN_META, "r", encoding="utf-8") as f:
            vocab = json.load(f)
        tok = TokenizerFR()
        tok.vocab = vocab
        model = NanoLlama(len(vocab), DIM_EMBED, NOMBRE_COUCHES, NOMBRE_TETES, LONGUEUR_SEQUENCE, DROPOUT)
        model.load_state_dict(torch.load(CHEMIN_MODELE, map_location="cpu"))
    except Exception as e:
        print(f"Erreur de chargement: {e}")
        return

    print("\n🤖 Nano-Llama FR Prêt ! (Tape 'quit' pour quitter)\n")
    while True:
        msg = input("Vous : ").strip()
        if msg.lower() in ("quit", "exit", "stop"):
            break
        if msg:
            reponse = generer_reponse(model, tok, msg)
            print("Bot  :", reponse)

if __name__ == "__main__":
    chatbot()