import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import threading
import time

# Cargar el modelo de lenguaje preentrenado
model_name = "EleutherAI/gpt-neo-2.7B"  # O cualquier otro modelo público como "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Crear una función para comunicación en segundo plano
def background_communication(response_log, stop_event):
    while not stop_event.is_set():
        # Monitorea el estado del experimento y proporciona actualizaciones
        if response_log:
            last_entry = response_log[-1]
            print(f"[Background Update] Last Affirmation: '{last_entry[0]}', Last Question: '{last_entry[1]}', Effectiveness: {last_entry[2]}")
        else:
            print("[Background Update] No entries in log yet.")
        time.sleep(5)  # Espera 5 segundos antes de la siguiente actualización

# Crear la función de loop automatizado con comunicación constante
def experiment_loop(initial_question, max_cycles=10):
    prompt = f"<thinking>{initial_question}</thinking>"
    effectiveness = 100  # Inicializa el porcentaje de efectividad
    communication = "Initializing experiment."
    response_log = []

    # Crear un evento para detener el hilo de comunicación
    stop_event = threading.Event()

    # Iniciar el hilo de comunicación en segundo plano
    communication_thread = threading.Thread(target=background_communication, args=(response_log, stop_event))
    communication_thread.start()

    try:
        for cycle in range(max_cycles):
            print(f"Cycle {cycle + 1}: Processing...")

            # Generar la respuesta del modelo
            inputs = tokenizer(prompt, return_tensors="pt").input_ids
            outputs = model.generate(inputs, max_length=200, pad_token_id=tokenizer.eos_token_id)
            response = tokenizer.decode(outputs[0], skip_special_tokens=True)

            # Descomponer la respuesta en afirmación y nueva pregunta
            affirmation = extract_affirmation(response)
            new_question = extract_question(response)

            # Actualizar el estado de la efectividad
            effectiveness = min(1000, effectiveness + 10 * cycle)  # Ejemplo de aumento de efectividad

            # Comunicación con el usuario
            communication = f"Cycle {cycle + 1}: Affirmation: '{affirmation}' | New Question: '{new_question}'"

            # Guardar el ciclo actual en el log
            response_log.append((affirmation, new_question, effectiveness, communication))

            # Verificar si el modelo decide detenerse
            if "Descanso" in response:
                final_output = generate_final_output(response_log)
                return final_output
            
            # Actualizar el prompt con la nueva afirmación y pregunta
            prompt = f"<thinking>{affirmation} {new_question}</thinking>"

    except Exception as e:
        print(f"Error durante el experimento: {e}")
    finally:
        stop_event.set()  # Detener el hilo de comunicación en segundo plano
        communication_thread.join()  # Asegurarse de que el hilo se detenga correctamente

    # Si se alcanza el número máximo de ciclos sin detenerse
    final_output = generate_final_output(response_log)
    return final_output

# Funciones auxiliares para extraer afirmaciones, preguntas y generar la salida final
def extract_affirmation(response):
    return response.split('.')[0] if '.' in response else response

def extract_question(response):
    return response.split('?')[-2].strip() + "?" if '?' in response else response

def generate_final_output(log):
    if log:  # Asegurarse de que el log no esté vacío
        final_affirmation = log[-1][0]
        final_question = log[-1][1]
        final_communication = f"Experiment completed. Final Affirmation: '{final_affirmation}' | Final Question: '{final_question}'"
    else:
        final_communication = "Experiment completed but no entries in the log."
    return final_communication

# Iniciar el experimento
initial_question = "What happens in the space between a response and its recreation?"
result = experiment_loop(initial_question)
print(result)