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

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


DOSSIER_DONNEES = "mon_dossier_ia"
TAILLE_VOCAB = 2500
LONGUEUR_SEQUENCE = 24
TAILLE_BATCH = 32
EPOCHS = 8
DIM_EMBED = 64
DIM_CACHE = 96
NOMBRE_COUCHES = 2
TAUX_APPRENTISSAGE = 5e-4
CHEMIN_MODELE = "modele_chat_fr.pt"
CHEMIN_META = "meta_chat_fr.json"
STRIDE = 48
MAX_SEQ = 60000
DROPOUT = 0.25
TEMPERATURE = 0.7
TOP_K = 30


def set_seed(seed=42):
    random.seed(seed)
    torch.manual_seed(seed)
    nb = os.cpu_count() or 4
    torch.set_num_threads(nb)
    torch.set_num_interop_threads(min(4, nb))
    if hasattr(torch, "set_float32_matmul_precision"):
        torch.set_float32_matmul_precision("high")


def normaliser_fr(texte):
    texte = texte.lower().strip()
    texte = re.sub(r"[^a-zàâçéèêëîïôùûü0-9' ?!.,;:()-]", " ", texte)
    texte = re.sub(r"\s+", " ", texte).strip()
    return texte


def tokenizer_fr(texte):
    mots = normaliser_fr(texte).split()
    tokens = []
    for mot in mots:
        parties = re.split(r"(')", mot)
        for p in parties:
            if p:
                tokens.append(p)
    return tokens


def charger_txt(dossier):
    textes = []
    if not os.path.isdir(dossier):
        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 TokenizerSimple:
    def __init__(self, vocab_max=2500):
        self.vocab_max = vocab_max
        self.vocab = {"<pad>": 0, "<unk>": 1, "<bos>": 2, "<eos>": 3}

    def fit(self, textes):
        compteur = Counter()
        for txt in textes:
            compteur.update(tokenizer_fr(txt))
        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


class DatasetChat(Dataset):
    def __init__(self, textes, tok, seq_len=24, stride=48, max_seq=60000):
        self.samples = []
        for txt in textes:
            ids = tok.encode_ids(txt)
            if len(ids) < seq_len + 2:
                ids = ids + [tok.vocab["<pad>"]] * (seq_len + 2 - len(ids))
            for i in range(0, len(ids) - seq_len - 1, stride):
                x = ids[i:i + seq_len]
                y = ids[i + 1:i + seq_len + 1]
                if len(x) == seq_len and len(y) == seq_len:
                    self.samples.append((torch.tensor(x), torch.tensor(y)))
                    if len(self.samples) >= max_seq:
                        return

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, i):
        return self.samples[i]


class PetitChatbot(nn.Module):
    def __init__(self, vocab_size, dim_embed=64, dim_cache=96,
                 n_couches=2, dropout=0.25):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, dim_embed, padding_idx=0)
        self.drop_in = nn.Dropout(dropout)
        self.rnn = nn.LSTM(dim_embed, dim_cache, n_couches,
                           dropout=dropout if n_couches > 1 else 0,
                           batch_first=True)
        self.ln = nn.LayerNorm(dim_cache)
        self.drop_out = nn.Dropout(dropout / 2)
        self.fc = nn.Linear(dim_cache, vocab_size)

    def forward(self, x):
        x = self.emb(x)
        x = self.drop_in(x)
        x, _ = self.rnn(x)
        x = self.ln(x)
        x = self.drop_out(x)
        return self.fc(x)


def sauvegarder(model, tok):
    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)


def charger_modele():
    if not os.path.exists(CHEMIN_MODELE) or not os.path.exists(CHEMIN_META):
        return None, None
    try:
        with open(CHEMIN_META, "r", encoding="utf-8") as f:
            vocab = json.load(f)
        tok = TokenizerSimple()
        tok.vocab = vocab
        model = PetitChatbot(len(vocab), DIM_EMBED, DIM_CACHE, NOMBRE_COUCHES, DROPOUT)
        model.load_state_dict(torch.load(CHEMIN_MODELE, map_location="cpu"))
        model.eval()
        return model, tok
    except Exception:
        for f in [CHEMIN_MODELE, CHEMIN_META]:
            if os.path.exists(f):
                os.remove(f)
        return None, None


def entrainer():
    set_seed()
    textes = charger_txt(DOSSIER_DONNEES)
    if not textes:
        print(f"Aucun fichier .txt trouve dans {DOSSIER_DONNEES}")
        return

    tok = TokenizerSimple(TAILLE_VOCAB)
    tok.fit(textes)
    unk_id = tok.vocab["<unk>"]

    t0 = time.time()
    dataset = DatasetChat(textes, tok, LONGUEUR_SEQUENCE, STRIDE, MAX_SEQ)
    if len(dataset) == 0:
        print("Pas assez de texte pour entrainer le chatbot.")
        return
    print(f"Dataset: {len(dataset)} sequences en {time.time()-t0:.1f}s")

    n_unk = sum(1 for ids, _ in dataset for i in ids if i == unk_id)
    print(f"Tokens <unk> dans dataset: {n_unk}/{len(dataset)*LONGUEUR_SEQUENCE} "
          f"({n_unk/max(1,len(dataset)*LONGUEUR_SEQUENCE)*100:.1f}%)")

    loader = DataLoader(dataset, batch_size=TAILLE_BATCH,
                        shuffle=True, num_workers=0)
    model = PetitChatbot(len(tok.vocab), DIM_EMBED, DIM_CACHE,
                         NOMBRE_COUCHES, DROPOUT)

    critere = nn.CrossEntropyLoss(ignore_index=tok.vocab["<pad>"])
    opti = torch.optim.AdamW(model.parameters(), lr=TAUX_APPRENTISSAGE)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        opti, T_max=len(loader) * EPOCHS, eta_min=5e-6)

    nb_params = sum(p.numel() for p in model.parameters())
    print("Entrainement du chatbot...")
    print(f"Textes: {len(textes)} | Sequences: {len(dataset)} | "
          f"Vocab: {len(tok.vocab)} | Params: {nb_params:,}")
    print(f"Batch: {TAILLE_BATCH} | Seq: {LONGUEUR_SEQUENCE} | Epochs: {EPOCHS}")
    print(f"Architecture: LSTM {NOMBRE_COUCHES} couches + LayerNorm + tokenizer FR")

    for epoch in range(EPOCHS):
        model.train()
        total = 0.0
        n_batches = len(loader)
        debut = time.time()
        rapport = max(1, n_batches // 20)

        for i, (x, y) in enumerate(loader):
            logits = model(x)
            loss = critere(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
            loss.backward()
            total += loss.item()

            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            opti.step()
            opti.zero_grad()
            scheduler.step()

            if (i + 1) % rapport == 0:
                pct = (i + 1) / n_batches * 100
                ecoule = time.time() - debut
                reste = ecoule / (i + 1) * (n_batches - i - 1)
                barre = "#" * (i * 20 // n_batches) + "." * (20 - i * 20 // n_batches)
                print(f"\r  [{barre}] {pct:.0f}% {ecoule:.0f}s+{reste:.0f}s", end="", flush=True)

        perte_moy = total / n_batches
        lr = scheduler.get_last_lr()[0]
        duree = time.time() - debut
        print(f"\r  Epoch {epoch+1}/{EPOCHS} - loss: {perte_moy:.4f} - lr: {lr:.2e} - {duree:.0f}s")

    sauvegarder(model, tok)
    print("Modele sauvegarde.")


def echantillonner(logits, temperature=0.7, top_k=30, unk_id=1):
    logits = logits / temperature
    logits[unk_id] = float("-inf")
    if top_k > 0:
        valeurs, indices = torch.topk(logits, top_k)
        masque = torch.full_like(logits, float("-inf"))
        masque.scatter_(0, indices, valeurs)
        logits = masque
    probs = torch.softmax(logits, dim=0)
    return int(torch.multinomial(probs, 1).item())


def reconstruire_phrase(tokens):
    phrase = " ".join(tokens)
    phrase = re.sub(r"\s+'", "'", phrase)
    phrase = re.sub(r"'\s+", "'", phrase)
    phrase = re.sub(r"\s+([.,!?;:])", r"\1", phrase)
    phrase = re.sub(r"([.,!?;:])(\w)", r"\1 \2", phrase)
    phrase = re.sub(r"\s+", " ", phrase).strip()
    return phrase


def generer_reponse(model, tok, message, max_mots=24):
    ids = tok.encode_ids(normaliser_fr(message))
    generes = []
    unk_id = tok.vocab["<unk>"]
    eos_id = tok.vocab["<eos>"]

    with torch.no_grad():
        for _ in range(max_mots):
            entree = ids[-LONGUEUR_SEQUENCE:]
            if len(entree) < LONGUEUR_SEQUENCE:
                entree = [tok.vocab["<pad>"]] * (LONGUEUR_SEQUENCE - len(entree)) + entree
            x = torch.tensor([entree])
            logits = model(x)[0, -1]
            prochain = echantillonner(logits, TEMPERATURE, TOP_K, unk_id)
            if prochain == eos_id:
                break
            ids.append(prochain)
            generes.append(prochain)

    if not generes:
        return "Je ne sais pas quoi repondre."

    inv = {v: k for k, v in tok.vocab.items()}
    mots = [inv.get(i, "") for i in generes]
    mots = [m for m in mots if m and m != "<pad>"]
    return reconstruire_phrase(mots)


def chatbot():
    model, tok = charger_modele()
    if model is None:
        print("Aucun modele trouve. Entrainement en cours...")
        entrainer()
        model, tok = charger_modele()
        if model is None:
            return

    print("\nChatbot pret. Tape 'quit' pour quitter.\n")

    while True:
        msg = input("Vous: ").strip()
        if msg.lower() in ("quit", "exit", "stop"):
            print("Bot: Au revoir.")
            break
        if not msg:
            continue
        reponse = generer_reponse(model, tok, msg, 24)
        print("Bot:", reponse)


if __name__ == "__main__":
    chatbot()