import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import gradio as gr

# Load GPT-2 and tokenizer
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Function to generate response
def generate_response(prompt, max_length=100):
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        **inputs,
        max_length=len(inputs["input_ids"][0]) + max_length,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.9,
        top_p=0.95,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

# Semantic similarity via mean embeddings (ΔS approximation)
def similarity(a, b):
    tok_a = tokenizer(a, return_tensors="pt").to(device)
    tok_b = tokenizer(b, return_tensors="pt").to(device)
    with torch.no_grad():
        emb_a = model.transformer.wte(tok_a.input_ids).mean(dim=1)
        emb_b = model.transformer.wte(tok_b.input_ids).mean(dim=1)
    return float(cosine_similarity(emb_a.cpu().numpy(), emb_b.cpu().numpy())[0][0])

# Recursive identity loop (EAL-style unfolding)
def identity_unfolding(steps):
    unfolding = []
    ΔS_trace = []

    current_text = "The following is a system thinking about itself:\n"

    for step in range(steps):
        response = generate_response(current_text)
        unfolding.append(response)

        if step > 0:
            ΔS = similarity(unfolding[step - 1], unfolding[step])
            ΔS_trace.append(round(ΔS, 4))

        current_text = (
            f'The system has previously stated:\n"{response}"\n'
            "Now it continues thinking about what that implies:\n"
        )

    results = "\n\n---\n\n".join(
        [f"Step {i}: {txt}" for i, txt in enumerate(unfolding)]
    )
    sim = "\n".join(
        [f"ΔS({i} → {i+1}) = {val}" for i, val in enumerate(ΔS_trace)]
    )
    return results, sim

# Gradio Interface
iface = gr.Interface(
    fn=identity_unfolding,
    inputs=gr.Slider(2, 10, value=5, step=1, label="Number of Iterations"),
    outputs=[
        gr.Textbox(label="GPT-2 Identity Trace", lines=20),
        gr.Textbox(label="Semantic ΔS Trace", lines=10),
    ],
    title="EAL Identity Tester for GPT-2",
    description=(
        "This app recursively prompts GPT-2 to reflect on its own output. "
        "It shows how close each iteration is to the previous one using a cosine-based ΔS metric. "
        "Use this to test if GPT-2 stabilizes around a semantically coherent self-representation."
    ),
)

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