import json
import numpy as np
import re

# On tente d'importer gplearn, sinon on fournit un message clair
try:
    from gplearn.genetic import SymbolicRegressor
    GPLEARN_AVAILABLE = True
except ImportError:
    GPLEARN_AVAILABLE = False

def parser_fichier_txt(chemin_fichier):
    """
    Parse un fichier .txt pour en extraire les paires (n, zero_t).
    Supporte les formats :
    - Lignes brutes : '2, 14.1347' ou '2 14.1347' ou '2;14.1347'
    - Format clé-valeur : 'n: 2, t: 14.1347'
    """
    X = []
    y = []
    
    with open(chemin_fichier, 'r', encoding='utf-8') as f:
        for num_ligne, ligne in enumerate(f, 1):
            ligne = ligne.strip()
            if not ligne or ligne.startswith('#'): # Ignore les lignes vides et commentaires
                continue
                
            # Extraction automatique de tous les nombres (entiers ou flottants) de la ligne
            nombres = re.findall(r"[-+]?\d*\.\d+|\d+", ligne)
            
            if len(nombres) >= 2:
                try:
                    n = float(nombres[0])
                    zero_t = float(nombres[1])
                    X.append([n])
                    y.append(zero_t)
                except ValueError:
                    print(f"[Attention] Ligne {num_ligne} ignorée (impossible de convertir) : {ligne}")
            else:
                print(f"[Attention] Ligne {num_ligne} ignorée (données insuffisantes) : {ligne}")
                
    return np.array(X), np.array(y)

# --- PROGRAMME PRINCIPAL ---
if __name__ == '__main__':
    chemin_txt = "zeros_riemann.txt"
    
    # 1. Création automatique d'un fichier .txt d'exemple pour la démonstration
    exemple_donnees = """# Fichier de zéros de Riemann pour l'IA
# Format : n (nombre premier ou index), valeur_t_du_zero
2, 14.134725
3, 21.022040
5, 25.010858
7, 30.424876
11, 32.935062
13, 37.586178
17, 40.918719
19, 43.327073
"""
    with open(chemin_txt, "w", encoding="utf-8") as f:
        f.write(exemple_donnees.strip())
    
    print(f"[Info] Fichier exemple '{chemin_txt}' créé pour la démonstration.\n")

    # 2. Chargement et parsing du fichier .txt
    print(f"Chargement des données depuis {chemin_txt}...")
    X_train, y_train = parser_fichier_txt(chemin_txt)
    
    if len(X_train) == 0:
        print("Erreur : Aucune donnée valide trouvée dans le fichier .txt.")
    else:
        print(f"Données chargées avec succès ! {len(X_train)} points trouvés.")
        print("Variables X (n) :", X_train.flatten())
        print("Cibles y (zeros) :", y_train)
        print("-" * 50)
        
        if not GPLEARN_AVAILABLE:
            print("\n[Outil IA] Pour exécuter la régression symbolique, installez gplearn :")
            print(">>> pip install gplearn")
            print("\nLe script a correctement traité ton .txt et attend que la bibliothèque soit installée.")
        else:
            # 3. Configuration et entraînement de l'IA
            print("\nL'IA analyse le fichier .txt pour découvrir l'équation de liaison...")
            
            # Fonctions mathématiques autorisées
            fonctions = ['add', 'sub', 'mul', 'div', 'log', 'sqrt', 'sin', 'cos']
            
            agent_ia = SymbolicRegressor(
                population_size=5000,
                generations=20,
                stopping_criteria=0.001,
                function_set=fonctions,
                p_crossover=0.7,
                p_subtree_mutation=0.1,
                p_hoist_mutation=0.05,
                p_point_mutation=0.1,
                verbose=1,
                random_state=42
            )
            
            agent_ia.fit(X_train, y_train)
            
            # 4. Affichage de la formule générée
            print("\n" + "="*50)
            print("FORMULE MATHÉMATIQUE DÉCOUVERTE PAR L'IA :")
            print(f"f(n) = {agent_ia._program}")
            print("="*50)
            
            # Test de prédiction sur un nombre premier absent du fichier (ex: n = 23)
            n_test = np.array([[23]])
            prediction = agent_ia.predict(n_test)
            print(f"Pour n=23, la formule de l'IA donne un zéro à t = {prediction[0]:.4f}")