import gradio as gr
import torch
import numpy as np
import random
import pandas as pd
import matplotlib.pyplot as plt
import time
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


# Sample problems database (you would expand this)
sample_problems = {
    "algebra": [
        {
            "problem": "Find all positive integers n such that n^2 + 20 is divisible by n + 5.",
            "difficulty": "medium",
            "solution": "Let's denote n^2 + 20 = k(n + 5) for some integer k.\nThis gives us n^2 + 20 = kn + 5k\nn^2 - kn - 5k + 20 = 0\nn^2 - kn = 5k - 20\nWe need to find values of n such that n^2 - kn = 5k - 20 has solutions.\nRearranging, we get n(n - k) = 5k - 20\nFor n > 0 and n + 5 to divide n^2 + 20, we need to check possible values.\nTrying n = 4: 4^2 + 20 = 36, 4 + 5 = 9, and 36 is divisible by 9. So n = 4 works.\nTrying n = 15: 15^2 + 20 = 245, 15 + 5 = 20, and 245 is not divisible by 20.\nAfter checking more values systematically, we find that n = 4 is the only positive integer solution."
        },
        {
            "problem": "Determine all real values of x such that log_(x-1)(x^2 - 5x + 7) = 2.",
            "difficulty": "hard",
            "solution": "For log_(x-1)(x^2 - 5x + 7) = 2 to be defined, we need:\n1) x - 1 > 0, so x > 1\n2) x - 1 ≠ 1, so x ≠ 2\n3) x^2 - 5x + 7 > 0\n\nNow, log_(x-1)(x^2 - 5x + 7) = 2 means (x^2 - 5x + 7) = (x-1)^2\n\nExpanding (x-1)^2 = (x-1)(x-1) = x^2 - 2x + 1\n\nSo we need to solve x^2 - 5x + 7 = x^2 - 2x + 1\n-5x + 7 = -2x + 1\n-3x = -6\nx = 2\n\nBut we already established x ≠ 2, so there are no solutions."
        }
    ],
    "geometry": [
        {
            "problem": "Points A, B, C, and D lie on a circle in that order. If AB = BC = CD and angle BAC = 30°, what is the measure of angle ADC in degrees?",
            "difficulty": "medium",
            "solution": "Since AB = BC = CD, we know that B divides arc AC into two equal parts, and C divides arc BD into two equal parts.\n\nLet's denote the center of the circle as O.\nSince AB = BC, triangles AOB and BOC are isosceles.\nThis means angle AOB = angle BOA and angle BOC = angle COB.\n\nWe know angle BAC = 30°.\nBy the inscribed angle theorem, angle BAC = (1/2) × (arc BC).\nSo arc BC = 60°.\n\nSince AB = BC = CD, arcs AB, BC, and CD all have the same length.\nThis means arc AB = arc BC = arc CD = 60°.\n\nBy the inscribed angle theorem, angle ADC = (1/2) × (arc AC).\nArc AC = arc AB + arc BC = 60° + 60° = 120°.\nTherefore, angle ADC = (1/2) × 120° = 60°."
        }
    ],
    "number_theory": [
        {
            "problem": "Find the sum of all positive integers n such that n^2 + n + 1 is divisible by 7.",
            "difficulty": "hard",
            "solution": "Let's consider n mod 7 and check when n^2 + n + 1 ≡ 0 (mod 7).\n\nFor n ≡ 0 (mod 7): 0^2 + 0 + 1 = 1 ≡ 1 (mod 7) ❌\nFor n ≡ 1 (mod 7): 1^2 + 1 + 1 = 3 ≡ 3 (mod 7) ❌\nFor n ≡ 2 (mod 7): 2^2 + 2 + 1 = 7 ≡ 0 (mod 7) ✓\nFor n ≡ 3 (mod 7): 3^2 + 3 + 1 = 13 ≡ 6 (mod 7) ❌\nFor n ≡ 4 (mod 7): 4^2 + 4 + 1 = 21 ≡ 0 (mod 7) ✓\nFor n ≡ 5 (mod 7): 5^2 + 5 + 1 = 31 ≡ 3 (mod 7) ❌\nFor n ≡ 6 (mod 7): 6^2 + 6 + 1 = 43 ≡ 1 (mod 7) ❌\n\nSo n^2 + n + 1 is divisible by 7 when n ≡ 2 (mod 7) or n ≡ 4 (mod 7).\n\nFor n ≤ 100, the positive integers that satisfy this are:\n2, 4, 9, 11, 16, 18, 23, 25, 30, 32, 37, 39, 44, 46, 51, 53, 58, 60, 65, 67, 72, 74, 79, 81, 86, 88, 93, 95, 100\n\nThe sum of these numbers is 1501."
        }
    ],
    "combinatorics": [
        {
            "problem": "How many different 4-digit numbers can be formed using the digits 1, 2, 3, 4, 5 without repetition?",
            "difficulty": "easy",
            "solution": "We need to create 4-digit numbers using 5 distinct digits without repetition.\n\nFor the first position, we have 5 choices (1, 2, 3, 4, or 5).\nFor the second position, we have 4 remaining choices.\nFor the third position, we have 3 remaining choices.\nFor the fourth position, we have 2 remaining choices.\n\nBy the multiplication principle, the total number of possible 4-digit numbers is:\n5 × 4 × 3 × 2 = 120"
        }
    ]
}

# Function to generate problem based on filters
def generate_problem(topic, difficulty):
    filtered_problems = [p for p in sample_problems.get(topic, []) if p["difficulty"] == difficulty]
    if filtered_problems:
        return random.choice(filtered_problems)["problem"]
    return "No problem found matching the criteria. Try a different combination."

# Function to solve problem using AI model
def solve_problem(problem_text):
    if not problem_text.strip():
        return "Please enter a problem first."
    
    prompt = f"Solve this math olympiad problem step by step:\n\n{problem_text}\n\nSolution:"
    
    # Add a small delay to simulate AI thinking (remove in production)
    time.sleep(2)
    
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # In a real system, you would use your AI model here
    # outputs = model.generate(inputs["input_ids"], max_length=1024, temperature=0.7)
    # solution = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Solution:")[1].strip()
    
    # For demo purposes, we'll provide a placeholder solution
    for topic in sample_problems:
        for problem in sample_problems[topic]:
            if problem["problem"] == problem_text:
                return problem["solution"]
    
    return "I'll solve this step-by-step:\n\n1. First, let's understand what the problem is asking...\n\n(This is a placeholder. In the actual implementation, the AI model would generate a detailed solution.)"

# Function to analyze student solution
def analyze_solution(problem, student_solution, ai_solution):
    if not student_solution.strip():
        return "Please enter your solution first."
    
    # In a real system, you would compare the solutions more intelligently
    # For demo purposes, we'll provide a placeholder analysis
    
    feedback = "Solution Analysis:\n\n"
    
    # Simple keyword checking (very basic, would be much more sophisticated in reality)
    ai_keywords = set([word.lower() for word in ai_solution.split() if len(word) > 4])
    student_keywords = set([word.lower() for word in student_solution.split() if len(word) > 4])
    
    common_keywords = ai_keywords.intersection(student_keywords)
    
    if len(common_keywords) / max(1, len(ai_keywords)) > 0.4:
        feedback += "✓ Your approach seems correct and contains many of the key concepts needed.\n\n"
    else:
        feedback += "⚠ Your approach may be missing some key concepts or taking a different direction.\n\n"
    
    # Check for solution steps
    if student_solution.count("\n") < 3:
        feedback += "⚠ Your solution could benefit from showing more steps and reasoning.\n\n"
    else:
        feedback += "✓ Good job showing your work step by step!\n\n"
    
    # Give general encouragement
    feedback += "Areas to focus on:\n"
    feedback += "- Consider whether you've addressed all constraints in the problem\n"
    feedback += "- Check if your solution is logically complete\n"
    feedback += "- Verify any algebraic manipulations\n\n"
    
    return feedback

# Function to generate practice schedule
def generate_schedule(topics, difficulty_level, hours_per_week, weeks):
    if not topics or not difficulty_level or not hours_per_week or not weeks:
        return "Please fill in all fields."
    
    # Create a DataFrame for the schedule
    schedule = []
    days = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
    
    # Distribute topics across the schedule
    topics_cycle = topics.copy()
    random.shuffle(topics_cycle)
    
    # Calculate hours per day (simple distribution)
    hours_per_day = [hours_per_week // 5 if d in ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday'] else 
                     hours_per_week // 10 for d in days]
    
    # Ensure the total equals hours per week
    while sum(hours_per_day) < hours_per_week:
        idx = random.randint(0, len(days)-1)
        hours_per_day[idx] += 1
    
    for week in range(1, weeks+1):
        for day_idx, day in enumerate(days):
            if hours_per_day[day_idx] > 0:
                topic = topics_cycle[week % len(topics_cycle)]
                schedule.append({
                    'Week': week,
                    'Day': day,
                    'Topic': topic,
                    'Hours': hours_per_day[day_idx],
                    'Difficulty': difficulty_level
                })
    
    df = pd.DataFrame(schedule)
    
    # Create a visualization
    fig, ax = plt.subplots(figsize=(10, 6))
    topic_hours = df.groupby('Topic')['Hours'].sum().reset_index()
    ax.bar(topic_hours['Topic'], topic_hours['Hours'])
    ax.set_title('Hours by Topic in Training Schedule')
    ax.set_xlabel('Topic')
    ax.set_ylabel('Total Hours')
    plt.xticks(rotation=45)
    plt.tight_layout()
    
    # Convert to HTML for display
    schedule_html = df.to_html(index=False)
    
    return fig, schedule_html

# Function to track progress
def update_progress(topic, correct, incorrect):
    # This would connect to a database in a real implementation
    # For now, we just return a visualization
    topics = ['Algebra', 'Geometry', 'Number Theory', 'Combinatorics', 'Calculus']
    correct_counts = [0, 0, 0, 0, 0]
    incorrect_counts = [0, 0, 0, 0, 0]
    
    # Update the counts based on input
    try:
        topic_idx = topics.index(topic)
        correct_counts[topic_idx] = int(correct)
        incorrect_counts[topic_idx] = int(incorrect)
    except:
        pass
    
    # Create the progress chart
    fig, ax = plt.subplots(figsize=(10, 6))
    
    x = np.arange(len(topics))
    width = 0.35
    
    ax.bar(x - width/2, correct_counts, width, label='Correct')
    ax.bar(x + width/2, incorrect_counts, width, label='Incorrect')
    
    # Add labels and legend
    ax.set_ylabel('Number of Problems')
    ax.set_title('Progress by Topic')
    ax.set_xticks(x)
    ax.set_xticklabels(topics)
    ax.legend()
    
    plt.tight_layout()
    
    # Calculate accuracy
    total = sum(correct_counts) + sum(incorrect_counts)
    accuracy = sum(correct_counts) / max(1, total) * 100
    
    return fig, f"Overall Accuracy: {accuracy:.1f}%"

# Function for the competition simulator
def simulate_competition(num_problems, difficulty, time_limit):
    if not num_problems or not difficulty or not time_limit:
        return "Please fill in all fields."
    
    # Generate a set of problems for the competition
    competition_problems = []
    for topic in sample_problems:
        filtered = [p for p in sample_problems[topic] if p["difficulty"] == difficulty]
        if filtered:
            competition_problems.extend(filtered[:min(2, len(filtered))])
    
    if len(competition_problems) > num_problems:
        competition_problems = random.sample(competition_problems, num_problems)
    
    # Format the problems
    formatted_problems = ""
    for i, p in enumerate(competition_problems, 1):
        formatted_problems += f"Problem {i}: {p['problem']}\n\n"
    
    # Calculate expected time per problem
    time_per_problem = time_limit / max(1, len(competition_problems))
    
    return f"Competition Simulation\n\nDifficulty: {difficulty}\nTime Limit: {time_limit} minutes\nRecommended time per problem: {time_per_problem:.1f} minutes\n\n{formatted_problems}"

# Create the Gradio interface
with gr.Blocks(title="AI Math Olympiad Trainer") as demo:
    gr.Markdown("# AI Math Olympiad Trainer System")
    
    with gr.Tab("Problem Generator"):
        gr.Markdown("### Generate and Solve Math Olympiad Problems")
        
        with gr.Row():
            with gr.Column():
                topic_dropdown = gr.Dropdown(
                    choices=["algebra", "geometry", "number_theory", "combinatorics"], 
                    label="Topic"
                )
                difficulty_dropdown = gr.Dropdown(
                    choices=["easy", "medium", "hard"], 
                    label="Difficulty"
                )
                generate_btn = gr.Button("Generate Problem")
            
            with gr.Column():
                problem_output = gr.Textbox(label="Problem", lines=5)
        
        with gr.Row():
            with gr.Column():
                solution_input = gr.Textbox(label="Your Solution", lines=10)
                analyze_btn = gr.Button("Analyze My Solution")
            
            with gr.Column():
                ai_solution_btn = gr.Button("Get AI Solution")
                ai_solution_output = gr.Textbox(label="AI Solution", lines=10)
                analysis_output = gr.Textbox(label="Analysis", lines=8)
    
    with gr.Tab("Training Schedule"):
        gr.Markdown("### Create a Personalized Training Schedule")
        
        with gr.Row():
            with gr.Column():
                topics_multiselect = gr.CheckboxGroup(
                    choices=["Algebra", "Geometry", "Number Theory", "Combinatorics", "Calculus"],
                    label="Select Topics"
                )
                difficulty_radio = gr.Radio(
                    choices=["easy", "medium", "hard", "mixed"],
                    label="Difficulty Level"
                )
                hours_slider = gr.Slider(
                    minimum=1, maximum=30, value=10, step=1,
                    label="Hours per Week"
                )
                weeks_slider = gr.Slider(
                    minimum=1, maximum=12, value=4, step=1,
                    label="Number of Weeks"
                )
                schedule_btn = gr.Button("Generate Schedule")
            
            with gr.Column():
                schedule_plot = gr.Plot(label="Hours Distribution")
                schedule_output = gr.HTML(label="Your Schedule")
    
    with gr.Tab("Progress Tracker"):
        gr.Markdown("### Track Your Progress")
        
        with gr.Row():
            with gr.Column():
                progress_topic = gr.Dropdown(
                    choices=["Algebra", "Geometry", "Number Theory", "Combinatorics", "Calculus"],
                    label="Topic"
                )
                correct_slider = gr.Slider(
                    minimum=0, maximum=50, value=0, step=1,
                    label="Correct Solutions"
                )
                incorrect_slider = gr.Slider(
                    minimum=0, maximum=50, value=0, step=1,
                    label="Incorrect Solutions"
                )
                update_btn = gr.Button("Update Progress")
            
            with gr.Column():
                progress_plot = gr.Plot(label="Progress Chart")
                accuracy_output = gr.Textbox(label="Accuracy")
    
    with gr.Tab("Competition Simulator"):
        gr.Markdown("### Simulate a Math Competition")
        
        with gr.Row():
            with gr.Column():
                problems_slider = gr.Slider(
                    minimum=1, maximum=10, value=3, step=1,
                    label="Number of Problems"
                )
                comp_difficulty = gr.Radio(
                    choices=["easy", "medium", "hard"],
                    label="Difficulty"
                )
                time_slider = gr.Slider(
                    minimum=15, maximum=180, value=60, step=15,
                    label="Time Limit (minutes)"
                )
                simulate_btn = gr.Button("Start Simulation")
            
            with gr.Column():
                simulation_output = gr.Textbox(label="Competition Problems", lines=15)
    
    # Connect the functions
    generate_btn.click(
        generate_problem, 
        inputs=[topic_dropdown, difficulty_dropdown], 
        outputs=problem_output
    )
    
    ai_solution_btn.click(
        solve_problem,
        inputs=[problem_output],
        outputs=ai_solution_output
    )
    
    analyze_btn.click(
        analyze_solution,
        inputs=[problem_output, solution_input, ai_solution_output],
        outputs=analysis_output
    )
    
    schedule_btn.click(
        generate_schedule,
        inputs=[topics_multiselect, difficulty_radio, hours_slider, weeks_slider],
        outputs=[schedule_plot, schedule_output]
    )
    
    update_btn.click(
        update_progress,
        inputs=[progress_topic, correct_slider, incorrect_slider],
        outputs=[progress_plot, accuracy_output]
    )
    
    simulate_btn.click(
        simulate_competition,
        inputs=[problems_slider, comp_difficulty, time_slider],
        outputs=simulation_output
    )

# Launch the app
demo.launch()