import streamlit as st
import pandas as pd
import numpy as np
import random
import os
from datetime import datetime
import ast

# Custom CSS for enhanced styling
def local_css():
    st.markdown("""
    <style>
    /* Global Styling */
    
    /* Game Container */
    .game-container {
        border-radius: 15px;
        padding: 20px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    /* Sentences Styling */
    .sentence-container {
        display: flex;
        align-items: center;
        border-left: 4px solid #3498db;
        padding: 10px;
        margin-bottom: 10px;
        transition: all 0.3s ease;
    }
    .sentence-container:hover {
        transform: translateX(5px);
    }
    /* Buttons */
    .stButton>button {
        color: black;
        border-radius: 20px;
        border: none;
        padding: 10px 20px;
        transition: all 0.3s ease;
    }
    .stButton>button:hover {
        transform: scale(1.05);
    }
    /* Radio Buttons */
    .stRadio>div {
        display: flex;
        flex-wrap: wrap;
        gap: 10px;
    }
    .stRadio>div>label {
        background-color: #000;
        padding: 10px;
        border-radius: 10px;
        transition: all 0.3s ease;
    }
    .stRadio>div>label:hover {
        background-color: #3498db;
        color: white;
    }
    /* Sidebar */
    .css-1aumxhk {
        background-color: #2c3e50;
        color: white;
    }
    /* Reason Validation */
    .reason-valid {
        color: #2ecc71;
        font-weight: bold;
    }
    .reason-invalid {
        color: #e74c3c;
        font-weight: bold;
    }
    </style>
    """, unsafe_allow_html=True)

class RoFTGame:
    def __init__(self, dataset_path):
        """
        Initialize the RoFT Game with the dataset
        
        :param dataset_path: Path to the roft.csv file
        """
        self.df = pd.read_csv(dataset_path)
        self.current_sample = None
        self.current_sentences = None
        self.true_boundary_index = None
        self.current_guess_index = None
        # Predefined reasons from the dataset description
        self.predefined_reasons = [
            "grammar",
            "repetition", 
            "irrelevant", 
            "contradicts_sentence", 
            "contradicts_knowledge", 
            "common_sense", 
            "coreference", 
            "generic"
        ]

    def load_random_sample(self):
        """
        Load a random sample from the dataset
        """
        # Filter for samples with valid generations and reasons
        valid_samples = self.df[
            (self.df['gen_body'].notna()) & 
            (self.df['reason'].notna()) & 
            (self.df['reason'] != '[]')
        ]
        
        # Select a random sample
        self.current_sample = valid_samples.sample(n=1).iloc[0]
        
        # Prepare sentences
        prompt_sentences = self.current_sample['prompt_body'].split('_SEP_')
        gen_sentences = self.current_sample['gen_body'].split('_SEP_')
        
        # Combine and truncate to 10 sentences
        combined_sentences = prompt_sentences + gen_sentences
        self.current_sentences = combined_sentences[:10]
        
        # Store true boundary
        self.true_boundary_index = self.current_sample['true_boundary_index']
        
        # Parse reasons from the dataset
        try:
            self.current_reasons = ast.literal_eval(self.current_sample['reason'])
        except:
            self.current_reasons = []
        
        # Reset current guess
        self.current_guess_index = None

    def check_guess(self, guess_index):
        """
        Check if the guess is correct
        
        :param guess_index: Index of the guessed boundary
        :return: Points earned
        """
        self.current_guess_index = guess_index
        
        # Calculate points based on closeness to true boundary
        if guess_index == self.true_boundary_index:
            return 5
        elif guess_index > self.true_boundary_index:
            return max(5 - (guess_index - self.true_boundary_index), 0)
        else:
            return 0

    def validate_reason(self, user_reason):
        """
        Validate user's reason against dataset reasons
        
        :param user_reason: Reason provided by user
        :return: Tuple of (is_valid, matching_reasons)
        """
        # Convert user reason to lowercase for matching
        user_reason_lower = user_reason.lower()
        
        # Check against predefined reasons and current sample's reasons
        matching_reasons = []
        
        # Check predefined reasons
        for reason in self.predefined_reasons:
            if reason.lower() in user_reason_lower:
                matching_reasons.append(reason)
        
        # Check original sample's reasons
        for orig_reason in self.current_reasons:
            if orig_reason.lower() in user_reason_lower:
                matching_reasons.append(orig_reason)
        
        return len(matching_reasons) > 0, matching_reasons

    def save_annotation(self, guess_index, reason, reason_validity):
        """
        Save annotation to a text file
        
        :param guess_index: Index of the guessed boundary
        :param reason: Reason for the guess
        :param reason_validity: Validity of the reason
        """
        # Ensure logs directory exists
        os.makedirs('logs', exist_ok=True)
        
        # Generate unique filename with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f'logs/annotation_{timestamp}.txt'
        
        # Prepare annotation details
        annotation_details = [
            f"Timestamp: {timestamp}",
            f"Model: {self.current_sample['model']}",
            f"Dataset: {self.current_sample['dataset']}",
            f"Guess Index: {guess_index + 1}",
            f"True Boundary Index: {self.true_boundary_index + 1}",
            f"Original Dataset Reasons: {self.current_reasons}",
            f"User Reason: {reason}",
            f"Reason Validity: {reason_validity[0]}",
            f"Matching Reasons: {reason_validity[1]}",
            "\nFull Text:\n" + "\n".join(f"{i+1}. {sent}" for i, sent in enumerate(self.current_sentences))
        ]
        
        # Write to file
        with open(filename, 'w') as f:
            f.write("\n".join(annotation_details))

def main():
    local_css()
    
    # Fancy title with animation
    st.markdown("""
    <h1 style='text-align: center; color: #2c3e50; 
    text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
    animation: fadeIn 2s;'>
    🕵️ Real or Fake Text Detective 🕵️‍♀️
    </h1>
    """, unsafe_allow_html=True)

    # Game introduction
    st.markdown("""
    <div class='game-container'>
    <p style='text-align: center; font-style: italic;'>
    Sharpen your AI detection skills! Read carefully and identify where human writing transforms into machine-generated text.
    </p>
    </div>
    """, unsafe_allow_html=True)
    
    
    # Initialize game session state
    if 'game' not in st.session_state:
        st.session_state.game = RoFTGame('roft.csv')
        st.session_state.game.load_random_sample()
        st.session_state.total_points = 0
        st.session_state.rounds_played = 0
    
    # Game container
    st.markdown("<div class='game-container'>", unsafe_allow_html=True)
    
    # Game information in sidebar with icons
    st.sidebar.markdown("## 🎮 Game Stats")
    st.sidebar.markdown(f"### 🏆 Total Points: {st.session_state.total_points}")
    st.sidebar.markdown(f"### 🎲 Rounds Played: {st.session_state.rounds_played}")

        # Animated difficulty indicator
    difficulty_map = {
        'gpt2': '🟢 Easy', 
        'gpt2-xl': '🟠 Medium', 
        'ctrl': '🔴 Hard'
    }
    current_model = st.session_state.game.current_sample['model']
    difficulty = difficulty_map.get(current_model, '⚪ Unknown')
    st.sidebar.markdown(f"### 🎯 Difficulty: {difficulty}")
    
    # Display sentences with enhanced styling and radio buttons
    st.subheader("🔍 Examine the Text Carefully")
    selected_guess = st.radio(
        "Where do you think the AI-generated text begins?",
        options=[f"{i+1}. {sentence}" for i, sentence in enumerate(st.session_state.game.current_sentences)],
        label_visibility="collapsed"
    )
    
    # Reason input with predefined options and visual enhancements
    st.markdown("### 🧐 Explain Your Reasoning")
    reason_options = st.session_state.game.predefined_reasons
    selected_predefined_reasons = st.multiselect(
        "Select indicators of AI generation",
        options=reason_options
    )
    
    # Custom reason input
    custom_reason = st.text_area("Additional detective notes (optional)")
    
    # Combine predefined and custom reasons
    full_reason = " ".join(selected_predefined_reasons)
    if custom_reason:
        full_reason += f" {custom_reason}"

        
    
    # Guess submission
    if st.button("Submit Guess"):
        if not full_reason.strip():
            st.warning("Please provide a reason for your guess.")
        else:
            # Convert guess to index (subtract 1 for 0-based indexing)
            guess_index = int(selected_guess.split('.')[0]) - 1
            
            # Check guess and update points
            points_earned = st.session_state.game.check_guess(guess_index)
            st.session_state.total_points += points_earned
            st.session_state.rounds_played += 1
            
            # If guess is correct, validate reason
            if guess_index == st.session_state.game.true_boundary_index:
                reason_validity = st.session_state.game.validate_reason(full_reason)
                st.subheader("Results")
                st.write(f"Your Guess: Sentence {selected_guess}")
                st.write(f"Actual Boundary: Sentence {st.session_state.game.true_boundary_index + 1}")
                st.write(f"Points Earned: {points_earned}")

                # Display reason validation
                st.write("Reason Validation:")
                st.write(f"Valid Reason: {reason_validity[0]}")
                if reason_validity[1]:
                    st.write("Matching Reasons:", ", ".join(reason_validity[1]))

                # Save annotation
                st.session_state.game.save_annotation(guess_index, full_reason, reason_validity)
            else:
                st.subheader("Results")
                st.write(f"Your Guess: Sentence {selected_guess}")
                st.write(f"Actual Boundary: Sentence {st.session_state.game.true_boundary_index + 1}")
                st.write(f"Points Earned: {points_earned}")
            
            # Option to continue
            if st.button("Next Round"):
                st.session_state.game.load_random_sample()
                st.experimental_rerun()
                
    st.markdown("</div>", unsafe_allow_html=True)

    # Optional: Show metadata for current sample
    if st.checkbox("Show Sample Metadata"):
        st.write("Current Sample Details:")
        sample = st.session_state.game.current_sample
        st.write(f"Model: {sample['model']}")
        st.write(f"Dataset: {sample['dataset']}")
        st.write(f"Sampling Strategy (p): {sample['dec_strat_value']}")
        st.write(f"Original Reasons: {st.session_state.game.current_reasons}")

if __name__ == "__main__":
    main()