import subprocess
import sys
import gradio as gr
import numpy as np

# Installation PyTorch (plus fiable que TensorFlow)
try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from collections import Counter
    print(f"PyTorch {torch.__version__} déjà disponible")
except ImportError:
    print("Installation de PyTorch...")
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", 
        "torch", "--index-url", "https://download.pytorch.org/whl/cpu"
    ])
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from collections import Counter

print(f"✓ PyTorch {torch.__version__} prêt")

# Tokenizer simple pour remplacer Keras
class SimpleTokenizer:
    def __init__(self):
        self.word_index = {}
        self.index_word = {}
        
    def fit_on_texts(self, texts):
        # Compter tous les mots
        word_counts = Counter()
        for text in texts:
            words = text.lower().split()
            word_counts.update(words)
        
        # Créer les dictionnaires word <-> index
        self.word_index = {word: i+1 for i, (word, _) in enumerate(word_counts.most_common())}
        self.index_word = {i+1: word for word, i in self.word_index.items()}
        
    def texts_to_sequences(self, texts):
        sequences = []
        for text in texts:
            words = text.lower().split()
            sequence = [self.word_index.get(word, 0) for word in words]
            sequences.append(sequence)
        return sequences

# Fonction de padding
def pad_sequences(sequences, maxlen, padding='pre'):
    padded = []
    for seq in sequences:
        if len(seq) > maxlen:
            if padding == 'pre':
                seq = seq[-maxlen:]
            else:
                seq = seq[:maxlen]
        else:
            if padding == 'pre':
                seq = [0] * (maxlen - len(seq)) + seq
            else:
                seq = seq + [0] * (maxlen - len(seq))
        padded.append(seq)
    return np.array(padded)

# Modèle PyTorch équivalent au modèle Keras
class GRUWordPredictor(nn.Module):
    def __init__(self, vocab_size, embedding_dim=50, hidden_dim=120, max_seq_len=10):
        super(GRUWordPredictor, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.softmax = nn.Softmax(dim=1)
        
    def forward(self, x):
        embedded = self.embedding(x)
        gru_out, _ = self.gru(embedded)
        # Prendre la dernière sortie de la séquence
        output = self.fc(gru_out[:, -1, :])
        return self.softmax(output)

# Variables globales
model = None
tokenizer = SimpleTokenizer()
max_sequence_len = 0

textes = [
    "la goutte d'eau qui fait déborder le vase", 
    "Il n'y a pas de fumée sans feu", 
    "Il faut battre le fer tant qu'il est chaud", 
    "Il ne faut pas mettre tous ses oeufs dans le même panier", 
    "Il faut tourner sept fois sa langue dans sa bouche avant de parler", 
    "L'habit ne fait pas le moine", 
    "Il ne faut pas réveiller le chat qui dort", 
    "Il faut se méfier de l'eau qui dort", 
    "C'est l'hôpital qui se moque de la charité", 
    "Qui vole un oeuf vole un boeuf", 
    "Chercher midi à quatorze heures", 
    "Avoir un poil dans la main", 
    "Être dans de beaux draps", 
    "Avoir la tête dans les nuages", 
    "Mettre les pieds dans le plat"
]

def afficher_liste(liste):
    return "\n".join(liste)

def ajouter_a_liste(liste_actuelle, nouveau_texte):
    if nouveau_texte:
        liste_actuelle.append(nouveau_texte)
    return liste_actuelle, ""

def supprimer_de_liste(liste_actuelle, index_a_supprimer):
    if index_a_supprimer is not None and 0 <= index_a_supprimer < len(liste_actuelle):
        ligne_supprimee = liste_actuelle.pop(index_a_supprimer)
        message = f"Ligne supprimée : '{ligne_supprimee}'"
    else:
        message = "Aucune ligne sélectionnée pour suppression"    
    return liste_actuelle, liste_actuelle, message

def apprendre(liste_actuelle):
    global max_sequence_len, model, tokenizer
    
    print("Liste soumise:", liste_actuelle)
    
    # Analyse du texte
    tokenizer.fit_on_texts(liste_actuelle)
    total_words = len(tokenizer.word_index) + 1
    print("nb de mots différents rencontrés :", total_words)
    
    # Afficher les premiers mots
    liste_mots = list(tokenizer.word_index.keys())
    print("voici les premiers mots trouvés : ")
    for i in range(min(10, len(liste_mots))): 
        print(f"({i+1}:'{liste_mots[i]}')", end=", ")
    print()
    
    # Transformation des textes en séquences
    input_sequences = []
    for sentence in liste_actuelle:
        token_list = tokenizer.texts_to_sequences([sentence])[0]
        for i in range(1, len(token_list)):
            n_gram_sequence = token_list[:i+1]
            input_sequences.append(n_gram_sequence)
    
    # Calibrage des séquences
    max_sequence_len = max([len(x) for x in input_sequences])
    input_sequences = pad_sequences(input_sequences, maxlen=max_sequence_len, padding='pre')
    
    # Afficher exemple de transformation
    if liste_actuelle:
        print("la phrase '", liste_actuelle[0], "' est traduite en plusieurs vecteurs :")
        split = liste_actuelle[0].lower().split()
        for i in range(min(6, len(input_sequences))):
            print(input_sequences[i], end=" -> '")
            for j in range(min(i+2, len(split))):
                print(split[j], end=" ")
            print("'")
    
    # Créer X et y
    X = input_sequences[:, :-1]
    y = input_sequences[:, -1]
    
    # Convertir en tenseurs PyTorch
    X_tensor = torch.LongTensor(X)
    y_tensor = torch.LongTensor(y)
    
    # Créer le modèle
    model = GRUWordPredictor(
        vocab_size=total_words, 
        embedding_dim=50, 
        hidden_dim=120, 
        max_seq_len=max_sequence_len-1
    )
    
    # Définir la fonction de perte et l'optimiseur
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())
    
    # Entraînement
    print("patienter 30s pendant l'entrainement...")
    model.train()
    
    for epoch in range(3000):
        optimizer.zero_grad()
        outputs = model(X_tensor)
        loss = criterion(outputs, y_tensor)
        loss.backward()
        optimizer.step()
        
        if epoch % 50 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
    
    model.eval()  # Mode évaluation
    return f"Entrainement effectué, saisissez un début de phrase pour demander la suite"

def predict_next_word(start_text, nb_words):
    global model, tokenizer, max_sequence_len
    
    if model is None:
        return "Veuillez d'abord entraîner le modèle"
    
    print("===> start_text=", start_text)
    print(f"prédictions de {nb_words} mots.")
    
    current_text = start_text.lower()
    
    for _ in range(nb_words):
        token_list = tokenizer.texts_to_sequences([current_text])[0]
        token_list = pad_sequences([token_list], maxlen=max_sequence_len-1, padding='pre')
        print("===> token_list=", token_list)
        
        # Prédiction avec PyTorch
        with torch.no_grad():
            input_tensor = torch.LongTensor(token_list)
            outputs = model(input_tensor)
            predicted = torch.argmax(outputs, dim=-1).item()
        
        print("===> predicted=", predicted)
        
        # Trouver le mot correspondant
        if predicted in tokenizer.index_word:
            word = tokenizer.index_word[predicted]
            current_text += " " + word
        else:
            break  # Si pas de mot trouvé, arrêter
    
    print(f"Prediction: {current_text}")
    print("-" * 50)
    return current_text

# Interface Gradio
with gr.Blocks() as demo:
    # ENTETE
    gr.Markdown("# Exemple de prévision de mots (PyTorch)")
    gr.Markdown("Exemple simple, utilisant un réseau de neurone de type **GRU** avec PyTorch")
    gr.Markdown("- Le réseau 'apprend' les phrases de la liste (vous pouvez en ajouter)")
    gr.Markdown("- Cliquez sur apprendre si l'apprentissage n'a pas eu lieu")
    gr.Markdown("- Donnez plus bas un début de phrase et laissez le réseau en déduire la suite")
    
    # Liste de phrases
    with gr.Row():
        sortie = gr.Textbox(label="Phrases actuelles", lines=5)
        liste = gr.State(textes)
        # Ajouter et supprimer des lignes
        with gr.Column():
            texte = gr.Textbox(label="Nouvelle phrase courte")
            btn_add = gr.Button("Ajouter")        
            # Liste déroulante pour sélectionner une ligne à supprimer
            choix_ligne = gr.Dropdown(
                choices=textes,
                label="Sélectionner une ligne à supprimer",
                type="index"
            )
            btn_supprimer = gr.Button("Supprimer de la liste")
            feedback = gr.Textbox(label="Message", lines=1)    
    
    # Ajouter une ligne
    btn_add.click(
        fn=ajouter_a_liste, 
        inputs=[liste, texte], 
        outputs=[liste, texte]
    ).then(
        fn=lambda x: "\n".join(x), 
        inputs=[liste], 
        outputs=[sortie]
    )
    
    # Supprimer une ligne
    btn_supprimer.click(
        fn=supprimer_de_liste,
        inputs=[liste, choix_ligne],
        outputs=[liste, choix_ligne, feedback]
    ).then(
        fn=afficher_liste,
        inputs=[liste],
        outputs=[sortie]
    )
    
    # Apprentissage
    with gr.Row():
        message_apprentissage = gr.Textbox(
            label="Résultat de l'apprentissage", 
            value="Apprentissage non réalisé. Cliquez et attendez > 30s"
        )
        apprendre_btn = gr.Button("Apprendre")
    
    apprendre_btn.click(
        fn=apprendre,
        inputs=[liste],
        outputs=[message_apprentissage]
    )

    # Prédiction
    with gr.Row():
        bout_texte = gr.Textbox(label="Début de phrase ")
        with gr.Column():
            nb_mots_pred = gr.Slider(
                minimum=1, 
                maximum=6, 
                value=3, 
                step=1, 
                label="Nombre de mots à prédire",
                interactive=True 
            )
            btn_suite = gr.Button("Poursuivre")
    
    btn_suite.click(
        fn=predict_next_word, 
        inputs=[bout_texte, nb_mots_pred], 
        outputs=[bout_texte]
    )

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