agents/practice_generator.py

"""
Practice Generator Agent - Creates customized practice problems
"""

import random
import json
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
from dataclasses import dataclass
import numpy as np


@dataclass
class PracticeProblem:
    """Practice problem structure"""
    id: str
    topic: str
    question: str
    options: List[str]  # For multiple choice
    correct_answer: str
    explanation: str
    difficulty: str  # easy, medium, hard
    solution_steps: List[str]
    hints: List[str]
    learning_objectives: List[str]


class PracticeGeneratorAgent:
    """Generates customized practice problems for CFA learning"""
    
    def __init__(self):
        self.problem_templates = self._load_problem_templates()
        self.generated_problems = {}
        self.user_performance = {}
        
    def _load_problem_templates(self) -> Dict[str, List[Dict]]:
        """Load problem templates by topic and difficulty"""
        return {
            'Time Value of Money': {
                'easy': [
                    {
                        'template': "Calculate the future value of ${amount} invested at {rate}% for {years} years.",
                        'variables': {
                            'amount': (1000, 10000, 1000),
                            'rate': (5, 10, 1),
                            'years': (5, 20, 5)
                        },
                        'solution_func': self._solve_fv_simple
                    }
                ],
                'medium': [
                    {
                        'template': "What is the present value of an annuity that pays ${payment} annually for {years} years with a discount rate of {rate}%?",
                        'variables': {
                            'payment': (500, 2000, 100),
                            'years': (10, 30, 5),
                            'rate': (6, 12, 1)
                        },
                        'solution_func': self._solve_pv_annuity
                    }
                ]
            },
            'Statistics': {
                'easy': [
                    {
                        'template': "Calculate the mean of the following dataset: {dataset}",
                        'variables': {
                            'dataset': self._generate_dataset
                        },
                        'solution_func': self._solve_mean
                    }
                ]
            }
        }
    
    def process(self,
                query: str,
                user_profile: Dict[str, Any],
                conversation_history: List[Dict] = None) -> Dict[str, Any]:
        """
        Generate practice problems based on user needs
        
        Args:
            query: User's request for practice
            user_profile: User's learning profile
            conversation_history: Previous conversation
        
        Returns:
            Generated practice problems
        """
        # Determine topics from query or profile
        topics = self._extract_topics_from_query(query, user_profile)
        
        # Determine difficulty based on user level
        difficulty = self._determine_difficulty(user_profile)
        
        # Generate problems
        problems = []
        for topic in topics[:2]:  # Limit to 2 topics
            topic_problems = self._generate_problems_for_topic(topic, difficulty, 2)  # 2 problems per topic
            problems.extend(topic_problems)
        
        # Format response
        response = self._format_problems_response(problems)
        
        # Store generated problems
        problem_id = f"practice_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.generated_problems[problem_id] = {
            'problems': [p.__dict__ for p in problems],
            'user_profile': user_profile,
            'timestamp': datetime.now().isoformat()
        }
        
        return {
            'response': response,
            'problems': [p.__dict__ for p in problems],
            'problem_count': len(problems),
            'topics': topics[:2],
            'difficulty': difficulty,
            'confidence': 0.9,
            'metadata': {
                'problem_id': problem_id,
                'estimated_time': len(problems) * 5  # 5 minutes per problem
            }
        }
    
    def _extract_topics_from_query(self, query: str, user_profile: Dict[str, Any]) -> List[str]:
        """Extract topics from query or use profile defaults"""
        query_lower = query.lower()
        
        # Topic keywords mapping
        topic_keywords = {
            'tvm': 'Time Value of Money',
            'time value': 'Time Value of Money',
            'npv': 'Time Value of Money',
            'irr': 'Time Value of Money',
            'statistics': 'Statistics',
            'probability': 'Statistics',
            'regression': 'Statistics',
            'correlation': 'Statistics',
            'economics': 'Economics',
            'supply demand': 'Economics'
        }
        
        found_topics = []
        for keyword, topic in topic_keywords.items():
            if keyword in query_lower:
                found_topics.append(topic)
        
        # If no topics found in query, use profile
        if not found_topics:
            found_topics = user_profile.get('study_focus', ['Time Value of Money', 'Statistics'])
        
        return list(set(found_topics))  # Remove duplicates
    
    def _determine_difficulty(self, user_profile: Dict[str, Any]) -> str:
        """Determine appropriate difficulty level"""
        experience = user_profile.get('experience_level', 'beginner')
        
        difficulty_map = {
            'beginner': 'easy',
            'intermediate': 'medium',
            'advanced': 'hard'
        }
        
        return difficulty_map.get(experience, 'medium')
    
    def _generate_problems_for_topic(self, topic: str, difficulty: str, count: int = 2) -> List[PracticeProblem]:
        """Generate problems for a specific topic and difficulty"""
        problems = []
        
        topic_templates = self.problem_templates.get(topic, {})
        difficulty_templates = topic_templates.get(difficulty, topic_templates.get('medium', []))
        
        if not difficulty_templates:
            # Fallback to any difficulty
            for diff in ['easy', 'medium', 'hard']:
                if diff in topic_templates:
                    difficulty_templates = topic_templates[diff]
                    break
        
        if not difficulty_templates:
            # Create a generic problem
            return [self._create_generic_problem(topic, difficulty)]
        
        for i in range(min(count, len(difficulty_templates))):
            template = random.choice(difficulty_templates)
            problem = self._instantiate_template(template, topic, difficulty, i+1)
            problems.append(problem)
        
        return problems
    
    def _instantiate_template(self, template: Dict, topic: str, difficulty: str, problem_num: int) -> PracticeProblem:
        """Instantiate a problem template with specific values"""
        # Generate variables
        variables = {}
        for var_name, var_spec in template['variables'].items():
            if callable(var_spec):
                variables[var_name] = var_spec()
            elif isinstance(var_spec, tuple) and len(var_spec) == 3:
                # (min, max, step)
                min_val, max_val, step = var_spec
                if step == 1:
                    variables[var_name] = random.randint(min_val, max_val)
                else:
                    choices = list(range(min_val, max_val + 1, step))
                    variables[var_name] = random.choice(choices)
            else:
                variables[var_name] = var_spec
        
        # Format question
        question = template['template'].format(**variables)
        
        # Generate solution
        solution_data = template['solution_func'](variables)
        
        # Generate options for multiple choice
        options = self._generate_options(solution_data['correct_answer'], variables)
        
        # Create problem ID
        problem_id = f"{topic[:3].upper()}_{difficulty[:1].upper()}_{problem_num}_{datetime.now().strftime('%H%M%S')}"
        
        return PracticeProblem(
            id=problem_id,
            topic=topic,
            question=question,
            options=options,
            correct_answer=str(solution_data['correct_answer']),
            explanation=solution_data.get('explanation', ''),
            difficulty=difficulty,
            solution_steps=solution_data.get('steps', []),
            hints=solution_data.get('hints', []),
            learning_objectives=solution_data.get('learning_objectives', [])
        )
    
    def _generate_dataset(self) -> str:
        """Generate a random dataset for statistics problems"""
        size = random.randint(5, 10)
        data = [random.randint(10, 100) for _ in range(size)]
        return ', '.join(str(x) for x in data)
    
    def _solve_fv_simple(self, variables: Dict) -> Dict:
        """Solve simple future value problem"""
        pv = variables['amount']
        r = variables['rate'] / 100
        n = variables['years']
        
        fv = pv * (1 + r) ** n
        fv_rounded = round(fv, 2)
        
        return {
            'correct_answer': f"${fv_rounded:,.2f}",
            'explanation': f"FV = PV × (1 + r)^n = ${pv} × (1 + {variables['rate']}%)^{n} = ${fv_rounded:,.2f}",
            'steps': [
                f"Identify PV = ${pv}",
                f"Convert rate to decimal: r = {variables['rate']}% = {r}",
                f"Identify n = {n} years",
                f"Apply formula: FV = {pv} × (1 + {r})^{n}",
                f"Calculate: FV = ${fv_rounded:,.2f}"
            ],
            'hints': [
                "Use the future value formula: FV = PV × (1 + r)^n",
                "Remember to convert percentage to decimal",
                "Be careful with the exponent"
            ],
            'learning_objectives': [
                "Understand future value calculation",
                "Apply compounding formula",
                "Convert percentages to decimals"
            ]
        }
    
    def _solve_pv_annuity(self, variables: Dict) -> Dict:
        """Solve present value of annuity problem"""
        pmt = variables['payment']
        r = variables['rate'] / 100
        n = variables['years']
        
        if r == 0:
            pv = pmt * n
        else:
            pv = pmt * (1 - (1 + r) ** -n) / r
        
        pv_rounded = round(pv, 2)
        
        return {
            'correct_answer': f"${pv_rounded:,.2f}",
            'explanation': f"PV of annuity = PMT × [1 - (1 + r)^-n] / r = ${pmt} × [1 - (1 + {variables['rate']}%)^-{n}] / {variables['rate']}% = ${pv_rounded:,.2f}",
            'steps': [
                f"Identify PMT = ${pmt}",
                f"Convert rate to decimal: r = {variables['rate']}% = {r}",
                f"Identify n = {n} years",
                f"Apply annuity formula: PV = {pmt} × [1 - (1 + {r})^-{n}] / {r}",
                f"Calculate: PV = ${pv_rounded:,.2f}"
            ],
            'hints': [
                "Use the present value of annuity formula",
                "Watch for the discount rate conversion",
                "The formula handles regular payments"
            ],
            'learning_objectives': [
                "Calculate present value of annuities",
                "Apply time value of money to multiple cash flows",
                "Understand annuity valuation"
            ]
        }
    
    def _solve_mean(self, variables: Dict) -> Dict:
        """Solve mean calculation problem"""
        dataset_str = variables['dataset']
        numbers = [float(x.strip()) for x in dataset_str.split(',')]
        mean = sum(numbers) / len(numbers)
        
        return {
            'correct_answer': f"{mean:.2f}",
            'explanation': f"Mean = Sum of values / Number of values = ({' + '.join(str(n) for n in numbers)}) / {len(numbers)} = {mean:.2f}",
            'steps': [
                f"Extract numbers: {numbers}",
                f"Calculate sum: {sum(numbers)}",
                f"Count values: {len(numbers)}",
                f"Divide: {sum(numbers)} / {len(numbers)} = {mean:.2f}"
            ],
            'hints': [
                "Mean = (Sum of all values) / (Number of values)",
                "Add all numbers first, then divide by count",
                "Check your arithmetic carefully"
            ],
            'learning_objectives': [
                "Calculate mean of a dataset",
                "Understand central tendency",
                "Perform basic statistical calculations"
            ]
        }
    
    def _generate_options(self, correct_answer: str, variables: Dict) -> List[str]:
        """Generate multiple choice options"""
        options = [correct_answer]
        
        # Generate plausible wrong answers
        if '$' in correct_answer:
            # Monetary answers
            value = float(correct_answer.replace('$', '').replace(',', ''))
            for factor in [0.8, 1.2, 0.5, 1.5]:
                wrong = value * factor
                options.append(f"${wrong:,.2f}")
        else:
            # Numeric answers
            try:
                value = float(correct_answer)
                for offset in [-value*0.1, value*0.1, -value*0.2, value*0.2]:
                    options.append(f"{value + offset:.2f}")
            except:
                # Text answers - add common misconceptions
                options.extend(["None of the above", "All of the above", "Cannot be determined"])
        
        # Shuffle and limit to 4 options
        random.shuffle(options)
        return options[:4]
    
    def _create_generic_problem(self, topic: str, difficulty: str) -> PracticeProblem:
        """Create a generic problem when no templates are available"""
        return PracticeProblem(
            id=f"GEN_{topic[:3].upper()}_{datetime.now().strftime('%H%M%S')}",
            topic=topic,
            question=f"Explain the key concept of {topic} in your own words.",
            options=["Concept explanation", "Formula application", "Real-world example", "All of the above"],
            correct_answer="Concept explanation",
            explanation="This question tests your conceptual understanding of the topic.",
            difficulty=difficulty,
            solution_steps=[
                "Review the core concepts of the topic",
                "Identify the most fundamental principle",
                "Explain it clearly in simple terms"
            ],
            hints=[
                "Focus on the foundational idea",
                "Think about why this concept matters",
                "Consider how it relates to other topics"
            ],
            learning_objectives=[
                "Articulate core concepts",
                "Demonstrate conceptual understanding",
                "Communicate financial ideas clearly"
            ]
        )
    
    def _format_problems_response(self, problems: List[PracticeProblem]) -> str:
        """Format problems for user presentation"""
        response = "## 📝 Practice Problems\n\n"
        response += f"I've generated {len(problems)} practice problems for you:\n\n"
        
        for i, problem in enumerate(problems, 1):
            response += f"### Problem {i}: {problem.topic} ({problem.difficulty.title()})\n"
            response += f"{problem.question}\n\n"
            
            if problem.options:
                response += "**Options:**\n"
                for j, option in enumerate(problem.options, 1):
                    response += f"{j}. {option}\n"
            
            response += f"\n**Hints available:** {len(problem.hints)}\n"
            response += f"**Learning objectives:** {', '.join(problem.learning_objectives[:2])}\n"
            response += "---\n\n"
        
        response += "**Instructions:**\n"
        response += "1. Try to solve each problem on your own first\n"
        response += "2. Use hints if you get stuck\n"
        response += "3. Review the solution steps carefully\n"
        response += "4. Ask for explanations of any unclear concepts\n"
        
        return response
    
    def submit_answer(self, problem_id: str, user_answer: str, user_id: str) -> Dict[str, Any]:
        """Submit and evaluate an answer to a practice problem"""
        if problem_id not in self.generated_problems:
            return {'error': 'Problem not found'}
        
        # Find the problem
        problem_data = self.generated_problems[problem_id]
        problem = None
        for p_dict in problem_data['problems']:
            if p_dict['id'] == problem_id:
                problem = PracticeProblem(**p_dict)
                break
        
        if not problem:
            return {'error': 'Problem not found in collection'}
        
        # Evaluate answer
        is_correct = self._evaluate_answer(user_answer, problem.correct_answer)
        
        # Update user performance
        if user_id not in self.user_performance:
            self.user_performance[user_id] = {
                'problems_attempted': 0,
                'problems_correct': 0,
                'by_topic': {},
                'by_difficulty': {}
            }
        
        self.user_performance[user_id]['problems_attempted'] += 1
        if is_correct:
            self.user_performance[user_id]['problems_correct'] += 1
        
        # Prepare feedback
        feedback = self._generate_answer_feedback(is_correct, problem)
        
        return {
            'is_correct': is_correct,
            'correct_answer': problem.correct_answer,
            'feedback': feedback,
            'explanation': problem.explanation,
            'solution_steps': problem.solution_steps,
            'performance': {
                'accuracy': self.user_performance[user_id]['problems_correct'] / 
                          self.user_performance[user_id]['problems_attempted'],
                'attempts': self.user_performance[user_id]['problems_attempted']
            }
        }
    
    def _evaluate_answer(self, user_answer: str, correct_answer: str) -> bool:
        """Evaluate if user answer is correct"""
        # Simple exact match for now
        # In production, use more sophisticated matching (numeric tolerance, case insensitive, etc.)
        return user_answer.strip().lower() == correct_answer.strip().lower()
    
    def _generate_answer_feedback(self, is_correct: bool, problem: PracticeProblem) -> str:
        """Generate feedback for user answer"""
        if is_correct:
            return f"✅ Correct! {random.choice(['Excellent work!', 'Well done!', 'Perfect!'])}"
        else:
            feedback = "❌ Not quite right. "
            feedback += f"The correct answer is: **{problem.correct_answer}**\n\n"
            feedback += "**Review these steps:**\n"
            for i, step in enumerate(problem.solution_steps[:3], 1):
                feedback += f"{i}. {step}\n"
            return feedback
    
    def get_user_performance(self, user_id: str) -> Dict[str, Any]:
        """Get user's practice performance"""
        return self.user_performance.get(user_id, {
            'problems_attempted': 0,
            'problems_correct': 0,
            'accuracy': 0
        })