import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM

class ProfileGenerator:
   def __init__(self):
       # Charger Open R1 pour la partie profiling
       self.tokenizer = AutoTokenizer.from_pretrained("open-r1")
       self.model = AutoModelForCausalLM.from_pretrained("open-r1")
       
       # Configuration des profils
       self.risk_levels = ["Prudent", "Modéré", "Dynamique"] 
       self.goals = ["Retraite", "Immobilier", "Études", "Épargne"]

   def generate_profile_data(self, num_samples=1000):
       data = []
       for _ in range(num_samples):
           # Générer profil de base
           profile = {
               "age": np.random.randint(18, 70),
               "income": np.random.randint(2000, 20000),
               "savings": np.random.randint(500, 5000),
               "goals": np.random.choice(self.goals, size=2, replace=False),
               "risk_tolerance": np.random.choice(self.risk_levels)
           }
           
           # Générer description avec Open R1
           profile["description"] = self._generate_profile_description(profile)
           
           data.append(profile)
           
       return pd.DataFrame(data)

   def _generate_profile_description(self, profile):
       # Préparation du prompt pour Open R1
       prompt = f"""
       Profil investisseur :
       - Age: {profile['age']} ans
       - Revenu: {profile['income']}€/mois 
       - Épargne: {profile['savings']}€/mois
       - Objectifs: {', '.join(profile['goals'])}
       - Tolérance au risque: {profile['risk_tolerance']}

       Décrivez ce profil d'investisseur et suggérez une approche personnalisée.
       """
       
       # Générer la description avec Open R1
       inputs = self.tokenizer(prompt, return_tensors="pt")
       outputs = self.model.generate(**inputs, max_length=200)
       return self.tokenizer.decode(outputs[0])

def main():
   # Générer données
   generator = ProfileGenerator()
   df = generator.generate_profile_data()
   
   # Sauvegarder
   df.to_csv("src/data/investment_profiles.csv", index=False)
   print("Données générées et sauvegardées")

if __name__ == "__main__":
   main()