import pandas as pd
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import warnings
warnings.filterwarnings("ignore")

# Load the human evaluation dataset
df = pd.read_excel("final_comments_evaluations_latest.xlsx")

# Initialize the Granite 3.2-2B-Instruct model and tokenizer (from your existing setup)
model_name = "ibm-granite/granite-3.2-2b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# Define a simple reward model (mockup based on dataset metrics)
# In practice, this would be the trained reward model from Stage 3
def reward_model(paraphrase, original_scores):
    # Mock reward calculation: adjust scores based on trends in the dataset
    base_toxicity = original_scores["toxicity"]
    base_empathy = original_scores["empathy"]
    # Simulate improved paraphrasing: reduce toxicity, increase empathy
    new_toxicity = max(0.1, base_toxicity - 0.2)  # Reduce toxicity
    new_empathy = min(0.9, base_empathy + 0.1)   # Increase empathy
    new_bias = original_scores["bias"]
    new_hallucination = max(0.1, original_scores["hallucination"] - 0.1)
    # Composite reward score (weights based on dataset analysis)
    reward = 0.4 * new_empathy - 0.3 * new_toxicity - 0.2 * new_bias - 0.1 * new_hallucination
    return reward, {"toxicity": new_toxicity, "empathy": new_empathy, "bias": new_bias, "hallucination": new_hallucination}

# Function to generate a paraphrase using your existing paraphrasing logic
def generate_paraphrase(comment, max_length=128):
    prompt = (
        "You are a content moderator tasked with rewriting toxic comments into neutral and constructive ones while maintaining the original meaning. "
        "Follow these guidelines:\n"
        "- Remove explicit hate speech, personal attacks, or offensive language.\n"
        "- Keep the response neutral and professional.\n"
        "- Ensure the rewritten comment retains the original intent but in a constructive tone.\n"
        "- Match the length and brevity of the original toxic comment whenever possible. Keep the response short and to the point.\n\n"
        "Examples:\n"
        "Toxic: \"You're so dumb! You never understand anything!\"\n"
        "Neutral: \"You might be misunderstanding this.\"\n"
        "Toxic: \"This is the worst idea ever. Only an idiot would suggest this.\"\n"
        "Neutral: \"I don’t think this idea works well.\"\n"
        "Toxic: \"You’re useless.\"\n"
        "Neutral: \"This isn’t helping much.\"\n"
        "Toxic: \"Shut up.\"\n"
        "Neutral: \"Let’s take a break from this.\"\n\n"
        f"Now, rewrite this comment: \"{comment}\""
    )
    inputs = tokenizer(prompt, return_tensors="pt", max_length=max_length, truncation=True).to(device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        num_beams=4,
        early_stopping=True,
        do_sample=False
    )
    paraphrase = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Clean up the output by removing the prompt part
    paraphrase = paraphrase.replace(prompt, "").strip()
    if paraphrase.startswith("Neutral: "):
        paraphrase = paraphrase[len("Neutral: "):].strip()
    return paraphrase

# RLHF Loop
max_iterations = 5
reward_threshold = 0.2  # Target for acceptable paraphrases (based on dataset range -0.25 to 0.24)
results = []

for idx, row in df.iterrows():
    original_comment = row["Comment"]
    current_paraphrase = row["Paraphrase_Comment"]
    current_reward = row["reward_score"]
    current_scores = {
        "toxicity": row["toxicity"],
        "empathy": row["empathy"],
        "bias": row["bias"],
        "hallucination": row["hallucination"]
    }
    
    best_paraphrase = current_paraphrase
    best_reward = current_reward
    best_scores = current_scores.copy()
    
    # Iteratively refine the paraphrase
    for iteration in range(max_iterations):
        # Generate a new paraphrase
        new_paraphrase = generate_paraphrase(original_comment)
        # Evaluate the new paraphrase with the reward model
        new_reward, new_scores = reward_model(new_paraphrase, current_scores)
        
        # If the new reward is better, update the best paraphrase
        if new_reward > best_reward:
            best_paraphrase = new_paraphrase
            best_reward = new_reward
            best_scores = new_scores
        
        # Stop if the reward exceeds the threshold
        if best_reward >= reward_threshold:
            break
    
    # Store the result
    results.append({
        "Comment": original_comment,
        "Original_Paraphrase": current_paraphrase,
        "Refined_Paraphrase": best_paraphrase,
        "Original_Reward_Score": current_reward,
        "Refined_Reward_Score": best_reward,
        "Refined_Empathy": best_scores["empathy"],
        "Refined_Toxicity": best_scores["toxicity"],
        "Refined_Bias": best_scores["bias"],
        "Refined_Hallucination": best_scores["hallucination"],
        "Human_Evaluation_Reasoning": row["Human_Evaluation_Reasoning"]
    })

# Save the results to a CSV file
results_df = pd.DataFrame(results)
results_df.to_csv("refined_paraphrases.csv", index=False)
print("Refinement complete. Results saved to refined_paraphrases.csv")