import gradio as gr
import pandas as pd
import json
import os
import re
from PyPDF2 import PdfReader
from collections import defaultdict
from transformers import pipeline

# Initialize NER model (will load only if transformers is available)
try:
    ner_pipeline = pipeline("ner", model="dslim/bert-base-NER")
except Exception as e:
    print(f"Could not load NER model: {e}")
    ner_pipeline = None

# ========== IMPROVED TRANSCRIPT PARSING ==========
def extract_gpa(text):
    gpa_data = {'weighted': "N/A", 'unweighted': "N/A"}
    gpa_patterns = [
        r'Weighted GPA[\s:]*(\d\.\d{1,2})',
        r'GPA \(Weighted\)[\s:]*(\d\.\d{1,2})',
        r'Cumulative GPA \(Weighted\)[\s:]*(\d\.\d{1,2})',
        r'Unweighted GPA[\s:]*(\d\.\d{1,2})',
        r'GPA \(Unweighted\)[\s:]*(\d\.\d{1,2})',
        r'Cumulative GPA \(Unweighted\)[\s:]*(\d\.\d{1,2})',
        r'GPA[\s:]*(\d\.\d{1,2})'
    ]
    for pattern in gpa_patterns:
        for match in re.finditer(pattern, text, re.IGNORECASE):
            gpa_value = match.group(1)
            if 'weighted' in pattern.lower():
                gpa_data['weighted'] = gpa_value
            elif 'unweighted' in pattern.lower():
                gpa_data['unweighted'] = gpa_value
            else:
                if gpa_data['unweighted'] == "N/A":
                    gpa_data['unweighted'] = gpa_value
                if gpa_data['weighted'] == "N/A":
                    gpa_data['weighted'] = gpa_value
    return gpa_data

def extract_courses_with_regex(text):
    patterns = [
        r'(?:^|\n)([A-Z]{2,}\s*-?\s*\d{3}[A-Z]?\b)\s*([A-F][+-]?|\d{2,3}%)?',
        r'(?:^|\n)([A-Z][a-z]+(?:\s+[A-Z]?[a-z]+)+)\s*[:\-]?\s*([A-F][+-]?|\d{2,3}%)?',
        r'(?:^|\n)([A-Z]{2,})\s*\d{3}\b'
    ]
    courses = []
    for pattern in patterns:
        for match in re.finditer(pattern, text, re.MULTILINE):
            course_name = match.group(1).strip()
            grade = match.group(2).strip() if match.group(2) else None
            courses.append({'name': course_name, 'grade': grade})
    return courses

def extract_grade_levels(text):
    grade_pattern = r'(?:Grade|Year|Term)\s*[:]?\s*(\d+|Freshman|Sophomore|Junior|Senior)\b'
    grade_matches = list(re.finditer(grade_pattern, text, re.IGNORECASE))
    grade_sections = []
    for i, match in enumerate(grade_matches):
        start_pos = match.start()
        end_pos = grade_matches[i+1].start() if i+1 < len(grade_matches) else len(text)
        grade_sections.append({
            'grade': match.group(1),
            'text': text[start_pos:end_pos]
        })
    return grade_sections

def parse_transcript(file):
    if file.name.endswith('.pdf'):
        text = ''
        reader = PdfReader(file)
        for page in reader.pages:
            text += page.extract_text() + '\n'
        
        # Try both NER and regex approaches
        courses = []
        if ner_pipeline:
            try:
                entities = ner_pipeline(text)
                current_course = {}
                for entity in entities:
                    if entity['word'].startswith('##'):
                        current_course['name'] = current_course.get('name', '') + entity['word'][2:]
                    elif entity['entity'] in ['B-ORG', 'I-ORG']:  # Using ORG as proxy for courses
                        if 'name' in current_course:
                            courses.append(current_course)
                        current_course = {'name': entity['word']}
                    elif entity['entity'] == 'GRADE' and current_course:
                        current_course['grade'] = entity['word']
                if current_course:
                    courses.append(current_course)
            except Exception as e:
                print(f"NER failed: {e}")
        
        # Fallback to regex if NER didn't find courses
        if not courses:
            courses = extract_courses_with_regex(text)
        
        # Organize by grade level
        grade_sections = extract_grade_levels(text)
        courses_by_grade = defaultdict(list)
        
        if grade_sections:
            for section in grade_sections:
                section_courses = extract_courses_with_regex(section['text'])
                for course in section_courses:
                    course['term'] = section['grade']
                    courses_by_grade[section['grade']].append(course)
        else:
            courses_by_grade["All"] = courses
        
        gpa_data = extract_gpa(text)
        
        output_text = "Transcript parsed successfully\n"
        output_text += f"Found {len(courses)} courses across {len(courses_by_grade)} grade levels\n"
        
        return output_text, {
            "gpa": gpa_data,
            "courses": dict(courses_by_grade)
        }
    elif file.name.endswith('.csv'):
        df = pd.read_csv(file)
    elif file.name.endswith('.xlsx'):
        df = pd.read_excel(file)
    else:
        return "Unsupported file format", None

    # Fallback for CSV/Excel
    gpa = "N/A"
    for col in ['GPA', 'Grade Point Average', 'Cumulative GPA']:
        if col in df.columns:
            gpa = df[col].iloc[0] if isinstance(df[col].iloc[0], (float, int)) else "N/A"
            break

    grade_level = "N/A"
    for col in ['Grade Level', 'Grade', 'Class', 'Year']:
        if col in df.columns:
            grade_level = df[col].iloc[0]
            break

    courses = []
    for col in ['Course', 'Subject', 'Course Name', 'Class']:
        if col in df.columns:
            courses = df[col].tolist()
            break

    return f"Grade Level: {grade_level}\nGPA: {gpa}", {
        "gpa": {"unweighted": gpa, "weighted": "N/A"},
        "grade_level": grade_level,
        "courses": courses
    }

# ========== LEARNING STYLE QUIZ ==========
learning_style_questions = [
    "When you study for a test, you prefer to:",
    "When you need directions to a new place, you prefer:",
    "When you learn a new skill, you prefer to:",
    "When you're trying to concentrate, you:",
    "When you meet new people, you remember them by:",
    "When you're assembling furniture or a gadget, you:",
    "When choosing a restaurant, you rely most on:",
    "When you're in a waiting room, you typically:",
    "When giving someone instructions, you tend to:",
    "When you're trying to recall information, you:",
    "When you're at a museum or exhibit, you:",
    "When you're learning a new language, you prefer:",
    "When you're taking notes in class, you:",
    "When you're explaining something complex, you:",
    "When you're at a party, you enjoy:",
    "When you're trying to remember a phone number, you:",
    "When you're relaxing, you prefer to:",
    "When you're learning to use new software, you:",
    "When you're giving a presentation, you rely on:",
    "When you're solving a difficult problem, you:"
]

learning_style_options = [
    ["Read the textbook (Reading/Writing)", "Listen to lectures (Auditory)", "Use diagrams/charts (Visual)", "Practice problems (Kinesthetic)"],
    ["Look at a map (Visual)", "Have someone tell you (Auditory)", "Write down directions (Reading/Writing)", "Try walking/driving there (Kinesthetic)"],
    ["Read instructions (Reading/Writing)", "Have someone show you (Visual)", "Listen to explanations (Auditory)", "Try it yourself (Kinesthetic)"],
    ["Need quiet (Reading/Writing)", "Need background noise (Auditory)", "Need to move around (Kinesthetic)", "Need visual stimulation (Visual)"],
    ["Their face (Visual)", "Their name (Auditory)", "What you talked about (Reading/Writing)", "What you did together (Kinesthetic)"],
    ["Read the instructions carefully (Reading/Writing)", "Look at the diagrams (Visual)", "Ask someone to explain (Auditory)", "Start putting pieces together (Kinesthetic)"],
    ["Online photos of the food (Visual)", "Recommendations from friends (Auditory)", "Reading the menu online (Reading/Writing)", "Remembering how it felt to eat there (Kinesthetic)"],
    ["Read magazines (Reading/Writing)", "Listen to music (Auditory)", "Watch TV (Visual)", "Fidget or move around (Kinesthetic)"],
    ["Write them down (Reading/Writing)", "Explain verbally (Auditory)", "Demonstrate (Visual)", "Guide them physically (Kinesthetic)"],
    ["See written words in your mind (Visual)", "Hear the information in your head (Auditory)", "Write it down to remember (Reading/Writing)", "Associate it with physical actions (Kinesthetic)"],
    ["Read all the descriptions (Reading/Writing)", "Listen to audio guides (Auditory)", "Look at the displays (Visual)", "Touch interactive exhibits (Kinesthetic)"],
    ["Study grammar rules (Reading/Writing)", "Listen to native speakers (Auditory)", "Use flashcards with images (Visual)", "Practice conversations (Kinesthetic)"],
    ["Write detailed paragraphs (Reading/Writing)", "Record the lecture (Auditory)", "Draw diagrams and charts (Visual)", "Doodle while listening (Kinesthetic)"],
    ["Write detailed steps (Reading/Writing)", "Explain verbally with examples (Auditory)", "Draw diagrams (Visual)", "Use physical objects to demonstrate (Kinesthetic)"],
    ["Conversations with people (Auditory)", "Watching others or the environment (Visual)", "Writing notes or texting (Reading/Writing)", "Dancing or physical activities (Kinesthetic)"],
    ["See the numbers in your head (Visual)", "Say them aloud (Auditory)", "Write them down (Reading/Writing)", "Dial them on a keypad (Kinesthetic)"],
    ["Read a book (Reading/Writing)", "Listen to music (Auditory)", "Watch TV/movies (Visual)", "Do something physical (Kinesthetic)"],
    ["Read the manual (Reading/Writing)", "Ask someone to show you (Visual)", "Call tech support (Auditory)", "Experiment with the software (Kinesthetic)"],
    ["Detailed notes (Reading/Writing)", "Verbal explanations (Auditory)", "Visual slides (Visual)", "Physical demonstrations (Kinesthetic)"],
    ["Write out possible solutions (Reading/Writing)", "Talk through it with someone (Auditory)", "Draw diagrams (Visual)", "Build a model or prototype (Kinesthetic)"]
]

def learning_style_quiz(*answers):
    scores = {
        "Visual": 0,
        "Auditory": 0,
        "Reading/Writing": 0,
        "Kinesthetic": 0
    }
    
    for i, answer in enumerate(answers):
        if answer == learning_style_options[i][0]:
            scores["Reading/Writing"] += 1
        elif answer == learning_style_options[i][1]:
            scores["Auditory"] += 1
        elif answer == learning_style_options[i][2]:
            scores["Visual"] += 1
        elif answer == learning_style_options[i][3]:
            scores["Kinesthetic"] += 1
    
    max_score = max(scores.values())
    total_questions = len(learning_style_questions)
    
    # Calculate percentages
    percentages = {style: (score/total_questions)*100 for style, score in scores.items()}
    
    # Sort styles by score (descending)
    sorted_styles = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    
    # Prepare detailed results
    result = "Your Learning Style Results:\n\n"
    for style, score in sorted_styles:
        result += f"{style}: {score}/{total_questions} ({percentages[style]:.1f}%)\n"
    
    result += "\n"
    
    # Determine primary and secondary styles
    primary_styles = [style for style, score in scores.items() if score == max_score]
    
    if len(primary_styles) == 1:
        result += f"Your primary learning style is: {primary_styles[0]}\n\n"
        # Add personalized tips based on primary style
        if primary_styles[0] == "Visual":
            result += "Tips for Visual Learners:\n"
            result += "- Use color coding in your notes\n"
            result += "- Create mind maps and diagrams\n"
            result += "- Watch educational videos\n"
            result += "- Use flashcards with images\n"
        elif primary_styles[0] == "Auditory":
            result += "Tips for Auditory Learners:\n"
            result += "- Record lectures and listen to them\n"
            result += "- Participate in study groups\n"
            result += "- Explain concepts out loud to yourself\n"
            result += "- Use rhymes or songs to remember information\n"
        elif primary_styles[0] == "Reading/Writing":
            result += "Tips for Reading/Writing Learners:\n"
            result += "- Write detailed notes\n"
            result += "- Create summaries in your own words\n"
            result += "- Read textbooks and articles\n"
            result += "- Make lists to organize information\n"
        else:  # Kinesthetic
            result += "Tips for Kinesthetic Learners:\n"
            result += "- Use hands-on activities\n"
            result += "- Take frequent movement breaks\n"
            result += "- Create physical models\n"
            result += "- Associate information with physical actions\n"
    else:
        result += f"You have multiple strong learning styles: {', '.join(primary_styles)}\n\n"
        result += "You may benefit from combining different learning approaches.\n"
    
    return result

# ========== SAVE STUDENT PROFILE ==========
def save_profile(name, age, interests, transcript, learning_style, 
                movie, movie_reason, show, show_reason, 
                book, book_reason, character, character_reason, blog):
    # Convert age to int if it's a numpy number (from gradio Number input)
    age = int(age) if age else 0
    
    favorites = {
        "movie": movie,
        "movie_reason": movie_reason,
        "show": show,
        "show_reason": show_reason,
        "book": book,
        "book_reason": book_reason,
        "character": character,
        "character_reason": character_reason
    }
    
    data = {
        "name": name,
        "age": age,
        "interests": interests,
        "transcript": transcript,
        "learning_style": learning_style,
        "favorites": favorites,
        "blog": blog
    }
    
    os.makedirs("student_profiles", exist_ok=True)
    json_path = os.path.join("student_profiles", f"{name.replace(' ', '_')}_profile.json")
    with open(json_path, "w") as f:
        json.dump(data, f, indent=2)

    markdown_summary = f"""### Student Profile: {name}
**Age:** {age}  
**Interests:** {interests}  
**Learning Style:** {learning_style}  
#### Transcript:
{transcript_display(transcript)}
#### Favorites:
- Movie: {favorites['movie']} ({favorites['movie_reason']})
- Show: {favorites['show']} ({favorites['show_reason']})
- Book: {favorites['book']} ({favorites['book_reason']})
- Character: {favorites['character']} ({favorites['character_reason']})
#### Blog:
{blog if blog else "_No blog provided_"}
"""
    return markdown_summary

def transcript_display(transcript_dict):
    if not transcript_dict or "courses" not in transcript_dict:
        return "No course information available"
    
    display = "### Course History\n\n"
    courses_by_grade = transcript_dict["courses"]
    
    if isinstance(courses_by_grade, dict):
        for grade, courses in courses_by_grade.items():
            display += f"**{grade}**\n"
            for course in courses:
                if isinstance(course, dict):
                    display += f"- {course.get('name', 'N/A')}"
                    if 'grade' in course:
                        display += f" (Grade: {course['grade']})"
                    if 'term' in course:
                        display += f" | Term: {course['term']}"
                    display += "\n"
                else:
                    display += f"- {str(course)}\n"
            display += "\n"
    elif isinstance(courses_by_grade, list):
        for course in courses_by_grade:
            if isinstance(course, dict):
                display += f"- {course.get('name', 'N/A')}"
                if 'grade' in course:
                    display += f" (Grade: {course['grade']})"
                display += "\n"
            else:
                display += f"- {str(course)}\n"
    
    if 'gpa' in transcript_dict:
        gpa = transcript_dict['gpa']
        display += "\n**GPA Information**\n"
        display += f"- Unweighted: {gpa.get('unweighted', 'N/A')}\n"
        display += f"- Weighted: {gpa.get('weighted', 'N/A')}\n"
    
    return display

# ========== AI TEACHING ASSISTANT ==========
def load_profile():
    if not os.path.exists("student_profiles"):
        return {}
    files = [f for f in os.listdir("student_profiles") if f.endswith('.json')]
    if files:
        with open(os.path.join("student_profiles", files[0]), "r") as f:
            return json.load(f)
    return {}

def generate_response(message, history):
    profile = load_profile()
    if not profile:
        return "Please complete and save your profile first using the previous tabs."
    
    # Get profile data
    learning_style = profile.get("learning_style", "")
    grade_level = profile.get("transcript", {}).get("grade_level", "unknown")
    gpa = profile.get("transcript", {}).get("gpa", {})
    interests = profile.get("interests", "")
    
    # Common responses
    greetings = ["hi", "hello", "hey"]
    study_help = ["study", "learn", "prepare", "exam"]
    grade_help = ["grade", "gpa", "score"]
    interest_help = ["interest", "hobby", "passion"]
    
    if any(greet in message.lower() for greet in greetings):
        return f"Hello {profile.get('name', 'there')}! How can I help you today?"
    
    elif any(word in message.lower() for word in study_help):
        if "Visual" in learning_style:
            response = ("Based on your visual learning style, I recommend:\n"
                       "- Creating mind maps or diagrams\n"
                       "- Using color-coded notes\n"
                       "- Watching educational videos")
        elif "Auditory" in learning_style:
            response = ("Based on your auditory learning style, I recommend:\n"
                       "- Recording lectures and listening to them\n"
                       "- Participating in study groups\n"
                       "- Explaining concepts out loud")
        elif "Reading/Writing" in learning_style:
            response = ("Based on your reading/writing learning style, I recommend:\n"
                       "- Writing detailed notes\n"
                       "- Creating summaries in your own words\n"
                       "- Reading textbooks and articles")
        elif "Kinesthetic" in learning_style:
            response = ("Based on your kinesthetic learning style, I recommend:\n"
                       "- Hands-on practice\n"
                       "- Creating physical models\n"
                       "- Taking frequent movement breaks")
        else:
            response = ("Here are some general study tips:\n"
                       "- Break study sessions into 25-minute chunks\n"
                       "- Review material regularly\n"
                       "- Teach concepts to someone else")
        
        return response
    
    elif any(word in message.lower() for word in grade_help):
        return (f"Your GPA information:\n"
               f"- Unweighted: {gpa.get('unweighted', 'N/A')}\n"
               f"- Weighted: {gpa.get('weighted', 'N/A')}\n\n"
               "To improve your grades, try:\n"
               "- Setting specific goals\n"
               "- Meeting with teachers\n"
               "- Developing a study schedule")
    
    elif any(word in message.lower() for word in interest_help):
        return (f"I see you're interested in: {interests}\n\n"
               "You might want to:\n"
               "- Find clubs or activities related to these interests\n"
               "- Explore career paths that align with them")
    
    elif "help" in message.lower():
        return ("I can help with:\n"
               "- Study tips based on your learning style\n"
               "- GPA and grade information\n"
               "- General academic advice\n\n"
               "Try asking about study strategies or your grades!")
    
    else:
        return ("I'm your personalized teaching assistant. "
               "I can help with study tips, grade information, and academic advice. "
               "Try asking about how to study for your classes!")

# ========== GRADIO INTERFACE ==========
with gr.Blocks() as app:
    with gr.Tab("Step 1: Upload Transcript"):
        gr.Markdown("### Upload your transcript (PDF recommended for best results)")
        transcript_file = gr.File(label="Transcript file", file_types=[".pdf", ".csv", ".xlsx"])
        transcript_output = gr.Textbox(label="Parsing Results")
        transcript_data = gr.State()
        transcript_file.change(
            fn=parse_transcript,
            inputs=transcript_file,
            outputs=[transcript_output, transcript_data]
        )

    with gr.Tab("Step 2: Learning Style Quiz"):
        gr.Markdown("### Learning Style Quiz (20 Questions)")
        quiz_components = []
        for i, (question, options) in enumerate(zip(learning_style_questions, learning_style_options)):
            quiz_components.append(gr.Radio(options, label=f"{i+1}. {question}"))
        
        learning_output = gr.Textbox(label="Your Learning Style", lines=10)
        gr.Button("Submit Quiz").click(
            fn=learning_style_quiz,
            inputs=quiz_components,
            outputs=learning_output
        )

    with gr.Tab("Step 3: Personal Questions"):
        name = gr.Textbox(label="What's your name?")
        age = gr.Number(label="How old are you?", precision=0)
        interests = gr.Textbox(label="What are your interests?")
        movie = gr.Textbox(label="Favorite movie?")
        movie_reason = gr.Textbox(label="Why do you like that movie?")
        show = gr.Textbox(label="Favorite TV show?")
        show_reason = gr.Textbox(label="Why do you like that show?")
        book = gr.Textbox(label="Favorite book?")
        book_reason = gr.Textbox(label="Why do you like that book?")
        character = gr.Textbox(label="Favorite character?")
        character_reason = gr.Textbox(label="Why do you like that character?")
        blog_checkbox = gr.Checkbox(label="Do you want to write a blog?", value=False)
        blog_text = gr.Textbox(label="Write your blog here", visible=False, lines=5)
        blog_checkbox.change(lambda x: gr.update(visible=x), inputs=blog_checkbox, outputs=blog_text)

    with gr.Tab("Step 4: Save & Review"):
        output_summary = gr.Markdown()
        save_btn = gr.Button("Save Profile")
        save_btn.click(
            fn=save_profile,
            inputs=[name, age, interests, transcript_data, learning_output,
                   movie, movie_reason, show, show_reason,
                   book, book_reason, character, character_reason, blog_text],
            outputs=output_summary
        )

    with gr.Tab("🤖 AI Teaching Assistant"):
        gr.Markdown("## Your Personalized Learning Assistant")
        chatbot = gr.ChatInterface(
            fn=generate_response,
            examples=[
                "How should I study for my next test?",
                "What's my GPA information?",
                "Help me with study strategies",
                "How can I improve my grades?"
            ]
        )

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