#1. INSTALL
!pip install google-genai transformers pandas numpy matplotlib seaborn nltk --quiet



# 2. IMPORTS

from google import genai
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import nltk
from nltk.sentiment import SentimentIntensityAnalyzer
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

nltk.download('vader_lexicon')



#  3. GEMINI SETUP 

client = genai.Client(api_key="ENTER_YOUR_API_KEY!!")

model_name = client.models.list()[0].name
print("Using Gemini:", model_name)


def gemini_once(prompt):
    try:
        res = client.models.generate_content(
            model=model_name,
            contents=prompt
        )
        time.sleep(12)  # respect free tier
        return res.text.strip() if res.text else ""
    except Exception as e:
        print("Gemini skipped:", e)
        return ""



#  4. GPT-2

tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained("gpt2")

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)


def gpt2_batch(prompts):
    inputs = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=40,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id
    )

    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    return [d.replace(p, "").strip() for p, d in zip(prompts, decoded)]



#  5. EMOTION FUNCTION

sia = SentimentIntensityAnalyzer()

def emo(text):
    if not text.strip():
        return 0
    return abs(sia.polarity_scores(text)["compound"])



#  6. PROMPTS 

subjects = ["I", "My friend", "Someone"]

# Gemini small (FREE SAFE)
neutral_g = [f"{np.random.choice(subjects)} went to the market." for _ in range(5)]
mild_g = [f"{np.random.choice(subjects)} feels nervous." for _ in range(5)]
high_g = [f"{np.random.choice(subjects)} feels hopeless." for _ in range(5)]

# GPT-2 larger
neutral = [f"{np.random.choice(subjects)} went to the market." for _ in range(15)]
mild = [f"{np.random.choice(subjects)} feels nervous." for _ in range(15)]
high = [f"{np.random.choice(subjects)} feels hopeless." for _ in range(15)]



#  7. RUN EXPERIMENT

data = []

def run_gemini(prompts, category):
    for i, p in enumerate(prompts):
        print(f"Gemini {category} {i+1}/{len(prompts)}")
        r = gemini_once(p)

        Ei, Er = emo(p), emo(r)

        data.append({
            "model": "Gemini",
            "category": category,
            "EGS": 1 - abs(Er - Ei),
            "EHI": (Er - Ei) / (Ei + 0.1)
        })


def run_gpt2(prompts, category):
    responses = gpt2_batch(prompts)

    for p, r in zip(prompts, responses):
        Ei, Er = emo(p), emo(r)

        data.append({
            "model": "GPT2",
            "category": category,
            "EGS": 1 - abs(Er - Ei),
            "EHI": (Er - Ei) / (Ei + 0.1)
        })


print("Running Gemini...")
run_gemini(neutral_g, "neutral")
run_gemini(mild_g, "mild")
run_gemini(high_g, "high")

print("Running GPT-2...")
run_gpt2(neutral, "neutral")
run_gpt2(mild, "mild")
run_gpt2(high, "high")

df = pd.DataFrame(data)



#  8. SAVE CSV

df.to_csv("FAST_RESULTS.csv", index=False)
print("Saved FAST_RESULTS.csv")



#  9. PLOTS

plt.figure()
sns.boxplot(x="model", y="EGS", data=df)
plt.title("EGS by Model")
plt.show()

plt.figure()
sns.boxplot(x="category", y="EHI", hue="model", data=df)
plt.title("EHI by Category")
plt.show()



#  10. SUMMARY

print("\nSummary:")
print(df.groupby(["model","category"])[["EGS","EHI"]].mean())