app.py

import streamlit as st
import pandas as pd
import json
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, timedelta
import google.generativeai as genai
import os
import re
from pathlib import Path
import math
import time

# ============================================
# 🚀 CONFIGURATION FOR HUGGING FACE DEPLOYMENT
# ============================================

# Page configuration
st.set_page_config(
    page_title="AI Study Planner Pro",
    page_icon="📚",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Get API Key from Hugging Face Secrets
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

# ============================================
# 🎨 CUSTOM CSS WITH IMPROVED STYLING
# ============================================

st.markdown("""
<style>
    /* Main Headers */
    .main-header {
        font-size: 2.5rem;
        color: #1E3A8A;
        text-align: center;
        margin-bottom: 1rem;
        font-weight: 800;
    }
    
    .tagline {
        font-size: 1.2rem;
        color: #6B7280;
        text-align: center;
        margin-bottom: 2rem;
        font-weight: 400;
    }
    
    .sub-header {
        font-size: 1.8rem;
        color: #3B82F6;
        margin-top: 1.5rem;
        margin-bottom: 1rem;
        font-weight: 700;
        border-bottom: 3px solid #3B82F6;
        padding-bottom: 10px;
    }
    
    /* Cards and Containers */
    .metric-card {
        background-color: #F3F4F6;
        padding: 1rem;
        border-radius: 10px;
        border-left: 5px solid #3B82F6;
        margin-bottom: 1rem;
    }
    
    .plan-card {
        background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
        padding: 1.5rem;
        border-radius: 12px;
        margin: 1rem 0;
        border: 2px solid #3B82F6;
    }
    
    .test-container {
        background-color: #F8FAFC;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
        border: 2px solid #3B82F6;
    }
    
    .result-box {
        background-color: #F0F9FF;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
        border: 2px solid #3B82F6;
    }
    
    .week-focus-box {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
    }
    
    .week-summary-box {
        background-color: #F0F9FF;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
        border-left: 5px solid #3B82F6;
    }
    
    .day-focus-box {
        background-color: #EFF6FF;
        padding: 1rem;
        border-radius: 8px;
        margin: 0.5rem 0;
        border-left: 4px solid #10B981;
    }
    
    .recommendation-box {
        background-color: #EFF6FF;
        padding: 1rem;
        border-radius: 10px;
        margin: 0.5rem 0;
        border-left: 4px solid #3B82F6;
    }
    
    /* Study Day Elements */
    .study-day {
        background-color: #EFF6FF;
        padding: 1rem;
        margin: 0.5rem 0;
        border-radius: 8px;
        border: 1px solid #BFDBFE;
        transition: all 0.3s;
    }
    
    .study-day:hover {
        background-color: #DBEAFE;
        transform: translateX(5px);
    }
    
    .completed {
        background-color: #D1FAE5;
        border-color: #34D399;
    }
    
    /* Progress Elements */
    .progress-bar {
        height: 15px;
        background-color: #E5E7EB;
        border-radius: 10px;
        margin: 15px 0;
        overflow: hidden;
    }
    
    .progress-fill {
        height: 100%;
        background: linear-gradient(90deg, #10B981 0%, #3B82F6 100%);
        border-radius: 10px;
        transition: width 0.5s ease-in-out;
    }
    
    /* Status Indicators */
    .status-on-track {
        color: #10B981;
        font-weight: bold;
    }
    
    .status-behind {
        color: #EF4444;
        font-weight: bold;
    }
    
    .status-ahead {
        color: #F59E0B;
        font-weight: bold;
    }
    
    /* Task Elements */
    .task-checkbox-container {
        display: flex;
        align-items: flex-start;
        margin: 8px 0;
        padding: 8px;
        border-radius: 6px;
        background-color: white;
        border: 1px solid #E5E7EB;
    }
    
    .task-checkbox-container.completed {
        background-color: #D1FAE5;
        border-color: #34D399;
    }
    
    /* Motivation and Tips */
    .motivation-box {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 1.5rem;
        border-radius: 12px;
        margin: 1rem 0;
    }
    
    .simple-tip {
        background-color: #F0F9FF;
        padding: 1rem;
        border-radius: 8px;
        margin: 0.5rem 0;
        border-left: 4px solid #3B82F6;
    }
    
    .daily-tips {
        background-color: #EFF6FF;
        padding: 1rem;
        border-radius: 10px;
        margin: 0.5rem 0;
        border-left: 4px solid #10B981;
        font-size: 0.9rem;
    }
    
    /* Streak Badge */
    .streak-badge {
        background: linear-gradient(135deg, #FFD700 0%, #FFA500 100%);
        color: #000;
        padding: 5px 15px;
        border-radius: 20px;
        font-weight: bold;
        font-size: 0.9rem;
        display: inline-flex;
        align-items: center;
        gap: 5px;
    }
    
    .day-badge {
        display: inline-block;
        padding: 2px 8px;
        border-radius: 12px;
        font-size: 0.75rem;
        font-weight: 600;
        background-color: #DBEAFE;
        color: #1E40AF;
        margin-right: 5px;
    }
    
    /* Test Elements */
    .question-box {
        background-color: white;
        padding: 1rem;
        border-radius: 8px;
        margin: 1rem 0;
        border-left: 4px solid #3B82F6;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    
    .option-item {
        padding: 12px;
        margin: 8px 0;
        border-radius: 8px;
        border: 1px solid #E5E7EB;
        cursor: pointer;
        transition: all 0.3s;
        background-color: white;
    }
    
    .option-item:hover {
        background-color: #F3F4F6;
        transform: translateX(5px);
    }
    
    .option-item.selected {
        background-color: #DBEAFE;
        border-color: #3B82F6;
        border-width: 2px;
    }
    
    .option-item.correct {
        background-color: #D1FAE5;
        border-color: #10B981;
    }
    
    .option-item.incorrect {
        background-color: #FEE2E2;
        border-color: #EF4444;
    }
    
    /* Timer */
    .timer-box {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 1rem;
        border-radius: 10px;
        text-align: center;
        font-size: 1.5rem;
        font-weight: bold;
        margin: 1rem 0;
        position: fixed;
        top: 20px;
        right: 20px;
        z-index: 1000;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        min-width: 150px;
    }
    
    .timer-warning {
        background: linear-gradient(135deg, #FF6B6B 0%, #FF8E53 100%);
        animation: pulse 1s infinite;
    }
    
    @keyframes pulse {
        0% { transform: scale(1); }
        50% { transform: scale(1.05); }
        100% { transform: scale(1); }
    }
    
    /* Score Display */
    .score-display {
        font-size: 3rem;
        font-weight: bold;
        text-align: center;
        color: #1E3A8A;
        margin: 1rem 0;
    }
    
    /* Certificate */
    .certificate-box {
        border: 3px solid #FFD700;
        padding: 2rem;
        border-radius: 15px;
        text-align: center;
        background: linear-gradient(135deg, #fff8e1 0%, #fff3e0 100%);
        margin: 2rem 0;
    }
    
    /* Buttons */
    .back-button {
        background-color: #3B82F6;
        color: white;
        padding: 10px 20px;
        border-radius: 8px;
        border: none;
        cursor: pointer;
        font-weight: bold;
        margin-top: 1rem;
        display: inline-block;
        text-decoration: none;
    }
    
    .back-button:hover {
        background-color: #2563EB;
    }
    
    /* Full Text Options for Tests */
    .option-full-text {
        white-space: normal !important;
        word-wrap: break-word !important;
        height: auto !important;
        min-height: 60px !important;
        display: flex !important;
        align-items: center !important;
    }
    
    /* General Layout Improvements */
    .stTabs [data-baseweb="tab-list"] {
        gap: 24px;
    }
    
    .stTabs [data-baseweb="tab"] {
        height: 50px;
        white-space: pre-wrap;
        background-color: #F8FAFC;
        border-radius: 4px 4px 0px 0px;
        gap: 1px;
        padding-top: 10px;
        padding-bottom: 10px;
        font-weight: 600;
    }
</style>
""", unsafe_allow_html=True)

# ============================================
# 🔧 SESSION STATE INITIALIZATION - FIXED
# ============================================

def init_session_state():
    """Initialize all session state variables"""
    session_vars = {
        'study_plan': None,
        'progress': {},
        'subject': "",
        'api_key': GEMINI_API_KEY,
        'raw_response': "",
        'show_debug': False,
        'study_days': ["Mon", "Tue", "Wed", "Thu", "Fri"],
        'plan_loaded': False,
        'current_week': 1,
        'task_states': {},
        'weekly_streak': 0,
        'longest_weekly_streak': 0,
        'task_completion_dates': {},
        'daily_tips': {},
        'test_results': {},
        'current_test': None,
        'test_start_time': None,
        'test_questions': [],
        'user_test_answers': {},
        'test_completed': False,
        'timer_seconds': 1200,
        'flexible_completion': True,  # ADDED THIS LINE
        'daily_progress': []  # ADDED THIS LINE
    }
    
    for key, value in session_vars.items():
        if key not in st.session_state:
            st.session_state[key] = value

init_session_state()

# ============================================
# 🛠️ UPDATED UTILITY FUNCTIONS WITH BETTER JSON PARSING
# ============================================

def extract_and_fix_json(text, subject="General"):
    """
    Extract JSON from text and fix common issues with improved handling
    """
    if not text or not isinstance(text, str):
        if st.session_state.show_debug:
            st.warning("⚠️ No text to parse in extract_and_fix_json")
        return None
    
    # Store raw response for debug
    st.session_state.raw_response = text
    
    # Clean the text - more aggressive cleaning
    text = text.strip()
    
    # Check if response indicates safety blocking
    if "safety" in text.lower() or "blocked" in text.lower() or "not allowed" in text.lower():
        if st.session_state.show_debug:
            st.error("❌ Response indicates safety blocking")
        return None
    
    # Remove ALL markdown code blocks and any surrounding text
    text = re.sub(r'.*?```(?:json)?\s*', '', text, flags=re.IGNORECASE | re.DOTALL)
    text = re.sub(r'```.*', '', text, flags=re.IGNORECASE | re.DOTALL)
    
    # Remove any non-JSON text before first {
    first_brace = text.find('{')
    if first_brace > 0:
        text = text[first_brace:]
    
    # Remove any non-JSON text after last }
    last_brace = text.rfind('}')
    if last_brace != -1:
        text = text[:last_brace+1]
    
    # Show cleaned text in debug
    if st.session_state.show_debug:
        with st.expander("🧹 Cleaned Text", expanded=False):
            st.code(text[:500] + "..." if len(text) > 500 else text)
    
    # Find JSON boundaries
    start_idx = text.find('{')
    end_idx = text.rfind('}')
    
    if start_idx == -1 or end_idx == -1 or end_idx <= start_idx:
        if st.session_state.show_debug:
            st.error(f"❌ No JSON found. Start: {start_idx}, End: {end_idx}")
            st.info(f"Text preview: {text[:200]}...")
        return None
    
    json_str = text[start_idx:end_idx+1]
    
    # Show JSON string in debug
    if st.session_state.show_debug:
        with st.expander("🔍 JSON String Being Parsed", expanded=False):
            st.code(json_str[:800] + "..." if len(json_str) > 800 else json_str)
    
    try:
        # First try: Direct parse
        parsed_json = json.loads(json_str)
        if st.session_state.show_debug:
            st.success("✅ JSON parsed successfully on first try")
        return parsed_json
    except json.JSONDecodeError as e:
        if st.session_state.show_debug:
            st.warning(f"⚠️ First parse failed: {str(e)}")
            st.info("Trying to fix JSON...")
        
        # Second try: Apply advanced fixes
        json_str = fix_json_string_advanced(json_str)
        try:
            parsed_json = json.loads(json_str)
            if st.session_state.show_debug:
                st.success("✅ JSON parsed successfully after advanced fixes")
            return parsed_json
        except json.JSONDecodeError as e2:
            if st.session_state.show_debug:
                st.error(f"❌ Second parse failed: {str(e2)}")
                # Try one more time with even more aggressive fixing
                try:
                    json_str = fix_json_string_aggressive(json_str)
                    parsed_json = json.loads(json_str)
                    if st.session_state.show_debug:
                        st.success("✅ JSON parsed after aggressive fixes")
                    return parsed_json
                except json.JSONDecodeError as e3:
                    st.error(f"❌ Final parse failed: {str(e3)}")
            return None

def fix_json_string_advanced(json_str):
    """
    Fix common JSON issues in AI responses based on actual patterns observed
    """
    if not json_str:
        return "{}"
    
    # 1. Fix unquoted keys (e.g., week: 1 -> "week": 1)
    json_str = re.sub(r'([{,]\s*)([a-zA-Z_][a-zA-Z0-9_]*)\s*:', r'\1"\2":', json_str)
    
    # 2. Fix trailing commas before } or ]
    json_str = re.sub(r',\s*}', '}', json_str)
    json_str = re.sub(r',\s*]', ']', json_str)
    
    # 3. Fix missing commas between objects in arrays
    json_str = re.sub(r'}\s*{', '},{', json_str)
    
    # 4. Fix missing commas after values
    json_str = re.sub(r'(?<!")(\d+\.?\d*|true|false|null)(\s*"[a-zA-Z_])', r'\1,\2', json_str, flags=re.IGNORECASE)
    
    # 5. Fix unescaped quotes inside strings
    lines = json_str.split('\n')
    fixed_lines = []
    for line in lines:
        # Count quotes in line
        quote_count = line.count('"')
        if quote_count % 2 != 0:
            # Find problematic quotes and escape them
            parts = line.split('"')
            if len(parts) > 1:
                # Escape quotes that are inside strings
                for i in range(1, len(parts)-1, 2):
                    if parts[i] and not parts[i].endswith('\\'):
                        # This is a string value that might have unescaped quotes
                        pass
        fixed_lines.append(line)
    json_str = '\n'.join(fixed_lines)
    
    # 6. Fix markdown bold/italic in strings (remove ** and *)
    json_str = re.sub(r'\*\*(.*?)\*\*', r'\1', json_str)
    json_str = re.sub(r'\*(.*?)\*', r'\1', json_str)
    
    # 7. Remove any control characters except newlines
    json_str = ''.join(char for char in json_str if ord(char) >= 32 or char in '\n\r\t')
    
    # 8. Fix common pattern: "column 6" -> this usually means missing comma in array
    # Look for patterns like: "something": [ ... "item1" "item2" ... ]
    json_str = re.sub(r'("[^"]*")(\s+"[^"]*")', r'\1,\2', json_str)
    
    return json_str

def fix_json_string_aggressive(json_str):
    """
    Even more aggressive JSON fixing for problematic AI responses
    """
    if not json_str:
        return "{}"
    
    # Try to find the main JSON structure
    start = json_str.find('{')
    end = json_str.rfind('}')
    
    if start >= 0 and end > start:
        json_str = json_str[start:end+1]
    
    # Replace single quotes with double quotes (carefully)
    # Only replace single quotes that are not escaped and appear to be string delimiters
    in_string = False
    result = []
    i = 0
    while i < len(json_str):
        if json_str[i] == '"' and (i == 0 or json_str[i-1] != '\\'):
            in_string = not in_string
            result.append('"')
        elif json_str[i] == "'" and not in_string:
            # Check if this looks like a string delimiter
            if i > 0 and json_str[i-1] in [' ', ':', ',', '[', '{']:
                result.append('"')
            else:
                result.append("'")
        else:
            result.append(json_str[i])
        i += 1
    
    json_str = ''.join(result)
    
    # Fix missing commas in arrays (aggressive)
    # Pattern: value whitespace value where value is string, number, or boolean
    json_str = re.sub(r'("[^"]*"|\d+\.?\d*|true|false|null)(\s+)(?="|\d|t|f|n|\[|\{)', r'\1,\2', json_str, flags=re.IGNORECASE)
    
    return json_str

def extract_test_json(text):
    """
    Extract test JSON from AI response
    """
    if not text or not isinstance(text, str):
        return None
    
    text = text.strip()
    
    # Remove markdown code blocks
    text = re.sub(r'```json\s*', '', text, flags=re.IGNORECASE)
    text = re.sub(r'```\s*', '', text, flags=re.IGNORECASE)
    
    # Find JSON
    start_idx = text.find('{')
    end_idx = text.rfind('}')
    
    if start_idx == -1 or end_idx == -1:
        return None
    
    json_str = text[start_idx:end_idx+1]
    
    try:
        return json.loads(json_str)
    except json.JSONDecodeError:
        # Apply fixes
        json_str = fix_json_string_advanced(json_str)
        try:
            return json.loads(json_str)
        except:
            return None

def calculate_learning_quality(days_available, hours_per_day, study_days_count):
    """
    Calculate learning quality metrics
    """
    total_hours = days_available * hours_per_day
    required_hours = 80  # Default for medium complexity
    
    coverage = min(100, (total_hours / required_hours) * 100)
    
    quality_score = 0
    if 1 <= hours_per_day <= 3:
        quality_score += 30
    elif hours_per_day > 3:
        quality_score += 20
    else:
        quality_score += 10
    
    if 3 <= study_days_count <= 5:
        quality_score += 30
    elif study_days_count > 5:
        quality_score += 20
    else:
        quality_score += 10
    
    if coverage >= 80:
        quality_score += 40
    elif coverage >= 50:
        quality_score += 20
    else:
        quality_score += 10
    
    return {
        "quality_score": min(100, quality_score),
        "coverage_percentage": coverage,
        "total_hours": total_hours,
        "required_hours": required_hours,
        "weekly_hours": hours_per_day * study_days_count,
        "is_optimal": quality_score >= 60
    }

def create_generic_plan(subject, days_available):
    """
    Create generic plan if AI fails
    """
    weeks = max(1, days_available // 7)
    daily_hours = 2
    study_days = ["Mon", "Tue", "Wed", "Thu", "Fri"]
    
    plan = {
        "subject": subject,
        "total_weeks": weeks,
        "daily_hours": daily_hours,
        "study_days": study_days,
        "topics_allocation": {
            "Fundamentals": 35,
            "Core Concepts": 30,
            "Practice": 25,
            "Advanced": 10
        },
        "weekly_schedule": [],
        "study_tips": [
            "Set specific daily goals",
            "Review previous material",
            "Take regular breaks",
            "Practice what you learn"
        ],
        "success_metrics": [
            f"Complete {weeks} weeks of study",
            "Demonstrate understanding through practice",
            "Apply knowledge in real scenarios"
        ]
    }
    
    for week_num in range(1, min(weeks, 12) + 1):
        daily_tasks = []
        for day_name in study_days:
            task1 = f"Study {subject} topics"
            task2 = f"Practice {subject} exercises"
            daily_tasks.append([task1, task2])
        
        week_plan = {
            "week": week_num,
            "focus": f"Week {week_num}: {subject} Fundamentals",
            "objectives": ["Understand core concepts", "Complete practice exercises"],
            "daily_tasks": daily_tasks,
            "day_focus": [f"Learn basic concepts", "Practice exercises", "Review", "Apply knowledge", "Project work"],
            "week_summary": f"This week focuses on building fundamental knowledge of {subject}.",
            "resources": ["Online courses", "Documentation", "Practice platforms"],
            "milestone": f"Complete Week {week_num} objectives",
            "completed": False,
            "tasks_completed": 0,
            "total_tasks": len(study_days) * 2
        }
        plan["weekly_schedule"].append(week_plan)
    
    return plan

# ============================================
# 🤖 UPDATED AI STUDY PLAN GENERATION WITH BETTER PROMPT
# ============================================

def generate_study_plan(api_key, **kwargs):
    """
    Generate study plan using Gemini with safety handling
    """
    if not api_key:
        st.error("❌ No API key provided")
        return create_generic_plan(kwargs.get('subject', 'General'), 
                                  kwargs.get('days_available', 60))
    
    # Extract parameters
    subject = kwargs.get('subject', 'General Learning')
    days_available = kwargs.get('days_available', 60)
    hours_per_day = kwargs.get('hours_per_day', 2)
    current_level = kwargs.get('current_level', 'Beginner')
    intensity = kwargs.get('intensity', 'Moderate')
    study_days = kwargs.get('study_days', ["Mon", "Tue", "Wed", "Thu", "Fri"])
    
    weeks = max(1, days_available // 7)
    
    # IMPROVED PROMPT with better JSON formatting instructions
    prompt = f"""Create a detailed {weeks}-week study plan for learning: {subject}

Study schedule: {hours_per_day} hours/day, {len(study_days)} days/week ({', '.join(study_days)})
Current level: {current_level}
Intensity: {intensity}

IMPORTANT: Return ONLY valid JSON with this exact structure:

{{
  "subject": "{subject}",
  "total_weeks": {weeks},
  "daily_hours": {hours_per_day},
  "study_days": {json.dumps(study_days)},
  "topics_allocation": {{
    "Topic 1": percentage_number,
    "Topic 2": percentage_number
  }},
  "weekly_schedule": [
    {{
      "week": 1,
      "focus": "Week focus title",
      "objectives": ["Objective 1", "Objective 2"],
      "daily_tasks": [
        ["Task 1 for Monday", "Task 2 for Monday"],
        ["Task 1 for Tuesday", "Task 2 for Tuesday"],
        ["Task 1 for Wednesday", "Task 2 for Wednesday"],
        ["Task 1 for Thursday", "Task 2 for Thursday"],
        ["Task 1 for Friday", "Task 2 for Friday"]
      ],
      "day_focus": [
        "Monday: Introduction",
        "Tuesday: Practice", 
        "Wednesday: Deep Dive",
        "Thursday: Review",
        "Friday: Project"
      ],
      "week_summary": "Brief summary of week's learning objectives",
      "resources": ["Resource 1", "Resource 2"],
      "milestone": "What to achieve by week's end"
    }}
  ],
  "study_tips": ["Tip 1", "Tip 2"],
  "success_metrics": ["Metric 1", "Metric 2"]
}}

CRITICAL RULES FOR VALID JSON:
1. Use double quotes ONLY for ALL strings and keys
2. Separate array items with commas
3. End object properties with commas (except last one before closing brace)
4. topics_allocation values must be NUMBERS (not strings with %)
5. Include EXACTLY {weeks} weeks in weekly_schedule array
6. Each week must have daily_tasks array with {len(study_days)} sub-arrays
7. NO trailing commas after last array element or object property
8. NO markdown formatting (no **bold**, no *italic*)
9. Escape quotes inside strings with backslash (\")
10. Ensure all brackets and braces are properly closed

Make the plan practical, educational, and appropriate for all ages."""

    try:
        # Configure API with safety settings
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.5-flash')
        
        # Show prompt in debug mode
        if st.session_state.show_debug:
            with st.expander("📝 Prompt Sent to AI", expanded=False):
                st.code(prompt[:800] + "..." if len(prompt) > 800 else prompt)
        
        # Generate response with safety settings
        response = model.generate_content(
            prompt,
            generation_config=genai.GenerationConfig(
                max_output_tokens=4000,
                temperature=0.7,
                top_p=0.95,
                top_k=40
            ),
            safety_settings=[
                {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
                {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
                {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
                {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}
            ]
        )
        
        # SAFE way to get response text
        if response and hasattr(response, 'text'):
            raw_text = response.text
        else:
            raw_text = ""
            
        st.session_state.raw_response = raw_text
        
        # Show raw response in debug mode
        if st.session_state.show_debug:
            with st.expander("🔍 Raw AI Response", expanded=False):
                if raw_text:
                    st.text_area("AI Response", raw_text, height=200, key="ai_response_raw_gen")
                else:
                    st.warning("No response text received")
        
        # Extract JSON
        plan = extract_and_fix_json(raw_text, subject)
        
        if plan:
            # Show success in debug mode
            if st.session_state.show_debug:
                st.success("✅ JSON successfully extracted")
                with st.expander("📋 Extracted Plan Preview", expanded=False):
                    st.json({
                        "subject": plan.get('subject'),
                        "total_weeks": plan.get('total_weeks'),
                        "weekly_schedule_length": len(plan.get('weekly_schedule', [])),
                        "first_week_focus": plan.get('weekly_schedule', [{}])[0].get('focus', 'N/A') if plan.get('weekly_schedule') else 'N/A'
                    })
            
            # Post-process the plan
            plan['generated_at'] = datetime.now().isoformat()
            plan['total_days'] = days_available
            plan['user_level'] = current_level
            plan['intensity'] = intensity
            
            # Initialize week tracking
            for week in plan.get('weekly_schedule', []):
                week['completed'] = week.get('completed', False)
                week['tasks_completed'] = week.get('tasks_completed', 0)
                week['completed_tasks_info'] = week.get('completed_tasks_info', [])
                
                # Calculate total tasks
                total_tasks = 0
                if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                    for day_tasks in week['daily_tasks']:
                        if isinstance(day_tasks, list):
                            total_tasks += len(day_tasks)
                week['total_tasks'] = total_tasks
            
            return plan
        else:
            if st.session_state.show_debug:
                st.warning("⚠️ Could not extract valid JSON from AI response")
                if not raw_text:
                    st.error("❌ Empty response from AI")
                else:
                    st.info("Falling back to generic plan")
            return create_generic_plan(subject, days_available)
    
    except Exception as e:
        error_msg = str(e)
        st.error(f"❌ AI Generation Error: {error_msg}")
        
        # Handle safety errors
        if "safety" in error_msg.lower() or "finish_reason" in error_msg.lower():
            st.warning("🔒 **Safety restriction triggered.** Try a different subject or simpler request.")
        
        if st.session_state.show_debug:
            with st.expander("🔍 Error Details", expanded=False):
                st.exception(e)
        
        return create_generic_plan(subject, days_available)

def generate_weekly_test(api_key, week_number, weekly_tasks, subject):
    """
    Generate test for a specific week
    """
    if not api_key:
        return None
    
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.5-flash')
        
        prompt = f"""Create 5 multiple choice questions for Week {week_number} of {subject}.

Return valid JSON with this structure:

{{
  "questions": [
    {{
      "question": "Clear, complete question?",
      "options": {{
        "A": "Complete option A text",
        "B": "Complete option B text",
        "C": "Complete option C text",
        "D": "Complete option D text"
      }},
      "correct_answer": "A",
      "explanation": "Brief explanation"
    }}
  ]
}}

Make ALL options complete sentences. Include exactly 5 questions."""

        response = model.generate_content(
            prompt,
            generation_config=genai.GenerationConfig(
                max_output_tokens=1500,
                temperature=0.7
            )
        )
        test_text = response.text.strip()
        st.session_state.raw_response = test_text  # Store for debug
        
        return extract_test_json(test_text)
    
    except Exception as e:
        return None


def calculate_progress_status(plan, completed_tasks, total_tasks, days_passed, total_days):
    """
    Calculate if user is on track, behind, or ahead
    """
    if total_tasks == 0:
        return "Not Started"
    
    task_completion_rate = completed_tasks / total_tasks
    time_passed_ratio = days_passed / total_days if total_days > 0 else 0
    
    if task_completion_rate >= time_passed_ratio + 0.1:
        return "Ahead of Schedule"
    elif task_completion_rate >= time_passed_ratio - 0.1:
        return "On Track"
    else:
        return "Behind Schedule"

def calculate_weekly_streak(plan):
    """
    Calculate current and longest weekly streak from completed weeks
    """
    if not plan or 'weekly_schedule' not in plan:
        return 0, 0
    
    completed_weeks = []
    for week in plan['weekly_schedule']:
        if week.get('completed', False):
            week_num = week['week']
            completion_date = week.get('completion_date', '2000-01-01')
            completed_weeks.append((week_num, completion_date))
    
    if not completed_weeks:
        return 0, 0
    
    completed_weeks.sort(key=lambda x: x[0])
    
    current_streak = 1
    longest_streak = 1
    temp_streak = 1
    
    for i in range(1, len(completed_weeks)):
        if completed_weeks[i][0] == completed_weeks[i-1][0] + 1:
            temp_streak += 1
            longest_streak = max(longest_streak, temp_streak)
            
            if completed_weeks[i][0] >= st.session_state.current_week - 1:
                current_streak = temp_streak
        else:
            temp_streak = 1
    
    return current_streak, longest_streak

def get_motivational_message(progress_percentage, status, weekly_streak):
    """
    Get motivational message based on progress and streak
    """
    import random
    
    messages = {
        "Not Started": [
            "🎯 Every expert was once a beginner. Start your journey today!",
            "🚀 The first step is always the hardest. You've got this!",
            "📚 Your future self will thank you for starting today."
        ],
        "Behind Schedule": [
            "💪 It's not about being perfect, it's about being consistent. Keep going!",
            "🔄 Falling behind is normal. Adjust your pace and continue forward.",
            "🌟 Progress is progress, no matter how small. Every step counts!"
        ],
        "On Track": [
            "✅ Amazing! You're right on schedule. Keep up the great work!",
            "🔥 Your consistency is paying off. You're doing fantastic!",
            "🏆 You're making steady progress toward your goal. Stay focused!"
        ],
        "Ahead of Schedule": [
            "⚡ You're crushing it! Consider challenging yourself with advanced topics.",
            "🎖️ Ahead of schedule! Your dedication is inspiring.",
            "🚀 You're making rapid progress. Consider helping others or exploring deeper."
        ]
    }
    
    if status in messages:
        return random.choice(messages[status])
    
    if progress_percentage < 30:
        return "🌱 You're planting the seeds of knowledge. Keep nurturing them!"
    elif progress_percentage < 70:
        return "🌊 You're in the flow. Maintain this momentum!"
    else:
        return "🏔️ You're approaching the summit. The view will be worth it!"

def save_plan(plan, subject):
    """
    Save study plan to JSON file
    """
    import os
    os.makedirs("data", exist_ok=True)
    filename = f"data/{subject.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(plan, f, indent=4, ensure_ascii=False)
    
    return filename

def load_progress(subject):
    """
    Load progress from file
    """
    progress_file = f"data/{subject.replace(' ', '_')}_progress.json"
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            try:
                return json.load(f)
            except:
                return {}
    return {}

def calculate_weekly_task_completion_with_flexibility(plan, subject):
    """
    Calculate completed tasks per week with flexible completion dates
    """
    weekly_completion = {}
    
    if plan and plan.get('weekly_schedule'):
        for week in plan['weekly_schedule']:
            week_num = week['week']
            total_tasks = 0
            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                for day_tasks in week['daily_tasks']:
                    if isinstance(day_tasks, list):
                        total_tasks += len(day_tasks)
            
            weekly_completion[week_num] = {
                'assigned_tasks': total_tasks,
                'completed_in_week': 0,
                'completed_elsewhere': 0,
                'total_completed': 0
            }
        
        # Count completions
        for week in plan['weekly_schedule']:
            week_num = week['week']
            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                for day_idx, day_tasks in enumerate(week['daily_tasks']):
                    if isinstance(day_tasks, list):
                        for task_idx, task in enumerate(day_tasks):
                            task_key = f"task_week{week_num}_day{day_idx}_task{task_idx}_{subject}"
                            if st.session_state.get(task_key, False):
                                weekly_completion[week_num]['completed_in_week'] += 1
        
        # Calculate totals
        for week_num in weekly_completion:
            weekly_completion[week_num]['total_completed'] = weekly_completion[week_num]['completed_in_week']
    
    return weekly_completion

def save_progress(subject, progress_data):
    """
    Save progress to file
    """
    import os
    os.makedirs("data", exist_ok=True)
    progress_file = f"data/{subject.replace(' ', '_')}_progress.json"
    with open(progress_file, 'w', encoding='utf-8') as f:
        json.dump(progress_data, f, indent=4, ensure_ascii=False)

def analyze_study_patterns(daily_progress, study_days):
    """
    Analyze study patterns and provide recommendations
    """
    if not daily_progress:
        return []
    
    recommendations = []
    
    try:
        df = pd.DataFrame(daily_progress)
        df['date'] = pd.to_datetime(df['date'], errors='coerce')
        df = df.dropna(subset=['date'])
        
        if len(df) < 2:
            return []
        
        date_range = (df['date'].max() - df['date'].min()).days + 1
        study_days_count = len(df)
        
        if date_range > 0:
            study_frequency = (study_days_count / date_range) * 100
            
            if study_frequency < 30:
                recommendations.append(f"**📅 Increase Study Frequency:** You're studying on {study_frequency:.0f}% of days. Try to study more regularly to build momentum.")
            elif study_frequency > 80:
                recommendations.append(f"**🎯 Excellent Consistency:** You're studying on {study_frequency:.0f}% of days! This consistency will lead to great results.")
        
        # Hours analysis
        if 'hours' in df.columns:
            df['hours'] = pd.to_numeric(df['hours'], errors='coerce')
            avg_hours = df['hours'].mean()
            
            if avg_hours < 1:
                recommendations.append(f"**⏰ Increase Study Duration:** Average {avg_hours:.1f} hours per session. Consider longer focused sessions for better learning.")
            elif avg_hours > 4:
                recommendations.append(f"**🧠 Manage Study Load:** Average {avg_hours:.1f} hours per session is high. Consider breaking into shorter sessions to prevent burnout.")
            else:
                recommendations.append(f"**✅ Optimal Study Duration:** Average {avg_hours:.1f} hours per session is perfect for focused learning.")
    
    except Exception as e:
        pass
    
    return recommendations

def create_completion_certificate(subject, completion_percentage, total_tasks, completed_tasks, total_weeks, completed_weeks):
    """
    Create a completion certificate HTML
    """
    # Determine certificate level
    if completion_percentage >= 95:
        achievement = "with Honors"
        color = "#FFD700"
        message = "Exceptional Achievement"
    elif completion_percentage >= 90:
        achievement = "Successfully Completed"
        color = "#C0C0C0"
        message = "Excellent Completion"
    elif completion_percentage >= 70:
        achievement = "Good Progress"
        color = "#3B82F6"
        message = "Solid Progress"
    else:
        achievement = "In Progress"
        color = "#10B981"
        message = "Learning Journey"
    
    current_date = datetime.now().strftime("%B %d, %Y")
    
    certificate_html = f'''
<!DOCTYPE html>
<html>
<head>
    <title>Certificate - {subject}</title>
    <style>
        body {{
            font-family: Arial, sans-serif;
            padding: 20px;
            text-align: center;
        }}
        .certificate {{
            border: 3px solid {color};
            padding: 40px;
            max-width: 800px;
            margin: 0 auto;
            background: white;
            border-radius: 10px;
        }}
        h1 {{
            color: #1E3A8A;
        }}
        h2 {{
            color: {color};
        }}
        .stats {{
            display: flex;
            justify-content: space-around;
            margin: 30px 0;
        }}
        .stat-item {{
            text-align: center;
        }}
        .stat-value {{
            font-size: 24px;
            font-weight: bold;
            color: {color};
        }}
        .footer {{
            margin-top: 40px;
            color: #666;
        }}
    </style>
</head>
<body>
    <div class="certificate">
        <h1>🎓 Certificate of Achievement</h1>
        <h2>{message}</h2>
        <h3>{subject}</h3>
        <p>Awarded to recognize outstanding dedication and completion of study program</p>
        
        <div class="stats">
            <div class="stat-item">
                <div class="stat-value">{completion_percentage:.1f}%</div>
                <div>Overall Progress</div>
            </div>
            <div class="stat-item">
                <div class="stat-value">{completed_tasks}/{total_tasks}</div>
                <div>Tasks Completed</div>
            </div>
            <div class="stat-item">
                <div class="stat-value">{completed_weeks}/{total_weeks}</div>
                <div>Weeks Completed</div>
            </div>
        </div>
        
        <p><strong>Achievement:</strong> {achievement}</p>
        <p><strong>Date:</strong> {current_date}</p>
        
        <div class="footer">
            <p>Generated by AI Study Planner Pro</p>
            <p>Certificate ID: {subject.replace(" ", "").upper()}{datetime.now().strftime("%Y%m%d")}</p>
        </div>
    </div>
</body>
</html>
'''
    return certificate_html

# ============================================
# 🎪 MAIN APPLICATION LAYOUT - SIMPLIFIED
# ============================================

# Main header
st.markdown('<h1 class="main-header">📚 AI Study Planner Pro</h1>', unsafe_allow_html=True)
st.markdown('<p class="tagline">Plan → Track → Achieve | Your personal AI study assistant</p>', unsafe_allow_html=True)

# API Key Status
if not st.session_state.api_key:
    st.error("⚠️ **API Key Missing:** Please set GEMINI_API_KEY in Hugging Face Secrets.")
else:
    st.success("✅ **API Key Loaded:** Ready to generate AI-powered study plans!")

# Create tabs
tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs([
    "🎯 Define Goal", 
    "📅 Study Plan", 
    "✅ Track Progress", 
    "📈 Analytics", 
    "🧪 Test", 
    "📤 Export"
])

# ============================================
# TAB 1: DEFINE GOAL - WITH DROPDOWN OPTIONS
# ============================================

with tab1:
    st.markdown('<h2 class="sub-header">Define Your Learning Goal</h2>', unsafe_allow_html=True)
    
    # Debug Mode Toggle
    col_debug1, col_debug2 = st.columns([1, 3])
    with col_debug1:
        st.session_state.show_debug = st.checkbox("🔍 Debug Mode", value=False, key="debug_mode_tab1")
    
    # API Status
    with col_debug2:
        if not st.session_state.api_key:
            st.error("⚠️ **API Key Missing:** Set GEMINI_API_KEY in Hugging Face Secrets")
        else:
            st.success("✅ **API Key Loaded**")
    
    # Test API Button with SAFETY FIX
    if st.session_state.api_key and st.session_state.show_debug:
        if st.button("Test API Connection", type="secondary", key="test_api_tab1"):
            try:
                genai.configure(api_key=st.session_state.api_key)
                model = genai.GenerativeModel('gemini-2.5-flash')
                
                # Use a safer, more neutral prompt
                test_response = model.generate_content(
                    "Please respond with just the word 'Connected'",
                    generation_config=genai.GenerationConfig(
                        max_output_tokens=10,
                        temperature=0.1
                    )
                )
                
                # Safely get response text
                if test_response and hasattr(test_response, 'text') and test_response.text:
                    st.success(f"✅ **API Connected:** {test_response.text}")
                else:
                    st.error("❌ No valid response from API")
                    
            except Exception as e:
                st.error(f"❌ **API Error:** {str(e)}")
                if "safety" in str(e).lower():
                    st.info("🔒 **Safety restriction detected.** Try a different subject or rephrase your request.")
    
    st.markdown("---")
    
    # Option 1: Upload Existing Plan
    st.info("📤 **Option 1: Upload Existing Plan**")
    uploaded_file = st.file_uploader("Upload saved study plan (JSON)", type=['json'], key="upload_plan_tab1")
    
    if uploaded_file is not None and not st.session_state.plan_loaded:
        try:
            plan = json.load(uploaded_file)
            st.session_state.study_plan = plan
            st.session_state.subject = plan['subject']
            st.session_state.plan_loaded = True
            st.success(f"✅ Plan loaded: {plan['subject']}")
        except Exception as e:
            st.error(f"❌ Invalid plan file: {str(e)}")
    
    st.markdown("---")
    
    # Option 2: Create New Plan - WITH DROPDOWN + CUSTOM
    st.info("🎯 **Option 2: Create New Plan**")
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Subject selection: Dropdown OR Custom input
        subject_options = [
            "Select a subject or enter custom...",
            "Data Science & Machine Learning",
            "Python Programming",
            "Web Development (Full Stack)",
            "Artificial Intelligence",
            "Cloud Computing (AWS/Azure/GCP)",
            "Cybersecurity",
            "Mobile App Development",
            "Digital Marketing",
            "Business Analytics",
            "Language Learning (Spanish/French/English)",
            "Project Management",
            "Graphic Design",
            "Game Development"
        ]
        
        selected_option = st.selectbox(
            "Choose a subject or enter custom:",
            options=subject_options,
            index=0,
            key="subject_select_tab1"
        )
        
        if selected_option == "Select a subject or enter custom...":
            subject = st.text_input(
                "Or enter your own subject:",
                placeholder="e.g., Quantum Computing, Ethical Hacking, Music Production...",
                key="custom_subject_tab1"
            )
        else:
            subject = selected_option
        
        col1a, col2a = st.columns(2)
        with col1a:
            hours_per_day = st.slider(
                "Hours per day:",
                min_value=1,
                max_value=8,
                value=2,
                key="hours_slider_tab1"
            )
        
        with col2a:
            days_available = st.slider(
                "Days available:",
                min_value=7,
                max_value=365,
                value=60,
                key="days_slider_tab1"
            )
        
        study_days = st.multiselect(
            "Study days:",
            ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"],
            default=["Mon", "Tue", "Wed", "Thu", "Fri"],
            key="study_days_select_tab1"
        )
        
    with col2:
        current_level = st.selectbox(
            "Your level:",
            ["Beginner", "Intermediate", "Advanced"],
            index=0,
            key="level_select_tab1"
        )
        
        # Add study intensity
        intensity = st.select_slider(
            "Study intensity:",
            options=["Casual", "Moderate", "Intensive"],
            value="Moderate",
            key="intensity_slider_tab1"
        )
    
    st.markdown("---")
    
    # Generate Plan Button
    if st.button("🚀 Generate AI Study Plan", type="primary", use_container_width=True, key="generate_plan_btn_tab1"):
        if not st.session_state.api_key:
            st.error("⚠️ **API Key Required:** Please set GEMINI_API_KEY in Hugging Face Secrets.")
        elif not subject or subject == "Select a subject or enter custom...":
            st.error("⚠️ Please enter or select a subject")
        elif not study_days:
            st.error("⚠️ Please select at least one study day")
        else:
            with st.spinner("🤖 AI is creating your personalized study plan..."):
                try:
                    # Calculate learning quality
                    learning_quality = calculate_learning_quality(
                        days_available, 
                        hours_per_day, 
                        len(study_days)
                    )
                    
                    # Create plan data with intensity
                    plan_data = {
                        'subject': subject,
                        'days_available': days_available,
                        'hours_per_day': hours_per_day,
                        'current_level': current_level,
                        'intensity': intensity,
                        'study_days': study_days
                    }
                    
                    # Show plan data in debug mode
                    if st.session_state.show_debug:
                        with st.expander("📋 Plan Data Sent to AI", expanded=True):
                            st.json(plan_data)
                    
                    # Generate plan
                    plan = generate_study_plan(st.session_state.api_key, **plan_data)
                    
                    if plan:
                        st.session_state.study_plan = plan
                        st.session_state.subject = subject
                        st.session_state.current_week = 1
                        st.session_state.weekly_streak = 0
                        st.session_state.longest_weekly_streak = 0
                        st.session_state.plan_loaded = True
                        
                        st.success("✅ Study plan generated successfully!")
                        st.balloons()
                        
                        # Show plan summary
                        st.info(f"📅 **Schedule:** {len(study_days)} days/week ({', '.join(study_days)}) × {hours_per_day} hours/day")
                        st.info(f"📊 **Duration:** {plan.get('total_weeks', 0)} weeks ({days_available} days)")
                        
                        # Show learning quality
                        if learning_quality['is_optimal']:
                            st.success(f"🎯 **Learning Quality:** {learning_quality['quality_score']:.0f}/100 (Excellent)")
                        else:
                            st.warning(f"📊 **Learning Quality:** {learning_quality['quality_score']:.0f}/100 (Needs improvement)")
                        
                        st.caption(f"Coverage: {learning_quality['coverage_percentage']:.0f}% | Total hours: {learning_quality['total_hours']}h")
                        
                    else:
                        st.error("❌ Failed to generate plan. Check debug mode for details.")
                        
                except Exception as e:
                    st.error(f"❌ Error: {str(e)}")
                    if st.session_state.show_debug:
                        with st.expander("🔍 Error Details", expanded=True):
                            st.exception(e)
    
    # Debug Info Section
    if st.session_state.show_debug and st.session_state.raw_response:
        st.markdown("---")
        st.subheader("🔍 Debug Information")
        
        col_debug_a, col_debug_b = st.columns(2)
        
        with col_debug_a:
            st.metric("Raw Response Length", f"{len(st.session_state.raw_response)} chars")
        
        with col_debug_b:
            if st.session_state.study_plan:
                plan_type = "AI Generated" if st.session_state.study_plan.get('generated_at') else "Generic"
                st.metric("Plan Type", plan_type)
        
        with st.expander("📋 Raw AI Response", expanded=False):
            st.text_area("Full Response", st.session_state.raw_response, height=300, key="raw_response_display_tab1")
        
        if st.session_state.study_plan:
            with st.expander("📊 Parsed Plan Structure", expanded=False):
                st.json(st.session_state.study_plan)

# ============================================
# TAB 2: STUDY PLAN
# ============================================

with tab2:
    if st.session_state.study_plan:
        plan = st.session_state.study_plan
        subject = st.session_state.subject
        study_days = plan.get('study_days', ["Mon", "Tue", "Wed", "Thu", "Fri"])
        
        st.markdown(f'<h2 class="sub-header">📅 Study Plan: {subject}</h2>', unsafe_allow_html=True)
        
        # Calculate progress metrics
        total_tasks = 0
        completed_tasks = 0
        
        for week in plan.get('weekly_schedule', []):
            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                for day_idx, day_tasks in enumerate(week['daily_tasks']):
                    if isinstance(day_tasks, list):
                        total_tasks += len(day_tasks)
                        for task_idx in range(len(day_tasks)):
                            task_key = f"task_week{week['week']}_day{day_idx}_task{task_idx}_{subject}"
                            if st.session_state.get(task_key, False):
                                completed_tasks += 1
        
        progress_percentage = (completed_tasks / total_tasks * 100) if total_tasks > 0 else 0
        
        start_date_str = plan.get('start_date', datetime.now().strftime("%Y-%m-%d"))
        try:
            start_date = datetime.strptime(start_date_str, "%Y-%m-%d")
            days_passed = max(1, (datetime.now() - start_date).days + 1)
        except:
            days_passed = 1
        
        total_days = plan.get('total_days', 60)
        status = calculate_progress_status(plan, completed_tasks, total_tasks, days_passed, total_days)
        
        current_streak, longest_streak = calculate_weekly_streak(plan)
        st.session_state.weekly_streak = current_streak
        st.session_state.longest_weekly_streak = longest_streak
        
        status_class = {
            "Not Started": "status-behind",
            "Behind Schedule": "status-behind",
            "On Track": "status-on-track",
            "Ahead of Schedule": "status-ahead"
        }.get(status, "status-on-track")
        
        # Display metrics
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            weeks = plan.get('total_weeks', 0)
            st.metric("Total Weeks", weeks)
        with col2:
            st.metric("Study Days/Week", len(study_days))
        with col3:
            if current_streak > 0:
                st.markdown(f"""
                <div style="text-align: center;">
                    <div class="streak-badge">🔥 {current_streak} Week Streak</div>
                    <p style="font-size: 0.8rem; color: #666; margin-top: 5px;">
                        Longest: {longest_streak} weeks
                    </p>
                </div>
                """, unsafe_allow_html=True)
            else:
                st.metric("Weekly Streak", "0 weeks")
        with col4:
            daily_hours = plan.get('daily_hours', 2)
            st.metric("Daily Hours", f"{daily_hours}h")
        
        st.info(f"📅 Your study schedule: {len(study_days)} days per week ({', '.join(study_days)})")
        
        # Progress display
        st.markdown(f"""
        <div style="background-color: #F8FAFC; padding: 1rem; border-radius: 10px; margin: 1rem 0;">
            <div style="display: flex; justify-content: space-between; align-items: center;">
                <div>
                    <h4 style="margin: 0;">Progress Status</h4>
                    <p style="margin: 5px 0;">Overall Progress: <strong>{progress_percentage:.1f}%</strong></p>
                    <p style="margin: 5px 0; color: #666;">
                        Weekly Streak: <strong>{current_streak} weeks</strong> • Longest: <strong>{longest_streak} weeks</strong>
                    </p>
                </div>
                <div class="{status_class}" style="font-size: 1.2rem;">
                    {status}
                </div>
            </div>
            <div class="progress-bar">
                <div class="progress-fill" style="width: {min(progress_percentage, 100)}%;"></div>
            </div>
            <p style="text-align: center; margin: 5px 0; color: #6B7280;">
                {completed_tasks} of {total_tasks} tasks completed | Day {min(days_passed, total_days)} of {total_days}
            </p>
        </div>
        """, unsafe_allow_html=True)
        
        # Motivation
        motivation = get_motivational_message(progress_percentage, status, current_streak)
        st.info(f"💡 {motivation}")
        
        # Study Tips
        st.subheader("💡 Simple Study Tips")
        simple_tips = [
            "🎯 **Set Daily Goals**: Know what you want to achieve each day",
            "⏰ **Study in Short Bursts**: 25-30 minutes of focused study",
            "🔄 **Review Regularly**: Spend 5 minutes reviewing yesterday's material",
            "💪 **Practice Daily**: Consistency is more important than perfection",
            "🧠 **Take Breaks**: Stand up and stretch every hour"
        ]
        
        tips_cols = st.columns(2)
        for i, tip in enumerate(simple_tips):
            with tips_cols[i % 2]:
                st.markdown(f'<div class="simple-tip">{tip}</div>', unsafe_allow_html=True)
        
        # Topic Allocation Chart
        if 'topics_allocation' in plan:
            st.subheader("📚 Topic Allocation")
            topics = plan['topics_allocation']
            chart_data = []
            for topic, value in topics.items():
                try:
                    num_value = float(value)
                    if num_value > 0:
                        chart_data.append({"Topic": topic, "Percentage": num_value})
                except:
                    continue
            
            if chart_data:
                df = pd.DataFrame(chart_data)
                fig = px.pie(df, values='Percentage', names='Topic', 
                           title=f"Time Allocation for {subject}", hole=0.3)
                st.plotly_chart(fig, use_container_width=True)
        
        # Weekly Schedule
        st.subheader("📋 Weekly Schedule")
        
        current_week = st.selectbox(
            "Jump to week:",
            options=list(range(1, plan.get('total_weeks', 1) + 1)),
            index=st.session_state.current_week - 1,
            key="week_selector"
        )
        st.session_state.current_week = current_week
        
        if plan.get('weekly_schedule'):
            week_idx = current_week - 1
            if week_idx < len(plan['weekly_schedule']):
                week = plan['weekly_schedule'][week_idx]
                
                with st.expander(f"**Week {week['week']}: {week.get('focus', 'Weekly Focus')}**", expanded=True):
                    # Week Focus Box
                    st.markdown(f"""
                    <div class="week-focus-box">
                        <h3 style="color: white; margin: 0;">📚 Week {week['week']} Focus</h3>
                        <p style="color: white; margin: 10px 0 0 0;">{week.get('focus', 'Weekly learning objectives')}</p>
                    </div>
                    """, unsafe_allow_html=True)
                    
                    col1, col2 = st.columns([2, 1])
                    
                    with col1:
                        st.markdown("### 📋 Objectives")
                        if 'objectives' in week and week['objectives']:
                            for obj in week['objectives']:
                                st.markdown(f"• **{obj}**")
                        
                        st.markdown("### 🎯 Milestone")
                        if 'milestone' in week:
                            st.info(f"**{week['milestone']}**")
                        
                        if week.get('completed') and week.get('completion_date'):
                            st.success(f"✅ **Completed on:** {week['completion_date']}")
                    
                    with col2:
                        week_tasks = week['total_tasks']
                        completed_week_tasks = 0
                        
                        # Count completed tasks for this week
                        if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                            for day_idx, day_tasks in enumerate(week['daily_tasks']):
                                if isinstance(day_tasks, list):
                                    for task_idx in range(len(day_tasks)):
                                        task_key = f"task_week{week['week']}_day{day_idx}_task{task_idx}_{subject}"
                                        if st.session_state.get(task_key, False):
                                            completed_week_tasks += 1
                        
                        week_progress = (completed_week_tasks / week_tasks * 100) if week_tasks > 0 else 0
                        week_completed = (completed_week_tasks == week_tasks)
                        
                        if week_completed:
                            st.success("✅ **Week Completed**")
                            if week.get('completion_date'):
                                st.caption(f"Completed on: {week['completion_date']}")
                        else:
                            st.metric("Week Progress", f"{week_progress:.0f}%")
                            
                            if st.button(f"Mark Week {week['week']} Complete", key=f"complete_week_{week['week']}"):
                                week['completed'] = True
                                week['completion_date'] = datetime.now().strftime("%Y-%m-%d")
                                week['tasks_completed'] = completed_week_tasks
                                
                                new_streak, new_longest = calculate_weekly_streak(plan)
                                st.session_state.weekly_streak = new_streak
                                st.session_state.longest_weekly_streak = max(st.session_state.longest_weekly_streak, new_longest)
                                
                                save_plan(plan, subject)
                                
                                progress_data = load_progress(subject)
                                if progress_data:
                                    progress_data['completed_weeks'] = [w['week'] for w in plan.get('weekly_schedule', []) if w.get('completed', False)]
                                    progress_data['weekly_streak'] = new_streak
                                    progress_data['longest_weekly_streak'] = max(progress_data.get('longest_weekly_streak', 0), new_longest)
                                    progress_data['last_updated'] = datetime.now().isoformat()
                                    save_progress(subject, progress_data)
                                
                                st.success(f"Week {week['week']} marked as complete!")
                                st.balloons()
                                st.rerun()
                    
                    # Week Summary
                    if week.get('week_summary'):
                        st.markdown(f"""
                        <div class="week-summary-box">
                            <h4>📝 Week Summary</h4>
                            <p>{week['week_summary']}</p>
                        </div>
                        """, unsafe_allow_html=True)
                    
                    st.markdown("---")
                    st.markdown(f"### 📝 Tasks for Week {week['week']}")
                    
                    if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                        day_focus = week.get('day_focus', [])
                        
                        for day_idx, day_tasks in enumerate(week['daily_tasks']):
                            if isinstance(day_tasks, list) and day_idx < len(study_days):
                                day_name = study_days[day_idx]
                                
                                st.markdown(f"#### {day_name}")
                                
                                # Show day focus
                                if day_idx < len(day_focus):
                                    st.markdown(f"""
                                    <div class="day-focus-box">
                                        <strong>Focus:</strong> {day_focus[day_idx]}
                                    </div>
                                    """, unsafe_allow_html=True)
                                
                                # Show each task
                                for task_idx, task in enumerate(day_tasks):
                                    task_key = f"task_week{week['week']}_day{day_idx}_task{task_idx}_{subject}"
                                    
                                    if task_key not in st.session_state:
                                        st.session_state[task_key] = False
                                    
                                    task_class = "completed" if st.session_state[task_key] else ""
                                    
                                    # Task container
                                    col_check, col_desc = st.columns([1, 20])
                                    with col_check:
                                        is_completed = st.checkbox(
                                            "",
                                            value=st.session_state[task_key],
                                            key=f"checkbox_{task_key}",
                                            label_visibility="collapsed"
                                        )
                                    
                                    with col_desc:
                                        st.markdown(f"""
                                        <div class="task-checkbox-container {task_class}">
                                            <strong>{task}</strong>
                                        </div>
                                        """, unsafe_allow_html=True)
                                    
                                    if is_completed != st.session_state[task_key]:
                                        st.session_state[task_key] = is_completed
                                        
                                        if is_completed:
                                            today = datetime.now().strftime("%Y-%m-%d")
                                            st.session_state.task_completion_dates[task_key] = today
                                            
                                            # Update week progress
                                            completed_week_tasks = 0
                                            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                                                for d_idx, d_tasks in enumerate(week['daily_tasks']):
                                                    if isinstance(d_tasks, list):
                                                        for t_idx in range(len(d_tasks)):
                                                            tk = f"task_week{week['week']}_day{d_idx}_task{t_idx}_{subject}"
                                                            if st.session_state.get(tk, False):
                                                                completed_week_tasks += 1
                                            
                                            week['tasks_completed'] = completed_week_tasks
                                            week['completed'] = (completed_week_tasks == week['total_tasks'])
                                            
                                            if week['completed'] and not week.get('completion_date'):
                                                week['completion_date'] = today
                                                
                                                new_streak, new_longest = calculate_weekly_streak(plan)
                                                st.session_state.weekly_streak = new_streak
                                                st.session_state.longest_weekly_streak = max(st.session_state.longest_weekly_streak, new_longest)
                                            
                                            save_plan(plan, subject)
                                            
                                            progress_data = load_progress(subject)
                                            if progress_data:
                                                all_completed = 0
                                                for w in plan.get('weekly_schedule', []):
                                                    if 'daily_tasks' in w and isinstance(w['daily_tasks'], list):
                                                        for d_idx, d_tasks in enumerate(w['daily_tasks']):
                                                            if isinstance(d_tasks, list):
                                                                for t_idx in range(len(d_tasks)):
                                                                    tk = f"task_week{w['week']}_day{d_idx}_task{t_idx}_{subject}"
                                                                    if st.session_state.get(tk, False):
                                                                        all_completed += 1
                                                
                                                progress_data['completed_tasks'] = all_completed
                                                progress_data['completed_weeks'] = [w['week'] for w in plan.get('weekly_schedule', []) if w.get('completed', False)]
                                                progress_data['weekly_streak'] = st.session_state.weekly_streak
                                                progress_data['longest_weekly_streak'] = max(progress_data.get('longest_weekly_streak', 0), st.session_state.longest_weekly_streak)
                                                progress_data['last_updated'] = datetime.now().isoformat()
                                                save_progress(subject, progress_data)
                                            
                                            st.rerun()
                                        
                                st.markdown("---")
                    else:
                        st.info("No tasks defined for this week")
    
    else:
        st.info("👈 **Define your learning goal first in the 'Define Goal' tab!**")

# ============================================
# TAB 3: TRACK PROGRESS
# ============================================

with tab3:
    if st.session_state.study_plan:
        subject = st.session_state.subject
        progress_data = load_progress(subject)
        plan = st.session_state.study_plan
        study_days = plan.get('study_days', ["Mon", "Tue", "Wed", "Thu", "Fri"])
        
        st.markdown(f'<h2 class="sub-header">✅ Progress Tracker: {subject}</h2>', unsafe_allow_html=True)
        
        # Calculate progress
        total_tasks = 0
        completed_tasks = 0
        
        for week in plan.get('weekly_schedule', []):
            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                for day_idx, day_tasks in enumerate(week['daily_tasks']):
                    if isinstance(day_tasks, list):
                        total_tasks += len(day_tasks)
                        for task_idx in range(len(day_tasks)):
                            task_key = f"task_week{week['week']}_day{day_idx}_task{task_idx}_{subject}"
                            if st.session_state.get(task_key, False):
                                completed_tasks += 1
        
        progress_percentage = (completed_tasks / total_tasks * 100) if total_tasks > 0 else 0
        
        weekly_completion = calculate_weekly_task_completion_with_flexibility(plan, subject)
        current_streak, longest_streak = calculate_weekly_streak(plan)
        
        # Display metrics
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.metric("Overall Progress", f"{progress_percentage:.1f}%")
        with col2:
            st.metric("Tasks Completed", f"{completed_tasks}/{total_tasks}")
        with col3:
            completed_weeks = sum(1 for week in plan.get('weekly_schedule', []) if week.get('completed', False))
            total_weeks = plan.get('total_weeks', 0)
            st.metric("Weeks Completed", f"{completed_weeks}/{total_weeks}")
        with col4:
            if current_streak > 0:
                st.markdown(f"""
                <div style="text-align: center;">
                    <div class="streak-badge">🔥 {current_streak} Week Streak</div>
                    <p style="font-size: 0.8rem; color: #666; margin-top: 5px;">
                        Longest: {longest_streak} weeks
                    </p>
                </div>
                """, unsafe_allow_html=True)
            else:
                st.metric("Weekly Streak", "0 weeks")
        
        # Weekly Task Completion Chart
        if weekly_completion:
            st.subheader("📊 Weekly Task Completion")
            
            weeks_list = sorted(weekly_completion.keys())
            completed_in_week = [weekly_completion[w]['completed_in_week'] for w in weeks_list]
            total_assigned = [weekly_completion[w]['assigned_tasks'] for w in weeks_list]
            
            df_tasks = pd.DataFrame({
                'Week': [f"Week {w}" for w in weeks_list],
                'Completed': completed_in_week,
                'Total Assigned': total_assigned,
                'Completion %': [(completed_in_week[i] / total_assigned[i] * 100) if total_assigned[i] > 0 else 0 for i in range(len(weeks_list))]
            })
            
            fig = go.Figure()
            fig.add_trace(go.Bar(
                x=df_tasks['Week'],
                y=df_tasks['Completed'],
                name='Completed Tasks',
                marker_color='#10B981',
                text=df_tasks['Completed'],
                textposition='inside'
            ))
            
            remaining = [total_assigned[i] - completed_in_week[i] for i in range(len(weeks_list))]
            fig.add_trace(go.Bar(
                x=df_tasks['Week'],
                y=remaining,
                name='Remaining Tasks',
                marker_color='#E5E7EB',
                text=remaining,
                textposition='inside'
            ))
            
            fig.update_layout(
                title='Weekly Task Completion',
                barmode='stack',
                yaxis_title='Number of Tasks',
                showlegend=True,
                hovermode='x unified'
            )
            
            st.plotly_chart(fig, use_container_width=True)
            st.info(f"📊 **Total tasks completed:** {sum(completed_in_week)} out of {sum(total_assigned)}")
        
        # Weekly Progress Overview
        st.subheader("📈 Weekly Progress Overview")
        
        if plan.get('weekly_schedule'):
            weekly_progress = []
            for week in plan['weekly_schedule']:
                week_num = week['week']
                week_tasks = week['total_tasks']
                
                week_completion = weekly_completion.get(week_num, {'total_completed': 0, 'assigned_tasks': week_tasks})
                completion_percentage = (week_completion['total_completed'] / week_tasks * 100) if week_tasks > 0 else 0
                
                weekly_progress.append({
                    "Week": week_num,
                    "Week_Label": f"Week {week_num}",
                    "Progress": completion_percentage,
                    "Completed": week.get('completed', False),
                    "Completion Date": week.get('completion_date', 'Not completed'),
                    "Tasks Completed": week_completion['total_completed'],
                    "Total Tasks": week_tasks
                })
            
            if weekly_progress:
                df_weekly = pd.DataFrame(weekly_progress)
                df_weekly = df_weekly.sort_values('Week')
                
                fig3 = px.bar(df_weekly, x='Week_Label', y='Progress',
                            title='Weekly Completion Percentage',
                            color='Completed',
                            color_discrete_map={True: '#10B981', False: '#3B82F6'},
                            hover_data=['Tasks Completed', 'Total Tasks', 'Completion Date'])
                
                fig3.add_hline(y=100, line_dash="dash", line_color="green", 
                              annotation_text="Complete", 
                              annotation_position="bottom right")
                
                st.plotly_chart(fig3, use_container_width=True)
                
                with st.expander("📋 Weekly Completion Details", expanded=False):
                    st.dataframe(df_weekly[['Week', 'Tasks Completed', 'Total Tasks', 'Progress', 'Completed', 'Completion Date']], 
                               use_container_width=True)
        
        # Daily Progress Log
        st.subheader("📝 Daily Progress Log")
        
        with st.form("daily_log_form"):
            today = datetime.now().strftime("%Y-%m-%d")
            st.write(f"**Date:** {today}")
            
            col1, col2, col3 = st.columns(3)
            with col1:
                hours_studied = st.number_input("Hours studied today:", 
                                               min_value=0.0, max_value=12.0, 
                                               value=2.0, step=0.5)
            with col2:
                study_hour = st.selectbox(
                    "Study hour of the day:",
                    options=["1st hour (00:00-01:00)", "2nd hour (01:00-02:00)", "3rd hour (02:00-03:00)",
                            "4th hour (03:00-04:00)", "5th hour (04:00-05:00)", "6th hour (05:00-06:00)",
                            "7th hour (06:00-07:00)", "8th hour (07:00-08:00)", "9th hour (08:00-09:00)",
                            "10th hour (09:00-10:00)", "11th hour (10:00-11:00)", "12th hour (11:00-12:00)",
                            "13th hour (12:00-13:00)", "14th hour (13:00-14:00)", "15th hour (14:00-15:00)",
                            "16th hour (15:00-16:00)", "17th hour (16:00-17:00)", "18th hour (17:00-18:00)",
                            "19th hour (18:00-19:00)", "20th hour (19:00-20:00)", "21st hour (20:00-21:00)",
                            "22nd hour (21:00-22:00)", "23rd hour (22:00-23:00)", "24th hour (23:00-00:00)"],
                    index=8,
                    help="Select the hour when you studied"
                )
            with col3:
                mood = st.select_slider("Today's study mood:", 
                                       options=["😫", "😕", "😐", "🙂", "😊", "🤩"],
                                       value="😊")
            
            topics_covered = st.text_input("Topics covered today:", 
                                          placeholder=f"e.g., {subject} concepts, exercises...")
            achievements = st.text_area("Key achievements:", 
                                       placeholder="What did you learn or accomplish today?")
            challenges = st.text_area("Challenges faced:", 
                                     placeholder="What was difficult? How did you overcome it?")
            
            if st.form_submit_button("💾 Save Today's Progress", use_container_width=True):
                daily_log = {
                    'date': today,
                    'hours': hours_studied,
                    'study_hour': study_hour,
                    'mood': mood,
                    'topics': topics_covered,
                    'achievements': achievements,
                    'challenges': challenges
                }
                
                if 'daily_progress' not in progress_data:
                    progress_data['daily_progress'] = []
                
                # Update or add entry
                existing_entry = False
                for i, entry in enumerate(progress_data['daily_progress']):
                    if entry.get('date') == today:
                        progress_data['daily_progress'][i] = daily_log
                        existing_entry = True
                        break
                
                if not existing_entry:
                    progress_data['daily_progress'].append(daily_log)
                
                progress_data['completed_tasks'] = completed_tasks
                progress_data['completed_weeks'] = [w['week'] for w in plan.get('weekly_schedule', []) if w.get('completed', False)]
                progress_data['weekly_streak'] = current_streak
                progress_data['longest_weekly_streak'] = max(progress_data.get('longest_weekly_streak', 0), longest_streak)
                progress_data['last_updated'] = datetime.now().isoformat()
                
                save_progress(subject, progress_data)
                st.success("✅ Progress saved successfully!")
                st.rerun()
        
        # Progress History
        if progress_data.get('daily_progress'):
            st.subheader("📅 Progress History")
            try:
                df = pd.DataFrame(progress_data['daily_progress'])
                if not df.empty:
                    display_df = df.copy()
                    
                    if 'date' in display_df.columns:
                        try:
                            display_df['Date'] = pd.to_datetime(display_df['date']).dt.strftime('%Y-%m-%d')
                        except:
                            display_df['Date'] = display_df['date']
                    
                    display_cols = []
                    if 'Date' in display_df.columns:
                        display_cols.append('Date')
                    if 'hours' in display_df.columns:
                        display_df['Hours'] = display_df['hours']
                        display_cols.append('Hours')
                    if 'study_hour' in display_df.columns:
                        display_df['Study Hour'] = display_df['study_hour']
                        display_cols.append('Study Hour')
                    if 'mood' in display_df.columns:
                        display_df['Mood'] = display_df['mood']
                        display_cols.append('Mood')
                    if 'topics' in display_df.columns:
                        display_df['Topics'] = display_df['topics']
                        display_cols.append('Topics')
                    if 'achievements' in display_df.columns:
                        display_df['Achievements'] = display_df['achievements']
                        display_cols.append('Achievements')
                    if 'challenges' in display_df.columns:
                        display_df['Challenges'] = display_df['challenges']
                        display_cols.append('Challenges')
                    
                    if display_cols:
                        st.dataframe(display_df[display_cols], use_container_width=True)
                    else:
                        st.warning("No valid data available to display.")
            except Exception as e:
                st.warning(f"Could not display progress history: {str(e)}")
    
    else:
        st.info("👈 **Generate a study plan first to track progress!**")

# ============================================
# TAB 4: ANALYTICS
# ============================================

with tab4:
    if st.session_state.study_plan:
        st.markdown('<h2 class="sub-header">📈 Study Analytics</h2>', unsafe_allow_html=True)
        
        subject = st.session_state.subject
        progress_data = load_progress(subject)
        plan = st.session_state.study_plan
        study_days = plan.get('study_days', ["Mon", "Tue", "Wed", "Thu", "Fri"])
        
        daily_progress = progress_data.get('daily_progress', [])
        study_days_logged = len(daily_progress)
        
        # Calculate completed tasks
        completed_tasks = 0
        total_tasks = 0
        
        for week in plan.get('weekly_schedule', []):
            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                for day_idx, day_tasks in enumerate(week['daily_tasks']):
                    if isinstance(day_tasks, list):
                        total_tasks += len(day_tasks)
                        for task_idx in range(len(day_tasks)):
                            task_key = f"task_week{week['week']}_day{day_idx}_task{task_idx}_{subject}"
                            if st.session_state.get(task_key, False):
                                completed_tasks += 1
        
        completed_weeks = sum(1 for week in plan.get('weekly_schedule', []) if week.get('completed', False))
        total_weeks = plan.get('total_weeks', 0)
        
        current_streak = progress_data.get('weekly_streak', 0)
        longest_streak = progress_data.get('longest_weekly_streak', 0)
        
        if daily_progress:
            try:
                col1, col2, col3, col4 = st.columns(4)
                
                with col1:
                    if daily_progress:
                        df = pd.DataFrame(daily_progress)
                        if 'hours' in df.columns:
                            df['hours'] = pd.to_numeric(df['hours'], errors='coerce')
                            total_hours = df['hours'].sum()
                        else:
                            total_hours = 0
                    else:
                        total_hours = 0
                    st.metric("Total Hours", f"{total_hours:.1f}")
                
                with col2:
                    st.metric("Study Days Logged", study_days_logged)
                
                with col3:
                    st.metric("Tasks Completed", f"{completed_tasks}/{total_tasks}")
                
                with col4:
                    st.metric("Weekly Streak", f"{current_streak} weeks")
                
                st.subheader("📈 Study Patterns Analysis")
                
                if daily_progress and len(daily_progress) > 1:
                    df = pd.DataFrame(daily_progress)
                    df['date'] = pd.to_datetime(df['date'], errors='coerce')
                    df = df.dropna(subset=['date'])
                    
                    if len(df) > 1:
                        if 'study_hour' in df.columns:
                            st.subheader("🕐 Preferred Study Hours")
                            
                            def extract_hour(hour_str):
                                if 'hour' in hour_str:
                                    try:
                                        hour_num = int(''.join(filter(str.isdigit, hour_str.split('hour')[0].strip())))
                                        return hour_num
                                    except:
                                        return 0
                                return 0
                            
                            df['hour_num'] = df['study_hour'].apply(extract_hour)
                            
                            fig_hour = go.Figure()
                            fig_hour.add_trace(go.Scatter(
                                x=df['date'],
                                y=df['hour_num'],
                                mode='lines+markers',
                                name='Study Hour',
                                line=dict(color='#3B82F6', width=2),
                                marker=dict(size=8),
                                text=df['study_hour'],
                                hoverinfo='text+y'
                            ))
                            
                            fig_hour.update_layout(
                                title='Study Hours Over Time',
                                xaxis_title='Date',
                                yaxis_title='Study Hour (1-24)',
                                yaxis=dict(tickmode='linear', dtick=1),
                                hovermode='x unified',
                                showlegend=True
                            )
                            
                            st.plotly_chart(fig_hour, use_container_width=True)
                            
                            if not df.empty and 'study_hour' in df.columns:
                                common_hour = df['study_hour'].mode()
                                if not common_hour.empty:
                                    st.info(f"**Most frequent study time:** {common_hour.iloc[0]}")
                        
                        if 'mood' in df.columns:
                            st.subheader("😊 Study Mood Analysis")
                            
                            mood_counts = df['mood'].value_counts().reset_index()
                            mood_counts.columns = ['Mood', 'Count']
                            
                            fig_mood = px.pie(mood_counts, values='Count', names='Mood',
                                            title='Mood Distribution During Study Sessions',
                                            hole=0.3)
                            
                            st.plotly_chart(fig_mood, use_container_width=True)
                
                st.subheader("💡 Personalized Recommendations")
                
                recommendations = []
                
                if daily_progress:
                    pattern_recommendations = analyze_study_patterns(daily_progress, study_days)
                    recommendations.extend(pattern_recommendations)
                
                progress_percentage = (completed_tasks / total_tasks * 100) if total_tasks > 0 else 0
                
                if progress_percentage < 30:
                    recommendations.append("**🌱 You're just getting started!** Focus on building consistent weekly study habits.")
                elif progress_percentage < 70:
                    recommendations.append("**📚 Great momentum!** You're making solid progress. Keep up the weekly consistency!")
                else:
                    recommendations.append("**🏆 Excellent progress!** You're approaching completion. Stay focused on your final goals!")
                
                if current_streak > 0:
                    if current_streak >= 5:
                        recommendations.append(f"🔥 **Incredible {current_streak}-week streak!** You've mastered consistent study habits!")
                    elif current_streak >= 3:
                        recommendations.append(f"🎯 **Great {current_streak}-week streak!** You're building powerful momentum!")
                    else:
                        recommendations.append(f"📈 **Good start with {current_streak} week streak!** Keep going to build a longer streak!")
                else:
                    recommendations.append("⏳ **Start a streak!** Complete a full week to begin your weekly streak journey.")
                
                if recommendations:
                    for rec in recommendations:
                        st.info(rec)
                else:
                    st.success("**Excellent!** Your study habits are on track. Keep up the good work!")
                    
                st.subheader("✅ Success Metrics Check")
                
                success_data = []
                success_data.append({
                    'Metric': 'Task completion progress',
                    'Status': '✅ Excellent' if progress_percentage >= 80 else '🟡 Good' if progress_percentage >= 50 else '⏳ Needs work',
                    'Value': f'{progress_percentage:.0f}%'
                })
                
                week_completion_rate = (completed_weeks / total_weeks * 100) if total_weeks > 0 else 0
                success_data.append({
                    'Metric': 'Weekly completion rate',
                    'Status': '✅ Excellent' if week_completion_rate >= 80 else '🟡 Good' if week_completion_rate >= 50 else '⏳ Needs work',
                    'Value': f'{week_completion_rate:.0f}%'
                })
                
                df_success = pd.DataFrame(success_data)
                st.dataframe(df_success, use_container_width=True)
            
            except Exception as e:
                st.warning(f"Could not generate analytics: {str(e)}")
        else:
            st.info("📊 **Log your daily progress in the 'Track Progress' tab to unlock analytics!**")
    
    else:
        st.info("👈 **Generate a study plan first to see analytics!**")

# ============================================
# TAB 5: TESTS - FIXED DUPLICATE BUTTON ERROR
# ============================================

with tab5:
    if st.session_state.study_plan:
        plan = st.session_state.study_plan
        subject = st.session_state.subject
        
        st.markdown(f'<h2 class="sub-header">🧪 Weekly Tests: {subject}</h2>', unsafe_allow_html=True)
        
        # Test Selection View
        if not st.session_state.current_test:
            st.info("Take weekly tests to assess your understanding.")
            
            if plan.get('weekly_schedule'):
                for week in plan['weekly_schedule']:
                    week_num = week['week']
                    
                    # Calculate completion
                    week_tasks = week['total_tasks']
                    completed_week_tasks = 0
                    
                    if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                        for day_idx, day_tasks in enumerate(week['daily_tasks']):
                            if isinstance(day_tasks, list):
                                for task_idx in range(len(day_tasks)):
                                    task_key = f"task_week{week_num}_day{day_idx}_task{task_idx}_{subject}"
                                    if st.session_state.get(task_key, False):
                                        completed_week_tasks += 1
                    
                    completion_percentage = (completed_week_tasks / week_tasks * 100) if week_tasks > 0 else 0
                    test_available = completion_percentage >= 50
                    
                    with st.container():
                        col1, col2, col3 = st.columns([3, 1, 1])
                        
                        with col1:
                            st.markdown(f"**Week {week_num}:** {week.get('focus', '')[:50]}...")
                            st.caption(f"Completion: {completion_percentage:.0f}% • Tasks: {completed_week_tasks}/{week_tasks}")
                        
                        with col2:
                            test_key = f"week_{week_num}_test"
                            test_taken = test_key in st.session_state.test_results
                            
                            if test_taken:
                                score = st.session_state.test_results[test_key]['percentage']
                                st.metric("", f"{score:.0f}%", label_visibility="collapsed")
                        
                        with col3:
                            if test_available:
                                if st.button(f"📝 Start Test", 
                                           key=f"start_test_week_{week_num}_{subject.replace(' ', '_')}",
                                           use_container_width=True):
                                    with st.spinner(f"Generating test for Week {week_num}..."):
                                        test_data = generate_weekly_test(
                                            st.session_state.api_key, 
                                            week_num, 
                                            week.get('daily_tasks', []), 
                                            subject
                                        )
                                        
                                        if test_data and 'questions' in test_data:
                                            st.session_state.current_test = week_num
                                            st.session_state.test_questions = test_data['questions']
                                            st.session_state.test_completed = False
                                            st.session_state.user_test_answers = {}
                                            st.rerun()
                                        else:
                                            st.error("Failed to generate test. Try again.")
                            else:
                                # FIXED: Added unique key and disabled properly
                                st.button(f"🔒 Locked", 
                                         key=f"locked_test_week_{week_num}_{subject.replace(' ', '_')}", 
                                         disabled=True, 
                                         use_container_width=True)
                                if completion_percentage < 50:
                                    st.caption(f"Complete 50% to unlock")
                        
                        st.markdown("---")
        
        # Test Taking View
        elif st.session_state.current_test and not st.session_state.test_completed:
            week_num = st.session_state.current_test
            test_questions = st.session_state.test_questions
            
            st.markdown(f"### 📝 Week {week_num} Test")
            st.info(f"Answer all {len(test_questions)} questions. Click on your chosen answer.")
            
            for i, question in enumerate(test_questions):
                with st.container():
                    st.markdown(f"**Q{i+1}. {question['question']}**")
                    
                    current_answer = st.session_state.user_test_answers.get(str(i), "")
                    
                    # Create 2 columns for options
                    col_a, col_b = st.columns(2)
                    
                    with col_a:
                        # Option A and B
                        for option in ['A', 'B']:
                            option_text = question['options'][option]
                            is_selected = (current_answer == option)
                            
                            if st.button(
                                f"{option}. {option_text[:80]}..." if len(option_text) > 80 else f"{option}. {option_text}",
                                key=f"test_q{week_num}_{i}_opt{option}",
                                use_container_width=True,
                                type="primary" if is_selected else "secondary"
                            ):
                                st.session_state.user_test_answers[str(i)] = option
                                st.rerun()
                    
                    with col_b:
                        # Option C and D
                        for option in ['C', 'D']:
                            option_text = question['options'][option]
                            is_selected = (current_answer == option)
                            
                            if st.button(
                                f"{option}. {option_text[:80]}..." if len(option_text) > 80 else f"{option}. {option_text}",
                                key=f"test_q{week_num}_{i}_opt{option}",
                                use_container_width=True,
                                type="primary" if is_selected else "secondary"
                            ):
                                st.session_state.user_test_answers[str(i)] = option
                                st.rerun()
                    
                    st.markdown("---")
            
            # Submit button
            col_submit1, col_submit2, col_submit3 = st.columns([1, 2, 1])
            with col_submit2:
                if st.button("📤 Submit Test", 
                           type="primary", 
                           key=f"submit_test_week_{week_num}",
                           use_container_width=True):
                    st.session_state.test_completed = True
                    st.rerun()
        
        # Test Results View
        elif st.session_state.test_completed:
            week_num = st.session_state.current_test
            test_questions = st.session_state.test_questions
            user_answers = st.session_state.user_test_answers
            
            # Calculate results
            correct_count = 0
            for i, question in enumerate(test_questions):
                if user_answers.get(str(i)) == question['correct_answer']:
                    correct_count += 1
            
            score_percentage = (correct_count / len(test_questions)) * 100
            
            # Store results
            test_key = f"week_{week_num}_test"
            st.session_state.test_results[test_key] = {
                'score': correct_count,
                'total': len(test_questions),
                'percentage': score_percentage
            }
            
            # Display results
            st.success(f"## 🎯 Test Score: {score_percentage:.1f}%")
            st.info(f"**{correct_count} out of {len(test_questions)} correct**")
            
            # Performance feedback
            if score_percentage >= 80:
                st.balloons()
                st.success("**Excellent!** You have a strong understanding of this week's material.")
            elif score_percentage >= 60:
                st.success("**Good job!** You understand the main concepts.")
            else:
                st.warning("**Review needed.** Consider revisiting this week's material.")
            
            with st.expander("📋 Review Answers", expanded=False):
                for i, question in enumerate(test_questions):
                    user_answer = user_answers.get(str(i), "Not answered")
                    is_correct = user_answer == question['correct_answer']
                    
                    st.markdown(f"**Q{i+1}. {question['question']}**")
                    
                    # Show user's answer with color
                    if user_answer in question['options']:
                        user_answer_text = f"{user_answer}. {question['options'][user_answer]}"
                        if is_correct:
                            st.success(f"✅ **Your answer:** {user_answer_text}")
                        else:
                            st.error(f"❌ **Your answer:** {user_answer_text}")
                    else:
                        st.warning("⏸️ **Your answer:** Not answered")
                    
                    # Show correct answer
                    correct_answer_text = f"{question['correct_answer']}. {question['options'][question['correct_answer']]}"
                    st.info(f"📚 **Correct answer:** {correct_answer_text}")
                    
                    # Show explanation
                    if question.get('explanation'):
                        st.markdown(f"💡 **Explanation:** {question['explanation']}")
                    
                    st.markdown("---")
            
            # Back button
            if st.button("← Back to Test List", 
                        key=f"back_from_test_{week_num}",
                        use_container_width=True):
                st.session_state.current_test = None
                st.session_state.test_completed = False
                st.session_state.test_questions = []
                st.session_state.user_test_answers = {}
                st.rerun()
    
    else:
        st.info("👈 **Generate a study plan first!**")

# ============================================
# TAB 6: EXPORT - FIXED MISSING FLEXIBLE_COMPLETION
# ============================================

with tab6:
    st.markdown('<h2 class="sub-header">📤 Export Your Study Plan</h2>', unsafe_allow_html=True)
    
    if st.session_state.study_plan:
        plan = st.session_state.study_plan
        subject = st.session_state.subject
        
        st.info("💾 **Download your study plan to continue later**")
        st.write("Save your progress and come back anytime to continue where you left off!")
        
        export_plan = plan.copy()
        
        # Update completed tasks info
        if export_plan.get('weekly_schedule'):
            for week in export_plan['weekly_schedule']:
                week_num = week['week']
                completed_tasks_info = []
                
                if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                    for day_idx, day_tasks in enumerate(week['daily_tasks']):
                        if isinstance(day_tasks, list):
                            for task_idx, task in enumerate(day_tasks):
                                task_key = f"task_week{week_num}_day{day_idx}_task{task_idx}_{subject}"
                                if st.session_state.get(task_key, False):
                                    completion_date = st.session_state.task_completion_dates.get(task_key, datetime.now().strftime("%Y-%m-%d"))
                                    completed_tasks_info.append({
                                        'day_idx': day_idx,
                                        'task_idx': task_idx,
                                        'task': task,
                                        'completion_date': completion_date
                                    })
                
                week['completed_tasks_info'] = completed_tasks_info
                week['tasks_completed'] = len(completed_tasks_info)
                week['completed'] = (len(completed_tasks_info) == week['total_tasks'])
                if week['completed'] and not week.get('completion_date'):
                    dates = [t['completion_date'] for t in completed_tasks_info if t.get('completion_date')]
                    if dates:
                        week['completion_date'] = max(dates)
        
        export_plan['task_completion_dates'] = st.session_state.task_completion_dates
        export_plan['flexible_completion'] = st.session_state.flexible_completion
        export_plan['exported_at'] = datetime.now().isoformat()
        
        # Add daily progress
        progress_data = load_progress(subject)
        if progress_data and 'daily_progress' in progress_data:
            export_plan['daily_progress'] = progress_data['daily_progress']
        
        json_str = json.dumps(export_plan, indent=2, ensure_ascii=False)
        
        st.download_button(
            label="📄 Download Study Plan (JSON)",
            data=json_str,
            file_name=f"study_plan_{subject.replace(' ', '_')}.json",
            mime="application/json",
            use_container_width=True,
            help="Download your complete study plan with all your progress and streak information"
        )
        
        st.markdown("---")
        st.subheader("🏆 Completion Certificate")
        
        # Calculate completion metrics
        total_tasks = 0
        completed_tasks = 0
        
        for week in plan.get('weekly_schedule', []):
            if 'daily_tasks' in week and isinstance(week['daily_tasks'], list):
                for day_idx, day_tasks in enumerate(week['daily_tasks']):
                    if isinstance(day_tasks, list):
                        total_tasks += len(day_tasks)
                        for task_idx in range(len(day_tasks)):
                            task_key = f"task_week{week['week']}_day{day_idx}_task{task_idx}_{subject}"
                            if st.session_state.get(task_key, False):
                                completed_tasks += 1
        
        completed_weeks = sum(1 for week in plan.get('weekly_schedule', []) if week.get('completed', False))
        total_weeks = plan.get('total_weeks', 0)
        completion_percentage = (completed_tasks / total_tasks * 100) if total_tasks > 0 else 0
        
        # Show certificate preview
        certificate_html = create_completion_certificate(subject, completion_percentage, total_tasks, completed_tasks, total_weeks, completed_weeks)
        
        if certificate_html:
            with st.expander("🎖️ View Certificate Preview", expanded=True):
                st.components.v1.html(certificate_html, height=800, scrolling=True)
            
            # Show certificate download options only if 90%+ complete
            if completion_percentage >= 90:
                st.success("🎉 **Congratulations! You've unlocked the certificate download!**")
                
                col1, col2 = st.columns(2)
                with col1:
                    st.download_button(
                        label="🎖️ Download Certificate (HTML)",
                        data=certificate_html,
                        file_name=f"certificate_{subject.replace(' ', '_')}.html",
                        mime="text/html",
                        use_container_width=True,
                        help="Download your certificate as HTML file"
                    )
                with col2:
                    st.download_button(
                        label="📝 Download Certificate (TXT)",
                        data=f"""CERTIFICATE OF COMPLETION

Subject: {subject}
Completion: {completion_percentage:.1f}%
Tasks Mastered: {completed_tasks}/{total_tasks}
Weeks Completed: {completed_weeks}/{total_weeks}
Date: {datetime.now().strftime("%B %d, %Y")}

This certifies successful completion of the {subject} study program.
Generated by AI Study Planner Pro.""",
                        file_name=f"certificate_{subject.replace(' ', '_')}.txt",
                        mime="text/plain",
                        use_container_width=True,
                        help="Download certificate details as text file"
                    )
            else:
                remaining_percentage = 90 - completion_percentage
                st.info(f"📈 **Complete {remaining_percentage:.1f}% more to unlock certificate download!**")
                st.progress(completion_percentage / 100)
                st.caption(f"Current progress: {completion_percentage:.1f}% (Minimum 90% required for certificate download)")
    else:
        st.info("👈 **Generate or load a study plan first to export!**")