app.py

# app.py - Enhanced Test Generator for Hugging Face
# Aligned with Pakistani National Curriculum (2006)
# Developer: Najaf Ali Sharqi
import os
import gradio as gr
from groq import Groq
from datetime import datetime
from docx import Document
from docx.shared import Pt, RGBColor, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.ns import qn
from docx.oxml import OxmlElement
import re
import json

# ---------------------------
# Initialize GROQ client with error handling
# ---------------------------
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY environment variable not set")
client = Groq(api_key=GROQ_API_KEY)

# ---------------------------
# Configuration Constants
# ---------------------------
APP_VERSION = "2.6"
MODEL_NAME = "llama-3.3-70b-versatile"

# ---------------------------
# Grades and subjects aligned with Pakistan National Curriculum
# ---------------------------
grades = [f"Grade {i}" for i in range(1, 13)]

subjects_by_grade = {
    "Grade 1": ["Mathematics", "English", "Urdu", "Islamiyat", "General Knowledge"],
    "Grade 2": ["Mathematics", "English", "Urdu", "Islamiyat", "General Knowledge"],
    "Grade 3": ["Mathematics", "English", "Urdu", "Science", "Islamiyat", "General Knowledge"],
    "Grade 4": ["Mathematics", "English", "Urdu", "Science", "Islamiyat", "General Knowledge"],
    "Grade 5": ["Mathematics", "English", "Urdu", "Science", "Islamiyat", "Social Studies"],
    "Grade 6": ["Mathematics", "English", "Urdu", "Science", "Islamiyat", "Social Studies"],
    "Grade 7": ["Mathematics", "English", "Urdu", "Science", "Islamiyat", "Social Studies", "Computer"],
    "Grade 8": ["Mathematics", "English", "Urdu", "Science", "Islamiyat", "Social Studies", "Computer"],
    "Grade 9": ["Mathematics", "English", "Urdu", "Biology", "Chemistry", "Physics", "Islamiyat", "Pakistan Studies"],
    "Grade 10": ["Mathematics", "English", "Urdu", "Biology", "Chemistry", "Physics", "Islamiyat", "Pakistan Studies"],
    "Grade 11": ["Mathematics", "English", "Urdu", "Biology", "Chemistry", "Physics", "Computer", "Islamiyat", "Pakistan Studies"],
    "Grade 12": ["Mathematics", "English", "Urdu", "Biology", "Chemistry", "Physics", "Computer", "Islamiyat", "Pakistan Studies"]
}

# ---------------------------
# Chapter mapping (National Curriculum 2006)
# ---------------------------
chapters_by_subject_and_grade = {
    "Grade 1": {
        "Mathematics": ["Counting & Numbers", "Basic Addition", "Basic Subtraction", "Shapes", "Measurements"],
        "English": ["Alphabet & Sounds", "Basic Words", "Simple Sentences", "Listening & Speaking", "Reading Short Texts"],
        "Urdu": ["حروف", "الفاظ", "سادہ جملے", "پڑھنا", "لکھنا"],
        "Islamiyat": ["Basic Beliefs", "Prophets Stories", "Good Manners", "Prayers & Worship"],
        "General Knowledge": ["My Body", "My Home", "My School", "Seasons"]
    },
    "Grade 2": {
        "Mathematics": ["Place Value", "Addition & Subtraction", "Money", "Time", "Shapes"],
        "English": ["Parts of Speech Basics", "Simple Grammar", "Short Comprehension", "Vocabulary Building"],
        "Urdu": ["الفاظ و جملے", "مختصر کہانیاں", "حروفِ صحیح"],
        "Islamiyat": ["Prophets", "Good Deeds", "Mosque Etiquette"],
        "General Knowledge": ["Plants & Animals", "Local Community", "Transport"]
    },
    "Grade 3": {
        "Mathematics": ["Multiplication", "Division", "Fractions", "Geometry Basics", "Measurement"],
        "English": ["Tenses (simple)", "Comprehension Passages", "Vocabulary", "Sentence Formation"],
        "Urdu": ["کالم و مضامین", "قواعدِ اردو", "کہانیاں"],
        "Science": ["Plants", "Animals", "Human Body", "Materials"],
        "Islamiyat": ["Basics of Iman", "Seerat", "Manners"],
        "General Knowledge": ["Environment", "Community Helpers", "Safety"]
    },
    "Grade 4": {
        "Mathematics": ["Fractions & Decimals", "Geometry", "Data Handling", "Word Problems"],
        "English": ["Grammar (intermediate)", "Comprehension", "Writing Short Paragraphs"],
        "Urdu": ["نظم و نثر", "قواعد", "مفردات"],
        "Science": ["Ecosystems", "Forces", "Heat & Energy"],
        "Islamiyat": ["Pillars of Islam", "Seerah Stories"],
        "General Knowledge": ["Maps & Directions", "Cultural Heritage", "Technology"]
    },
    "Grade 5": {
        "Mathematics": ["Numbers & Operations", "Geometry", "Statistics", "Ratio & Proportion"],
        "English": ["Advanced Grammar Basics", "Paragraph Writing", "Comprehension"],
        "Urdu": ["ادبی مطالعہ", "تعبیر و تشریح"],
        "Science": ["Cells & Life", "Matter", "Energy"],
        "Islamiyat": ["Quranic Stories", "Akhlaq"],
        "Social Studies": ["Pakistan Geography", "History Basics", "Civics Introduction"]
    },
    "Grade 6": {
        "Mathematics": ["Number System", "Algebra Introduction", "Geometry", "Data Handling"],
        "English": ["Grammar (complex)", "Essay Writing", "Comprehension Skills"],
        "Urdu": ["ادب و نحو", "تحریر"],
        "Science": ["Atoms & Molecules", "Human Body", "Plants"],
        "Islamiyat": ["Seerah & Ahkam"],
        "Social Studies": ["World Geography", "Historical Events", "Government"]
    },
    "Grade 7": {
        "Mathematics": ["Algebra", "Geometry", "Mensuration", "Statistics"],
        "English": ["Poetry & Prose", "Grammar", "Comprehension"],
        "Urdu": ["نثر و شاعری", "قواعد"],
        "Science": ["Cell Biology", "Human Systems", "Electricity Basics"],
        "Islamiyat": ["Islamic History", "Moral Teachings"],
        "Social Studies": ["Geography Basics", "History", "Civics"],
        "Computer": ["Basic ICT", "Computer Hardware", "Introduction to Coding"]
    },
    "Grade 8": {
        "Mathematics": ["Linear Equations", "Geometry", "Pythagoras", "Graphs"],
        "English": ["Advanced Comprehension", "Composition", "Grammar"],
        "Urdu": ["تجزیہ ادب", "تحقیقِ مختصر"],
        "Science": ["Forces & Motion", "Heat", "Waves Basics"],
        "Islamiyat": ["Islamic Law Basics", "Seerah"],
        "Social Studies": ["Regional Geography", "Historical Events"],
        "Computer": ["Algorithms", "Basic Programming"]
    },
    "Grade 9": {
        "Mathematics": ["Number Systems & Algebra", "Coordinate Geometry", "Trigonometry", "Graphs"],
        "English": ["Advanced Grammar", "Comprehension", "Writing Skills"],
        "Urdu": ["ادب اور نثر", "نظم", "مضامین"],
        "Biology": [
            "Introduction to Biology", "Solving a Biological Problem", "Biodiversity",
            "Cells and Tissues", "Cell Cycle", "Enzymes", "Bioenergetics", "Nutrition", "Transport"
        ],
        "Chemistry": [
            "Fundamentals of Chemistry", "Structure of Atoms", "Periodic Table and Periodicity of Properties",
            "Structure of Molecules", "Physical States of Matter", "Solutions", "Electrochemistry", "Chemical Reactivity"
        ],
        "Physics": [
            "Physical Quantities and Measurement", "Kinematics", "Dynamics", "Turning Effect of Forces",
            "Gravitation", "Work and Energy", "Properties of Matter", "Thermal Properties of Matter", "Transfer of Heat"
        ],
        "Islamiyat": ["Seerat of Prophet", "Faith & Beliefs", "Worship Practices"],
        "Pakistan Studies": ["Ideological Basis of Pakistan", "Making of Pakistan", "Land and Environment", "History of Pakistan"]
    },
    "Grade 10": {
        "Mathematics": ["Advanced Algebra", "Trigonometry", "Geometry", "Probability"],
        "English": ["Literature", "Advanced Composition", "Comprehension"],
        "Urdu": ["ادبِ جدید", "تحقیقی مضمون"],
        "Biology": [
            "Gaseous Exchange", "Homeostasis", "Coordination", "Support and Movement",
            "Reproduction", "Inheritance", "Man and His Environment", "Biotechnology", "Pharmacology"
        ],
        "Chemistry": [
            "Chemical Equilibrium", "Acids, Bases and Salts", "Organic Chemistry", "Hydrocarbons",
            "Biochemistry", "Environmental Chemistry I: The Atmosphere", "Environmental Chemistry II: Water", "Chemical Industries"
        ],
        "Physics": [
            "Simple Harmonic Motion and Waves", "Sound", "Geometrical Optics", "Electrostatics",
            "Current Electricity", "Electromagnetism", "Introductory Electronics", 
            "Information and Communication Technology", "Radioactivity"
        ],
        "Islamiyat": ["Islamic Studies II", "Ethics & Society", "Fiqh Basics"],
        "Pakistan Studies": [
            "History of Pakistan-II", "Pakistan in World Affairs", "Economic Developments",
            "Population, Society and Culture of Pakistan"
        ]
    },
    "Grade 11": {
        "Mathematics": ["Advanced Algebra & Calculus Intro", "Coordinate Geometry", "Trigonometry", "Vectors"],
        "English": ["Academic Writing", "Literature Studies", "Critical Reading"],
        "Urdu": ["ادبی مطالعہ", "نظم و نثر کی تشریح"],
        "Biology": [
            "Cell Structure and Functions", "Biological Molecules", "Enzymes", "Bioenergetics",
            "Acellular Life", "Prokaryotes", "Protists and Fungi", "Diversity among Plants",
            "Diversity among Animals", "Form and Functions in Plants", "Digestion", "Circulation", "Immunity"
        ],
        "Chemistry": [
            "Stoichiometry", "Atomic Structure", "Theories of Covalent Bonding and Shapes of Molecules",
            "States of Matter I: Gases", "States of Matter II: Liquids", "States of Matter III: Solids",
            "Chemical Equilibrium", "Acids, Bases and Salts", "Chemical Kinetics",
            "Solutions and Colloids", "Thermochemistry", "Oxidation, Reduction and Electrochemistry"
        ],
        "Physics": [
            "Measurement", "Vectors and Equilibrium", "Forces and Motion", "Work and Energy",
            "Rotational and Circular Motion", "Fluid Dynamics", "Oscillations", "Waves",
            "Physical Optics", "Thermodynamics", "Electrostatics", "Current Electricity"
        ],
        "Computer": ["Programming Fundamentals", "Data Structures Intro", "Databases"],
        "Islamiyat": ["Advanced Seerah", "Comparative Religion"],
        "Pakistan Studies": [
            "Pakistan in Geographical Perspective", "Pakistan Resources", "Historical Perspective (Ancient to Modern)",
            "Islam and Pakistan", "Administrative System", "Political and Constitutional Developments", "Human Rights"
        ]
    },
    "Grade 12": {
        "Mathematics": ["Calculus", "Advanced Algebra", "Statistics", "Analytical Geometry"],
        "English": ["Research Writing", "Advanced Literature", "Critical Analysis"],
        "Urdu": ["تحقیق و تنقید", "ادبی شخصیات"],
        "Biology": [
            "Respiration", "Homeostasis", "Support and Movement", "Nervous Coordination",
            "Chemical Coordination", "Behavior", "Reproduction", "Development and Aging",
            "Inheritance", "Chromosome and DNA", "Evolution", "Man and His Environment",
            "Biotechnology", "Biology and Human Welfare"
        ],
        "Chemistry": [
            "s and p Block Elements", "d-Block Elements", "Organic Compounds", "Hydrocarbons",
            "Alkyl Halides and Amines", "Alcohols and Phenols", "Aldehydes and Ketones",
            "Carboxylic Acids and Functional Derivatives", "Biochemistry", "Industrial Chemistry",
            "Environmental Chemistry", "Analytical Chemistry"
        ],
        "Physics": [
            "Electromagnetism", "Electromagnetic Induction", "Alternating Current", "Physics of Solids",
            "Electronics", "Dawn of Modern Physics", "Atomic Spectra", "Nuclear Physics"
        ],
        "Computer": ["Web Technologies", "Advanced Programming Concepts"],
        "Islamiyat": ["Advanced Islamic Thought", "Islamic Jurisprudence"],
        "Pakistan Studies": [
            "Society and Culture", "Foreign Relations of Pakistan", "Economic Development",
            "Sports, Tourism and National Identity", "Contemporary Challenges and Policies"
        ]
    }
}

# ---------------------------
# Question types for comprehensive assessment
# ---------------------------
QUESTION_TYPES = {
    "Multiple Choice Questions (MCQs)": {
        "code": "mcq",
        "description": "Four-option questions testing recall and application",
        "format": "Question with options A, B, C, D"
    },
    "Fill in the Blanks": {
        "code": "fill_blank",
        "description": "Complete sentences with missing key terms",
        "format": "Sentences with underlined blanks"
    },
    "Match the Column": {
        "code": "match_column",
        "description": "Connect related items from two columns",
        "format": "Two columns with items to match"
    },
    "Short Response Questions": {
        "code": "short_response",
        "description": "Brief written answers (2-3 sentences)",
        "format": "Questions requiring concise explanations"
    },
    "Essay Type Questions": {
        "code": "essay",
        "description": "Extended written responses with detailed analysis",
        "format": "Open-ended questions requiring paragraphs"
    }
}

# ---------------------------
# Cognitive levels for item writing (Bloom's Taxonomy)
# ---------------------------
DIFFICULTY_LEVELS = {
    "Remembering": {
        "description": "Recall facts, terms, basic concepts",
        "keywords": ["define", "list", "identify", "name", "state", "recall"]
    },
    "Understanding": {
        "description": "Explain ideas or concepts",
        "keywords": ["explain", "describe", "summarize", "interpret", "paraphrase"]
    },
    "Applying": {
        "description": "Use information in new situations",
        "keywords": ["apply", "demonstrate", "solve", "use", "calculate", "implement"]
    },
    "Analyzing": {
        "description": "Draw connections, examine relationships",
        "keywords": ["analyze", "compare", "contrast", "differentiate", "examine"]
    },
    "Evaluating": {
        "description": "Justify decisions, make judgments",
        "keywords": ["evaluate", "justify", "critique", "assess", "judge"]
    },
    "Creating": {
        "description": "Produce new or original work",
        "keywords": ["create", "design", "construct", "develop", "formulate"]
    }
}

# ---------------------------
# Enhanced question generation with Bloom's Taxonomy distribution
# ---------------------------
def generate_questions_multi_level(grade, subject, chapter, test_type, school_name, question_type, 
                                   bloom_distribution, total_questions):
    """Generate questions with mixed cognitive levels based on user distribution"""
    
    all_questions = []
    all_answers = []
    
    try:
        # Calculate number of questions per level
        for level, percentage in bloom_distribution.items():
            if percentage > 0:
                num_q = int((percentage / 100) * total_questions)
                if num_q > 0:
                    questions, answers, error, q_type_code = generate_questions(
                        grade, subject, chapter, level,
                        num_q, test_type, school_name, question_type
                    )
                    
                    if error:
                        return [], [], error, "mcq"
                    
                    all_questions.extend(questions)
                    all_answers.extend(answers)
        
        # Ensure we have the right number of questions
        if len(all_questions) < total_questions:
            # Generate additional questions at first non-zero level
            shortage = total_questions - len(all_questions)
            first_level = next((k for k, v in bloom_distribution.items() if v > 0), "Understanding")
            
            extra_q, extra_a, error, q_type_code = generate_questions(
                grade, subject, chapter, first_level,
                shortage, test_type, school_name, question_type
            )
            
            if not error:
                all_questions.extend(extra_q)
                all_answers.extend(extra_a)
        
        return all_questions[:total_questions], all_answers[:total_questions], None, q_type_code
        
    except Exception as e:
        return [], [], f"Error in multi-level generation: {str(e)}", "mcq"

# ---------------------------
# Enhanced question generation with proper item writing principles
# ---------------------------
def generate_questions(grade, subject, chapter, difficulty, num_questions, test_type, school_name, question_type):
    """Generate questions using AI with enhanced item writing guidelines"""
    
    try:
        q_type_code = QUESTION_TYPES.get(question_type, {}).get("code", "mcq")
        is_mathematics = subject.lower() == "mathematics"
        
        # Extract grade number for level-specific prompts
        grade_num = int(grade.split()[-1]) if grade.startswith("Grade") else 6
        
        # Build type-specific prompts
        if q_type_code == "mcq":
            system_msg = """You are an expert educational item writer specializing in Multiple Choice Questions aligned with the Pakistani National Curriculum (2006).
CRITICAL FORMATTING RULES:
- Start DIRECTLY with question 1
- DO NOT write introductory text like "Here are the questions" or similar
- DO NOT number your output with "Q1." - the system will add numbering
- Each question must be on its own, clean format
Professional MCQ Principles:
1. STEM CLARITY: Write clear, concise question stems
2. PLAUSIBLE DISTRACTORS: Wrong options based on common misconceptions
3. GRAMMATICAL CONSISTENCY: All options match the stem
4. NO CUEING: Avoid "all/none of the above"
5. EQUAL LENGTH: Similar option lengths
6. AGE-APPROPRIATE: Vocabulary for grade level
Format EXACTLY as:
[Question stem]
A) [Option 1]
B) [Option 2]
C) [Option 3]
D) [Option 4]
[Next question stem]
A) [Option 1]
B) [Option 2]
C) [Option 3]
D) [Option 4]
After all questions, provide:
ANSWER KEY:
1. B
2. A"""

            user_msg = f"""Create {num_questions} MCQs for {grade} {subject}, Chapter: {chapter}
Cognitive Level: {difficulty}
START DIRECTLY WITH THE FIRST QUESTION - NO INTRODUCTION TEXT.
Align with Pakistani National Curriculum (2006).
Use {difficulty} skills: {DIFFICULTY_LEVELS.get(difficulty, {}).get('description', '')}"""

        elif q_type_code == "fill_blank":
            system_msg = """You are an expert at creating Fill in the Blanks questions for Pakistani curriculum.
CRITICAL: Start DIRECTLY with question 1. NO introductory text.
Format EXACTLY as:
[Sentence with __________ for the blank]
[Next sentence with __________ for the blank]
After all questions:
ANSWERS:
1. [correct word/phrase]
2. [correct word/phrase]"""

            user_msg = f"""Create {num_questions} Fill in the Blanks for {grade} {subject}, Chapter: {chapter}
Cognitive Level: {difficulty}
Rules:
- One blank per sentence marked with __________
- Blank should test key concepts
- Provide clear context clues
- NO introduction - start with question 1 directly"""

        elif q_type_code == "match_column":
            system_msg = """You are an expert at creating Match the Column questions for Pakistani curriculum.
CRITICAL: Start DIRECTLY with the matching sets. NO introductory text.
Format EXACTLY as:
Column A | Column B
1. [Item A1] | a. [Item B1]
2. [Item A2] | b. [Item B2]
3. [Item A3] | c. [Item B3]
ANSWERS:
1-[letter], 2-[letter], 3-[letter]"""

            user_msg = f"""Create {num_questions//2} matching pairs for {grade} {subject}, Chapter: {chapter}
Cognitive Level: {difficulty}
Create two columns with related items (terms-definitions, concepts-examples, etc.)
NO introduction - start directly with the table"""

        elif q_type_code == "short_response":
            if is_mathematics:
                # MATHEMATICS-SPECIFIC SHORT QUESTIONS
                system_msg = f"""You are an expert Mathematics teacher creating SHORT NUMERICAL/COMPUTATIONAL questions for {grade} aligned with Pakistani National Curriculum.

CRITICAL RULES FOR MATHEMATICS SHORT QUESTIONS:
1. Create ACTUAL COMPUTATIONAL PROBLEMS - NOT descriptive/theoretical questions
2. Each question must require CALCULATIONS, SOLVING, or MATHEMATICAL WORKING
3. Questions must be SOLVABLE with specific numerical answers
4. NO questions like "Describe...", "Explain the significance...", "What is the importance..."
5. START DIRECTLY - NO introductory text

QUESTION TYPES TO USE (based on grade level):
"""
                # Add grade-specific guidance
                if grade_num <= 5:
                    system_msg += """
- Simple arithmetic calculations (addition, subtraction, multiplication, division)
- Word problems with numbers
- Basic fraction operations
- Simple geometry (find perimeter, area)
- Money calculations
"""
                elif grade_num <= 8:
                    system_msg += """
- Algebraic expressions and simplification
- Solving linear equations (1-2 steps)
- Geometry problems (angles, areas, perimeters)
- Ratio and proportion calculations
- Percentage problems
- Statistics (mean, median, mode)
"""
                else:  # Grade 9-12
                    system_msg += """
- Solving algebraic equations and inequalities
- Trigonometric calculations (find angles, sides)
- Coordinate geometry problems
- Calculus problems (differentiation, integration for Grade 11-12)
- Matrix operations
- Sequence and series
- Probability calculations
- Vector operations (for higher grades)
"""
                
                system_msg += """
FORMAT EXACTLY AS:
Solve: 3x + 5 = 20
Find the area of a triangle with base 8 cm and height 6 cm.
Calculate: (2/3) + (5/6)

After all questions provide:
SOLUTIONS WITH WORKING:
1. 3x + 5 = 20
   3x = 20 - 5
   3x = 15
   x = 5

2. Area = (1/2) × base × height
   Area = (1/2) × 8 × 6
   Area = 24 cm²
"""
                
                user_msg = f"""Create {num_questions} SHORT MATHEMATICAL PROBLEMS for {grade}, Chapter: {chapter}
Cognitive Level: {difficulty}

CRITICAL REQUIREMENTS:
- Each question must be a SOLVABLE mathematical problem
- Include specific numbers and require calculations
- Mix of: equations, word problems, geometry, calculations
- Difficulty appropriate for {grade}
- Related to chapter: {chapter}
- NO descriptive/theoretical questions

START DIRECTLY WITH FIRST PROBLEM - NO INTRODUCTION"""

            else:
                # Non-mathematics subjects (original behavior)
                system_msg = """You are an expert at creating Short Response questions for Pakistani curriculum.
CRITICAL: Start DIRECTLY with question 1. NO introductory text.
Format EXACTLY as:
[Question requiring 2-3 sentence answer]
[Next question]
After all questions:
MARKING SCHEME:
1. Key points: [point 1], [point 2], [point 3]
2. Key points: [point 1], [point 2]"""

                user_msg = f"""Create {num_questions} Short Response questions for {grade} {subject}, Chapter: {chapter}
Cognitive Level: {difficulty}
Each answer should be 2-3 sentences (3-5 marks each)
NO introduction - start with question 1 directly"""

        elif q_type_code == "essay":
            if is_mathematics:
                # MATHEMATICS-SPECIFIC LONG QUESTIONS
                system_msg = f"""You are an expert Mathematics teacher creating LONG NUMERICAL/PROOF questions for {grade} aligned with Pakistani National Curriculum.

CRITICAL RULES FOR MATHEMATICS LONG QUESTIONS:
1. Create MULTI-STEP MATHEMATICAL PROBLEMS requiring detailed working
2. Each question must involve CALCULATIONS, PROOFS, DERIVATIONS, or COMPLEX PROBLEM-SOLVING
3. Questions must require showing complete mathematical working
4. NO descriptive essays like "Write an essay on...", "Discuss the history of..."
5. START DIRECTLY - NO introductory text

QUESTION TYPES TO USE (based on grade level):
"""
                # Add grade-specific guidance
                if grade_num <= 5:
                    system_msg += """
- Multi-step word problems
- Combined operations problems
- Detailed geometry problems (multiple shapes)
- Problems requiring multiple calculations
"""
                elif grade_num <= 8:
                    system_msg += """
- Solve and verify equations
- Word problems with multiple steps
- Geometric proofs (basic)
- Complete problems with parts (a), (b), (c)
- Data handling with multiple calculations
- Graph drawing and interpretation
"""
                else:  # Grade 9-12
                    system_msg += """
- Multi-part algebraic problems with proof
- Trigonometric identities to prove
- Calculus problems (derivatives, integrals for Grade 11-12)
- Coordinate geometry with complete derivation
- Vector problems with multiple parts
- Matrix operations and determinants
- Complex word problems requiring modeling
- Geometric constructions and proofs
- Sequence and series with proofs
- Probability with multiple scenarios
"""
                
                system_msg += """
FORMAT FOR MULTI-PART QUESTIONS:
Question 1:
(a) Solve the equation: 2x² - 5x - 3 = 0
(b) Verify your solutions by substituting back
(c) Graph the quadratic function

Question 2:
Prove that: sin²θ + cos²θ = 1
[Show complete algebraic proof]

After all questions provide:
DETAILED SOLUTIONS:
1. (a) Using quadratic formula:
   x = [-b ± √(b² - 4ac)] / 2a
   [complete working shown step by step]
   
   (b) Verification:
   [substitution and checking]
   
   (c) [Graph description with key points]

MARKING SCHEME:
- Method: X marks
- Calculations: Y marks  
- Final answer: Z marks
"""
                
                user_msg = f"""Create {num_questions} LONG MATHEMATICAL PROBLEMS for {grade}, Chapter: {chapter}
Cognitive Level: {difficulty}

CRITICAL REQUIREMENTS:
- Each question must be a COMPLEX SOLVABLE mathematical problem
- Include multi-step calculations, proofs, or derivations
- Can have multiple parts (a), (b), (c)
- Appropriate difficulty for {grade}
- Related to chapter: {chapter}
- Show complete working in solutions
- NO essay-type descriptive questions

Types to include:
- Numerical problems with multiple steps
- Proofs of identities/theorems (for higher grades)
- Word problems requiring mathematical modeling
- Problems combining multiple concepts

START DIRECTLY WITH FIRST PROBLEM - NO INTRODUCTION"""

            else:
                # Non-mathematics subjects (original behavior)
                system_msg = """You are an expert at creating Essay questions for Pakistani curriculum.
CRITICAL: Start DIRECTLY with question 1. NO introductory text.
Format EXACTLY as:
[Essay question requiring detailed analysis]
[Next essay question]
After all questions:
MARKING CRITERIA:
1. Expected coverage: [main points to cover]
2. Expected coverage: [main points to cover]"""

                user_msg = f"""Create {num_questions} Essay questions for {grade} {subject}, Chapter: {chapter}
Cognitive Level: {difficulty}
Each answer should be 1-2 paragraphs (8-10 marks each)
Require analysis, evaluation, or synthesis
NO introduction - start with question 1 directly"""

        # Call Groq API
        response = client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_msg},
                {"role": "user", "content": user_msg}
            ],
            model=MODEL_NAME,
            temperature=0.7,
            max_tokens=4000
        )

        text = response.choices[0].message.content
        
        # Clean up any introductory text
        text = clean_ai_response(text)
        
        # Parse based on question type
        if q_type_code == "mcq":
            questions, answers = parse_mcqs_and_answers(text, num_questions)
        elif q_type_code == "fill_blank":
            questions, answers = parse_fill_blanks(text, num_questions)
        elif q_type_code == "match_column":
            questions, answers = parse_match_column(text)
        elif q_type_code == "short_response":
            questions, answers = parse_short_response(text, num_questions)
        elif q_type_code == "essay":
            questions, answers = parse_essay(text, num_questions)
        else:
            questions, answers = [], []
        
        return questions, answers, None, q_type_code
        
    except Exception as e:
        error_msg = f"Error generating questions: {str(e)}"
        return [], [], error_msg, "mcq"

def clean_ai_response(text):
    """Remove common AI introductory phrases"""
    # Patterns to remove
    patterns = [
        r'^Here\s+(are|is)\s+the\s+\d+\s+.*?:?\s*\n*',
        r'^Below\s+(are|is)\s+.*?:?\s*\n*',
        r'^I\'ve\s+created\s+.*?:?\s*\n*',
        r'^The\s+following\s+(are|is)\s+.*?:?\s*\n*',
        r'^Q1\.\s+Here.*?:\s*\n*',
    ]
    
    for pattern in patterns:
        text = re.sub(pattern, '', text, flags=re.IGNORECASE | re.MULTILINE)
    
    return text.strip()

def parse_mcqs_and_answers(text, expected_count):
    """Parse AI response to extract MCQs and answer key"""
    
    # Split by answer key section
    parts = re.split(r'ANSWER\s+KEY:?', text, flags=re.IGNORECASE)
    
    mcqs = []
    answer_key = []
    
    if len(parts) >= 1:
        # Extract MCQs
        mcq_text = parts[0]
        
        # Remove any "Q#." prefixes and split by question patterns
        mcq_text = re.sub(r'^\s*Q\d+\.\s*', '', mcq_text, flags=re.MULTILINE)
        
        # Split by numbered questions (1., 2., etc.) but keep the content
        questions = re.split(r'\n\s*\d+\.\s+', mcq_text)
        questions = [q.strip() for q in questions if q.strip() and len(q.strip()) > 20]
        mcqs = questions[:expected_count]
    
    if len(parts) >= 2:
        # Extract answer key
        answer_text = parts[1]
        # Find patterns like "1. B" or "1) B" or "1: B"
        answers = re.findall(r'\d+[\.\):\s]+([A-D])', answer_text, re.IGNORECASE)
        answer_key = [ans.upper() for ans in answers[:expected_count]]
    
    # If parsing failed, create placeholder
    if len(answer_key) < len(mcqs):
        answer_key.extend(['A'] * (len(mcqs) - len(answer_key)))
    
    return mcqs, answer_key

def parse_fill_blanks(text, expected_count):
    """Parse Fill in the Blanks questions"""
    parts = re.split(r'ANSWERS?:?', text, flags=re.IGNORECASE)
    
    questions = []
    answers = []
    
    if len(parts) >= 1:
        q_text = parts[0]
        # Split by line breaks, filter empty lines
        lines = [line.strip() for line in q_text.split('\n') if line.strip() and '_' in line]
        questions = lines[:expected_count]
    
    if len(parts) >= 2:
        a_text = parts[1]
        # Extract answers
        ans_lines = re.findall(r'\d+[\.\):\s]+(.+)', a_text)
        answers = [ans.strip() for ans in ans_lines[:expected_count]]
    
    if len(answers) < len(questions):
        answers.extend(['_____'] * (len(questions) - len(answers)))
    
    return questions, answers

def parse_match_column(text):
    """Parse Match the Column questions"""
    parts = re.split(r'ANSWERS?:?', text, flags=re.IGNORECASE)
    
    questions = []
    answers = []
    
    if len(parts) >= 1:
        # Extract the table
        questions.append(parts[0].strip())
    
    if len(parts) >= 2:
        # Extract matching answers
        answers.append(parts[1].strip())
    
    return questions, answers

def parse_short_response(text, expected_count):
    """Parse Short Response questions - separating questions from expected answers"""
    # Split by expected answers section
    parts = re.split(r'EXPECTED\s+ANSWERS?:?|MARKING\s+SCHEME:?|SOLUTIONS?\s+WITH\s+WORKING:?', text, flags=re.IGNORECASE)
    
    questions = []
    answers = []
    
    if len(parts) >= 1:
        q_text = parts[0]
        # Split by numbered patterns
        qs = re.split(r'\n\s*\d+[\.\)]\s+', q_text)
        questions = [q.strip() for q in qs if q.strip() and len(q.strip()) > 10][:expected_count]
    
    if len(parts) >= 2:
        a_text = parts[1]
        # Extract expected answers/key points
        schemes = re.findall(r'\d+[\.\)]\s*(.+?)(?=\d+[\.\)]|$)', a_text, re.DOTALL)
        answers = [s.strip() for s in schemes[:expected_count]]
    
    if len(answers) < len(questions):
        answers.extend(['Key points: Provide comprehensive explanation with examples'] * (len(questions) - len(answers)))
    
    return questions, answers

def parse_essay(text, expected_count):
    """Parse Essay questions - separating questions from expected answers"""
    # Split by expected answers section
    parts = re.split(r'EXPECTED\s+ANSWERS?:?|MARKING\s+CRITERIA:?|DETAILED\s+SOLUTIONS?:?', text, flags=re.IGNORECASE)
    
    questions = []
    answers = []
    
    if len(parts) >= 1:
        q_text = parts[0]
        # Split by numbered patterns or "Question X:"
        qs = re.split(r'\n\s*(?:\d+[\.\)]|Question\s+\d+:)\s+', q_text, flags=re.IGNORECASE)
        questions = [q.strip() for q in qs if q.strip() and len(q.strip()) > 15][:expected_count]
    
    if len(parts) >= 2:
        a_text = parts[1]
        # Extract criteria/expected coverage
        criteria = re.findall(r'(?:\d+[\.\)]|Question\s+\d+:)\s*(.+?)(?=(?:\d+[\.\)]|Question\s+\d+:)|$)', a_text, re.DOTALL | re.IGNORECASE)
        answers = [c.strip() for c in criteria[:expected_count]]
    
    if len(answers) < len(questions):
        answers.extend(['Main points: Detailed analysis with evidence, examples, and logical argumentation'] * (len(questions) - len(answers)))
    
    return questions, answers

def create_answer_table(doc, answers, num_cols=5):
    """Helper function to create formatted answer table"""
    num_rows = len(answers) + 1
    ans_table = doc.add_table(rows=num_rows, cols=num_cols)
    ans_table.style = 'Medium Grid 1 Accent 1'
    
    # Set column widths
    for row in ans_table.rows:
        for idx, cell in enumerate(row.cells):
            if idx == 0:
                cell.width = Inches(0.6)
            else:
                cell.width = Inches(1.2)
    
    return ans_table

# ---------------------------
# Enhanced DOCX export supporting all question types
# ---------------------------
def export_to_word(questions, answers, grade, subject, chapter, difficulty, test_type, school_name, total_marks, question_type_code):
    """Export questions to professionally formatted Word document"""
    
    try:
        doc = Document()
        
        # Set margins
        sections = doc.sections
        for section in sections:
            section.top_margin = Inches(0.8)
            section.bottom_margin = Inches(0.8)
            section.left_margin = Inches(0.9)
            section.right_margin = Inches(0.9)
        
        # Header with school name
        if school_name and school_name.strip():
            school_heading = doc.add_heading(school_name.upper(), level=0)
            school_heading.alignment = WD_ALIGN_PARAGRAPH.CENTER
            for run in school_heading.runs:
                run.font.size = Pt(16)
                run.font.color.rgb = RGBColor(0, 51, 102)
        
        # Test title
        test_title = doc.add_heading(f"{test_type.upper()}", level=1)
        test_title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        for run in test_title.runs:
            run.bold = True
            run.font.color.rgb = RGBColor(0, 102, 51)
            run.font.size = Pt(16)
        
        # Information table
        table = doc.add_table(rows=7, cols=2)
        table.style = 'Light Grid Accent 1'
        
        # Get question type name
        q_type_name = next((k for k, v in QUESTION_TYPES.items() if v['code'] == question_type_code), "MCQs")
        
        info_data = [
            ("Grade:", grade),
            ("Subject:", subject),
            ("Chapter:", chapter),
            ("Question Type:", q_type_name),
            ("Cognitive Level:", difficulty),
            ("Total Marks:", str(total_marks)),
            ("Date:", datetime.now().strftime("%d-%m-%Y"))
        ]
        
        for idx, (label, value) in enumerate(info_data):
            row = table.rows[idx]
            row.cells[0].text = label
            row.cells[1].text = value
            row.cells[0].paragraphs[0].runs[0].bold = True
            row.cells[0].paragraphs[0].runs[0].font.size = Pt(10)
            row.cells[1].paragraphs[0].runs[0].font.size = Pt(10)
        
        # Type-specific instructions
        doc.add_paragraph()
        instructions = doc.add_paragraph()
        instructions.add_run("INSTRUCTIONS:").bold = True
        
        if question_type_code == "mcq":
            doc.add_paragraph("• Read each question carefully before selecting your answer")
            doc.add_paragraph("• Each question carries equal marks")
            doc.add_paragraph("• Select the MOST appropriate answer from the given options")
        elif question_type_code == "fill_blank":
            doc.add_paragraph("• Fill in the blanks with appropriate words/phrases")
            doc.add_paragraph("• Write clearly and legibly")
            doc.add_paragraph("• Each blank carries equal marks")
        elif question_type_code == "match_column":
            doc.add_paragraph("• Match items from Column A with Column B")
            doc.add_paragraph("• Write the letter of the correct match")
            doc.add_paragraph("• Each correct match carries equal marks")
        elif question_type_code == "short_response":
            if subject.lower() == "mathematics":
                doc.add_paragraph("• Show all your working clearly")
                doc.add_paragraph("• Write the final answer clearly")
                doc.add_paragraph("• Marks are awarded for method and accuracy")
            else:
                doc.add_paragraph("• Answer in 2-3 complete sentences")
                doc.add_paragraph("• Be concise and clear")
                doc.add_paragraph("• Each question carries 3-5 marks")
        elif question_type_code == "essay":
            if subject.lower() == "mathematics":
                doc.add_paragraph("• Show complete working for all calculations")
                doc.add_paragraph("• Write proofs/derivations step-by-step")
                doc.add_paragraph("• Label diagrams clearly if required")
                doc.add_paragraph("• Each question may have multiple parts")
            else:
                doc.add_paragraph("• Write detailed, well-organized paragraphs")
                doc.add_paragraph("• Support your answers with relevant examples")
                doc.add_paragraph("• Each question carries 8-10 marks")
        
        # Questions section
        doc.add_paragraph()
        section_title = "SECTION A: " + q_type_name.upper()
        questions_heading = doc.add_paragraph()
        questions_heading.add_run(section_title).bold = True
        questions_heading.runs[0].font.size = Pt(12)
        doc.add_paragraph()
        
        # Add questions based on type
        if question_type_code == "mcq":
            for i, q in enumerate(questions, start=1):
                q_para = doc.add_paragraph()
                q_para.add_run(f"{i}. ").bold = True
                q_para.add_run(q)
                q_para.space_after = Pt(8)
                doc.add_paragraph()
                
        elif question_type_code == "fill_blank":
            for i, q in enumerate(questions, start=1):
                q_para = doc.add_paragraph()
                q_para.add_run(f"{i}. ").bold = True
                q_para.add_run(q)
                q_para.space_after = Pt(8)
                doc.add_paragraph()
                
        elif question_type_code == "match_column":
            for q in questions:
                doc.add_paragraph(q)
                
        elif question_type_code in ["short_response", "essay"]:
            marks_per_q = 5 if question_type_code == "short_response" else 10
            
            # For mathematics, provide space for working
            working_lines = 5 if question_type_code == "short_response" else 12
            
            for i, q in enumerate(questions, start=1):
                q_para = doc.add_paragraph()
                q_para.add_run(f"Q{i}. ").bold = True
                q_para.add_run(q)
                q_para.add_run(f" [{marks_per_q} marks]").italic = True
                doc.add_paragraph()
                
                # Add space for answer
                if subject.lower() == "mathematics":
                    # Add "Working:" label for math questions
                    working_para = doc.add_paragraph()
                    working_para.add_run("Working:").italic = True
                    
                for _ in range(working_lines):
                    doc.add_paragraph("_" * 85)
                doc.add_paragraph()
        
        # Answer key (on separate page)
        doc.add_page_break()
        
        ans_heading = doc.add_heading("ANSWER KEY / MARKING SCHEME", level=2)
        ans_heading.alignment = WD_ALIGN_PARAGRAPH.CENTER
        for run in ans_heading.runs:
            run.font.color.rgb = RGBColor(153, 0, 0)
        
        doc.add_paragraph("(For Teacher's Use Only)", style='Intense Quote')
        doc.add_paragraph()
        
        # Add answers based on type
        if question_type_code == "mcq":
            # Answer table for MCQs
            ans_table = create_answer_table(doc, answers, 5)
            
            headers = ["Q#", "A", "B", "C", "D"]
            for idx, header in enumerate(headers):
                cell = ans_table.cell(0, idx)
                cell.text = header
                cell.paragraphs[0].runs[0].bold = True
                cell.paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            for r, ans in enumerate(answers, start=1):
                ans_table.cell(r, 0).text = str(r)
                ans_table.cell(r, 0).paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
                
                for c, option in enumerate(["A", "B", "C", "D"], start=1):
                    cell = ans_table.cell(r, c)
                    if ans == option:
                        cell.text = "✓"
                        cell.paragraphs[0].runs[0].bold = True
                        cell.paragraphs[0].runs[0].font.color.rgb = RGBColor(0, 128, 0)
                        cell.paragraphs[0].runs[0].font.size = Pt(14)
                    else:
                        cell.text = ""
                    cell.paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        elif question_type_code == "fill_blank":
            # Tabular answer key for Fill in the Blanks
            ans_table = create_answer_table(doc, answers, 2)
            
            # Headers
            headers = ["Q#", "Correct Answer"]
            for idx, header in enumerate(headers):
                cell = ans_table.cell(0, idx)
                cell.text = header
                cell.paragraphs[0].runs[0].bold = True
                cell.paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            # Fill answers
            for r, ans in enumerate(answers, start=1):
                ans_table.cell(r, 0).text = str(r)
                ans_table.cell(r, 0).paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
                ans_table.cell(r, 1).text = str(ans)
                
        elif question_type_code in ["short_response", "essay"]:
            # Expected answers for subjective questions
            for i, ans in enumerate(answers, start=1):
                ans_para = doc.add_paragraph()
                if subject.lower() == "mathematics":
                    ans_para.add_run(f"Q{i}. Solution / Marking Scheme:").bold = True
                else:
                    ans_para.add_run(f"Q{i}. Expected Answer:").bold = True
                doc.add_paragraph(str(ans))
                doc.add_paragraph()
                
        else:
            # Text-based answers for other types
            for i, ans in enumerate(answers, start=1):
                ans_para = doc.add_paragraph()
                ans_para.add_run(f"{i}. ").bold = True
                ans_para.add_run(str(ans))
                doc.add_paragraph()
        
        # Footer
        doc.add_paragraph()
        footer = doc.add_paragraph()
        footer.add_run(f"Generated by Test Generator Pro v{APP_VERSION} | Developer: Najaf Ali Sharqi | {datetime.now().strftime('%d-%m-%Y %H:%M')}").italic = True
        footer.runs[0].font.size = Pt(8)
        footer.runs[0].font.color.rgb = RGBColor(128, 128, 128)
        footer.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        # Save file
        safe_chapter = re.sub(r'[\\/*?:"<>|]', "_", chapter)
        safe_school = re.sub(r'[\\/*?:"<>|]', "_", school_name) if school_name else "School"
        safe_qtype = re.sub(r'[\\/*?:"<>|]', "_", q_type_name)
        filename = f"{safe_school}_{grade}_{subject}_{safe_chapter}_{safe_qtype}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
        
        doc.save(filename)
        return filename, None
        
    except Exception as e:
        return None, f"Error creating document: {str(e)}"

# ---------------------------
# Main generation handler with Bloom's Taxonomy distribution
# ---------------------------
def on_generate(grade, subject, chapter, test_type, school_name, question_type,
                use_bloom_mix, bloom_remember, bloom_understand, bloom_apply, 
                bloom_analyze, bloom_evaluate, bloom_create, num_questions):
    """Handle question generation request with optional Bloom's distribution"""
    
    # Validation
    if not all([grade, subject, chapter, question_type]):
        return "⚠️ Please fill all required fields", None, None
    
    if not school_name or not school_name.strip():
        school_name = "Educational Institution"
    
    try:
        # Check if using Bloom's mix
        if use_bloom_mix:
            # Validate percentages
            total_percent = bloom_remember + bloom_understand + bloom_apply + bloom_analyze + bloom_evaluate + bloom_create
            if total_percent != 100:
                return f"❌ Bloom's Taxonomy percentages must total 100% (currently {total_percent}%)", None, None
            
            bloom_distribution = {
                "Remembering": bloom_remember,
                "Understanding": bloom_understand,
                "Applying": bloom_apply,
                "Analyzing": bloom_analyze,
                "Evaluating": bloom_evaluate,
                "Creating": bloom_create
            }
            
            # Show progress
            status_msg = f"🔄 Generating {num_questions} {question_type} with mixed cognitive levels...\n"
            status_msg += f"📊 Bloom's Distribution:\n"
            for level, pct in bloom_distribution.items():
                if pct > 0:
                    status_msg += f"   • {level}: {pct}% ({int(pct * num_questions / 100)} questions)\n"
            
            # Generate with mixed levels
            questions, answers, error, q_type_code = generate_questions_multi_level(
                grade, subject, chapter, test_type, school_name, question_type,
                bloom_distribution, num_questions
            )
        else:
            # Single level generation (default to Understanding)
            status_msg = f"🔄 Generating {num_questions} {question_type}...\n"
            status_msg += f"📖 Grade: {grade} | Subject: {subject}\n"
            status_msg += f"📚 Chapter: {chapter}\n"
            status_msg += f"🎯 Cognitive Level: Understanding\n"
            status_msg += "⏳ Please wait..."
            
            questions, answers, error, q_type_code = generate_questions(
                grade, subject, chapter, "Understanding",
                num_questions, test_type, school_name, question_type
            )
        
        if error:
            return f"❌ Error: {error}", None, None
        
        if not questions or len(questions) == 0:
            return "❌ No questions generated. Please try again.", None, None
        
        # Calculate marks based on question type
        if q_type_code == "mcq":
            marks_per_q = 1
        elif q_type_code == "fill_blank":
            marks_per_q = 2
        elif q_type_code == "match_column":
            marks_per_q = len(questions) * 0.5
        elif q_type_code == "short_response":
            marks_per_q = 5
        elif q_type_code == "essay":
            marks_per_q = 10
        else:
            marks_per_q = 1
            
        total_marks = int(len(questions) * marks_per_q)
        
        # Export to Word
        file_path, doc_error = export_to_word(
            questions, answers, grade, subject, chapter, 
            "Mixed Cognitive Levels" if use_bloom_mix else "Understanding", 
            test_type, school_name, total_marks, q_type_code
        )
        
        if doc_error:
            return f"❌ {doc_error}", None, None
        
        # Success message
        success_msg = f"✅ Test generated successfully!\n\n"
        success_msg += f"📊 Test Statistics:\n"
        success_msg += f"• Question Type: {question_type}\n"
        success_msg += f"• Total Questions: {len(questions)}\n"
        success_msg += f"• Total Marks: {total_marks}\n"
        success_msg += f"• Grade: {grade} | Subject: {subject}\n"
        success_msg += f"• Chapter: {chapter}\n"
        
        if use_bloom_mix:
            success_msg += f"\n🎯 Bloom's Taxonomy Distribution:\n"
            for level, pct in bloom_distribution.items():
                if pct > 0:
                    success_msg += f"   • {level}: {pct}%\n"
        
        success_msg += f"\n📥 Download your test document below"
        
        # Preview
        preview_count = min(3, len(questions))
        preview = f"\n\n📝 PREVIEW (First {preview_count} Questions):\n" + "="*50 + "\n\n"
        
        for i, q in enumerate(questions[:preview_count], 1):
            if q_type_code == "mcq":
                preview += f"{i}. {q[:200]}...\n\n" if len(q) > 200 else f"{i}. {q}\n\n"
            else:
                preview += f"{i}. {q[:200]}...\n\n" if len(q) > 200 else f"{i}. {q}\n\n"
        
        return success_msg + preview, file_path, f"🎉 File ready: {file_path}"
        
    except Exception as e:
        return f"❌ Unexpected error: {str(e)}", None, None

# ---------------------------
# Enhanced Gradio UI
# ---------------------------
def create_ui():
    """Create professional Gradio interface"""
    
    with gr.Blocks(
        theme=gr.themes.Soft(
            primary_hue="green",
            secondary_hue="blue",
        ),
        css="""
        .header {text-align: center; margin-bottom: 20px;}
        .description {text-align: center; color: #666; margin-bottom: 30px;}
        .footer {text-align: center; margin-top: 30px; padding: 20px; color: #888;}
        .warning {background-color: #fff3cd; padding: 10px; border-radius: 5px; margin: 10px 0;}
        """
    ) as demo:
        
        # Header
        gr.Markdown(
            """
            # 📚 Test Generator Pro
            ### **Developer: Najaf Ali Sharqi**
            
            Professional Assessment Tool Aligned with Pakistani National Curriculum (2006)
            """,
            elem_classes="header"
        )
        
        gr.Markdown(
            """
            Generate high-quality assessment questions with multiple formats and Bloom's Taxonomy integration.
            Perfect for educators creating comprehensive tests for Grades 1-12.
            
            **✨ New: Enhanced Mathematics Question Generation** - Creates actual computational problems, equations, and proofs!
            """,
            elem_classes="description"
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📋 Test Configuration")
                
                school_name = gr.Textbox(
                    label="School/Institution Name",
                    placeholder="Enter your school name (optional)",
                    value="",
                    info="Will appear on the test document header"
                )
                
                test_type = gr.Dropdown(
                    choices=[
                        "Weekly Test",
                        "Monthly Test", 
                        "Mid-Term Examination",
                        "Final Examination",
                        "Practice Test",
                        "Mock Examination"
                    ],
                    value="Monthly Test",
                    label="Test Type",
                    info="Select the type of assessment"
                )
                
                question_type = gr.Dropdown(
                    choices=list(QUESTION_TYPES.keys()),
                    value="Multiple Choice Questions (MCQs)",
                    label="Question Type",
                    info="Select the format of questions to generate"
                )
                
                # Show description of selected question type
                question_type_info = gr.Markdown(
                    """
                    **Description:** Four-option questions testing recall and application
                    
                    **Format:** Question with options A, B, C, D
                    """,
                    visible=True
                )
                
                with gr.Row():
                    grade_dropdown = gr.Dropdown(
                        choices=grades,
                        label="Grade Level",
                        info="Select student grade level"
                    )
                    
                    num_questions = gr.Slider(
                        minimum=5,
                        maximum=50,
                        value=20,
                        step=5,
                        label="Number of Questions",
                        info="How many questions to generate"
                    )
                
                subject_dropdown = gr.Dropdown(
                    choices=[],
                    label="Subject",
                    info="Select subject based on grade"
                )
                
                chapter_dropdown = gr.Dropdown(
                    choices=[],
                    label="Chapter/Unit",
                    info="Select specific chapter or unit"
                )
                
                # Bloom's Taxonomy Distribution Section
                use_bloom_mix = gr.Checkbox(
                    label="📊 Use Mixed Bloom's Taxonomy Levels",
                    value=False,
                    info="Create tests with questions from multiple cognitive levels"
                )
                
                bloom_percentages = gr.Column(visible=False)
                
                with bloom_percentages:
                    gr.Markdown("### Set Percentage for Each Cognitive Level")
                    gr.Markdown("*(Total must equal 100%)*")
                    
                    with gr.Row():
                        bloom_remember = gr.Slider(
                            minimum=0, maximum=100, value=30, step=5,
                            label="Remembering %",
                            info="Recall facts and basic concepts"
                        )
                        bloom_understand = gr.Slider(
                            minimum=0, maximum=100, value=30, step=5,
                            label="Understanding %",
                            info="Explain ideas and concepts"
                        )
                    
                    with gr.Row():
                        bloom_apply = gr.Slider(
                            minimum=0, maximum=100, value=20, step=5,
                            label="Applying %",
                            info="Use information in new situations"
                        )
                        bloom_analyze = gr.Slider(
                            minimum=0, maximum=100, value=10, step=5,
                            label="Analyzing %",
                            info="Draw connections and relationships"
                        )
                    
                    with gr.Row():
                        bloom_evaluate = gr.Slider(
                            minimum=0, maximum=100, value=10, step=5,
                            label="Evaluating %",
                            info="Justify decisions and make judgments"
                        )
                        bloom_create = gr.Slider(
                            minimum=0, maximum=100, value=0, step=5,
                            label="Creating %",
                            info="Produce new or original work"
                        )
                    
                    bloom_total = gr.Textbox(
                        label="Total Percentage",
                        value="100%",
                        interactive=False
                    )
                
                difficulty_dropdown = gr.Dropdown(
                    choices=list(DIFFICULTY_LEVELS.keys()),
                    value="Understanding",
                    label="Single Cognitive Level (if not using mix)",
                    info="Select one level for entire test",
                    visible=True
                )
                
                gr.Markdown(
                    """
                    <div class="warning">
                    ⚠️ <strong>Note:</strong> Ensure GROQ_API_KEY is set in environment variables
                    </div>
                    """
                )
                
                generate_btn = gr.Button(
                    "🚀 Generate Test",
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("---")
                
                gr.Markdown(
                    """
                    ### 💡 Tips for Best Results:
                    - **MCQs**: Best for quick assessment and objective grading
                    - **Fill in the Blanks**: Tests specific terminology and facts
                    - **Match the Column**: Assesses relationships and connections
                    - **Short Response**: Evaluates understanding and explanation
                      - *Mathematics*: Generates actual calculation problems and equations
                    - **Essay Questions**: Measures critical thinking and synthesis
                      - *Mathematics*: Creates multi-step problems, proofs, and derivations
                    - **Bloom's Mix**: Use for comprehensive assessment across cognitive levels
                    - Review generated questions before final use
                    
                    ### 🔢 Mathematics Question Enhancement:
                    - Short questions now include actual numerical problems, equations, and calculations
                    - Long questions feature multi-step problems, proofs, and detailed derivations
                    - Aligned with FBISE and provincial board exam patterns
                    - Grade-appropriate complexity (Middle, Secondary, Higher Secondary)
                    
                    ### 📧 Custom App Development
                    **Need a customized assessment tool?**
                    Contact developer **Najaf Ali Sharqi** for:
                    - Custom question banks
                    - Institution-specific features
                    - Advanced analytics
                    - Integration with LMS platforms
                    """
                )
            
            with gr.Column(scale=1):
                gr.Markdown("### 📊 Output & Preview")
                
                output_text = gr.Textbox(
                    label="Generation Status & Preview",
                    lines=20,
                    placeholder="Click 'Generate Test' to create questions...",
                    interactive=False
                )
                
                file_status = gr.Textbox(
                    label="File Status",
                    lines=1,
                    interactive=False,
                    visible=False
                )
                
                output_file = gr.File(
                    label="📥 Download Test Document (.docx)",
                    file_types=[".docx"]
                )
                
                gr.Markdown(
                    """
                    ### 📄 Document Includes:
                    - ✅ Professional header with school name
                    - ✅ Complete test information table
                    - ✅ Type-specific instructions
                    - ✅ Well-formatted questions
                    - ✅ Separate answer key / marking scheme
                    - ✅ Answer spaces (for written questions)
                    - ✅ **NEW**: Proper mathematical notation and working space
                    """
                )
        
        # Footer
        gr.Markdown(
            """
            ---
            <div class="footer">
            <strong>Test Generator Pro v2.6</strong><br>
            Comprehensive Assessment Tool with Enhanced Mathematics Support<br>
            Now featuring actual computational problems, equations, and proofs for Mathematics!<br>
            Aligned with Pakistani National Curriculum (2006) & FBISE Standards<br>
            Developer: <strong>Najaf Ali Sharqi</strong><br>
            Based on Professional Item Writing Standards<br><br>
            📧 For custom app development or support: Contact Najaf Ali Sharqi
            </div>
            """,
            elem_classes="footer"
        )
        
        # Event handlers
        def update_subjects(selected_grade):
            """Update subject dropdown when grade changes"""
            if selected_grade:
                subjects = subjects_by_grade.get(selected_grade, [])
                return gr.update(choices=subjects, value=None)
            return gr.update(choices=[], value=None)
        
        def update_chapters(selected_subject, selected_grade):
            """Update chapter dropdown when subject or grade changes"""
            if selected_grade and selected_subject:
                chapters = chapters_by_subject_and_grade.get(selected_grade, {}).get(selected_subject, [])
                return gr.update(choices=chapters, value=None)
            return gr.update(choices=[], value=None)
        
        def update_question_type_info(q_type):
            """Update question type description"""
            if q_type in QUESTION_TYPES:
                info = QUESTION_TYPES[q_type]
                desc = info.get('description', '')
                format_info = info.get('format', '')
                return f"""
                **Description:** {desc}
                
                **Format:** {format_info}
                """
            return ""
        
        def toggle_bloom_sliders(use_mix):
            """Show/hide Bloom's taxonomy sliders"""
            return gr.update(visible=use_mix), gr.update(visible=not use_mix)
        
        def calculate_bloom_total(r, u, ap, an, e, c):
            """Calculate total percentage"""
            total = r + u + ap + an + e + c
            color = "green" if total == 100 else "red"
            return f'<span style="color:{color}; font-weight:bold;">{total}%</span>'
        
        # Connect event handlers
        grade_dropdown.change(
            update_subjects,
            inputs=[grade_dropdown],
            outputs=[subject_dropdown]
        )
        
        subject_dropdown.change(
            update_chapters,
            inputs=[subject_dropdown, grade_dropdown],
            outputs=[chapter_dropdown]
        )
        
        grade_dropdown.change(
            update_chapters,
            inputs=[subject_dropdown, grade_dropdown],
            outputs=[chapter_dropdown]
        )
        
        question_type.change(
            update_question_type_info,
            inputs=[question_type],
            outputs=[question_type_info]
        )
        
        use_bloom_mix.change(
            toggle_bloom_sliders,
            inputs=[use_bloom_mix],
            outputs=[bloom_percentages, difficulty_dropdown]
        )
        
        # Update total when sliders change
        for slider in [bloom_remember, bloom_understand, bloom_apply, bloom_analyze, bloom_evaluate, bloom_create]:
            slider.change(
                calculate_bloom_total,
                inputs=[bloom_remember, bloom_understand, bloom_apply, bloom_analyze, bloom_evaluate, bloom_create],
                outputs=[bloom_total]
            )
        
        generate_btn.click(
            on_generate,
            inputs=[
                grade_dropdown,
                subject_dropdown,
                chapter_dropdown,
                test_type,
                school_name,
                question_type,
                use_bloom_mix,
                bloom_remember,
                bloom_understand,
                bloom_apply,
                bloom_analyze,
                bloom_evaluate,
                bloom_create,
                num_questions
            ],
            outputs=[output_text, output_file, file_status]
        )
    
    return demo

# ---------------------------
# Launch application
# ---------------------------
try:
    print("🚀 Starting Test Generator Pro v2.6")
    print(f"📦 Gradio version: {gr.__version__}")
    
    # Check API key
    if not GROQ_API_KEY:
        print("⚠️ WARNING: GROQ_API_KEY not found in environment variables")
        print("Please set it in Hugging Face Space Settings > Repository secrets")
    else:
        print(f"✅ API Key configured (length: {len(GROQ_API_KEY)})")
    
    # Create and launch
    demo = create_ui()
    print("✅ UI created successfully")
    
    demo.queue()  # Enable queue for better handling
    demo.launch()
    
except Exception as e:
    print(f"❌ Critical Error: {str(e)}")
    import traceback
    traceback.print_exc()
    
    # Create minimal error interface
    with gr.Blocks() as error_demo:
        gr.Markdown(f"""
        # ❌ Application Error
        
        **Error Message:** {str(e)}
        
        **Common Solutions:**
        1. Check if GROQ_API_KEY is set in Space settings
        2. Verify all dependencies are installed
        3. Check the build logs for errors
        
        **Contact:** Developer - Najaf Ali Sharqi
        """)
    
    error_demo.launch()