evaluation/training_recommender.py

"""
Training Recommender Module

Provides AI-driven recommendations for next training session based on gap analysis.
"""

from typing import List, Dict, Optional, Any, Tuple
import json
from pathlib import Path
from collections import defaultdict


class TrainingRecommender:
    """
    Recommends training strategies based on performance gaps.

    Features:
    - Targeted training recommendations
    - Data generation suggestions
    - Priority-based training plans
    - Progress tracking
    """

    def __init__(self, gap_analyzer: Optional[Any] = None):
        """
        Initialize training recommender.

        Args:
            gap_analyzer: GapAnalyzer instance with performance data
        """
        self.gap_analyzer = gap_analyzer
        self.recommendations: List[Dict[str, Any]] = []

    def generate_recommendations(
        self,
        max_recommendations: int = 5,
        focus_on_weak: bool = True
    ) -> List[Dict[str, Any]]:
        """
        Generate training recommendations based on gaps.

        Args:
            max_recommendations: Maximum number of recommendations
            focus_on_weak: Prioritize weak areas over moderate ones

        Returns:
            List of training recommendations
        """
        if not self.gap_analyzer or not self.gap_analyzer.gaps:
            return [{
                'category': 'General',
                'priority': 'MEDIUM',
                'action': 'No performance data available. Start with general training.',
                'estimated_examples': 100,
                'topics': ['General training data']
            }]

        recommendations = []

        # Get gaps to address
        gaps = self.gap_analyzer.gaps

        if focus_on_weak:
            # Focus on weak and declining areas first
            priority_gaps = [
                g for g in gaps
                if g['level'] == 'WEAK' or g['trend'] == 'declining'
            ]
            if not priority_gaps:
                # Fall back to moderate areas
                priority_gaps = [g for g in gaps if g['level'] == 'MODERATE']
        else:
            # Include all non-strong areas
            priority_gaps = [g for g in gaps if g['level'] != 'STRONG']

        # Generate recommendations for top gaps
        for gap in priority_gaps[:max_recommendations]:
            recommendation = self._create_recommendation(gap)
            recommendations.append(recommendation)

        self.recommendations = recommendations
        return recommendations

    def _create_recommendation(self, gap: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create a detailed recommendation for a gap.

        Args:
            gap: Gap analysis data

        Returns:
            Training recommendation
        """
        category = gap['category']
        avg_score = gap['avg_score']
        level = gap['level']

        # Determine number of examples needed
        if avg_score < 40:
            estimated_examples = 100
            intensity = "intensive"
        elif avg_score < 60:
            estimated_examples = 50
            intensity = "moderate"
        else:
            estimated_examples = 25
            intensity = "light"

        # Generate specific action items
        action_items = self._generate_action_items(category, level, gap['trend'])

        # Suggest topics
        topics = self._suggest_topics(category)

        recommendation = {
            'category': category,
            'priority': gap['priority'],
            'current_score': avg_score,
            'trend': gap['trend'],
            'intensity': intensity,
            'estimated_examples': estimated_examples,
            'action': f"Focus on {category} with {intensity} training",
            'action_items': action_items,
            'suggested_topics': topics,
            'expected_improvement': self._estimate_improvement(avg_score, estimated_examples)
        }

        return recommendation

    def _generate_action_items(
        self,
        category: str,
        level: str,
        trend: str
    ) -> List[str]:
        """
        Generate specific action items for a category.

        Args:
            category: Category name
            level: Performance level
            trend: Performance trend

        Returns:
            List of action items
        """
        items = []

        # Base recommendations based on level
        if level == 'WEAK':
            items.append(f"Add 50-100 {category} examples to training data")
            items.append(f"Review fundamental {category} concepts")
            items.append("Include diverse question types and difficulty levels")
        elif level == 'MODERATE':
            items.append(f"Add 25-50 {category} examples focusing on edge cases")
            items.append(f"Review intermediate {category} topics")
        else:
            items.append(f"Maintain current {category} performance with 10-20 examples")

        # Add trend-specific items
        if trend == 'declining':
            items.append("⚠️ Address declining performance immediately")
            items.append(f"Review recent {category} training data for quality issues")
        elif trend == 'improving':
            items.append("✅ Continue current training approach")

        # Add testing recommendation
        items.append(f"Test specifically on {category} after training")

        return items

    def _suggest_topics(self, category: str) -> List[str]:
        """
        Suggest specific topics for a category.

        Args:
            category: Category name

        Returns:
            List of suggested topics
        """
        # Topic suggestions by common categories
        topic_map = {
            'Estate Planning': [
                'Revocable living trusts',
                'Wills and probate',
                'Power of attorney',
                'Estate tax strategies',
                'Charitable giving',
                'Trust structures'
            ],
            'Retirement Planning': [
                '401(k) and IRA strategies',
                'Required minimum distributions',
                'Social Security optimization',
                'Pension planning',
                'Retirement income strategies',
                'Healthcare in retirement'
            ],
            'Tax Planning': [
                'Tax-efficient investing',
                'Capital gains strategies',
                'Tax-loss harvesting',
                'Deductions and credits',
                'Alternative minimum tax',
                'Estate and gift taxes'
            ],
            'Investment Planning': [
                'Asset allocation',
                'Portfolio diversification',
                'Risk management',
                'Modern portfolio theory',
                'Performance evaluation',
                'Rebalancing strategies'
            ],
            'Insurance Planning': [
                'Life insurance types',
                'Disability insurance',
                'Long-term care insurance',
                'Property and casualty',
                'Umbrella policies',
                'Insurance needs analysis'
            ],
            'Education Planning': [
                '529 plans',
                'Coverdell ESAs',
                'Financial aid strategies',
                'Student loan planning',
                'Education tax benefits'
            ]
        }

        # Return specific topics if available, otherwise generic suggestions
        if category in topic_map:
            return topic_map[category]
        else:
            return [
                f"Fundamental {category} concepts",
                f"Intermediate {category} topics",
                f"Advanced {category} strategies",
                f"{category} best practices",
                f"Common {category} scenarios"
            ]

    def _estimate_improvement(
        self,
        current_score: float,
        num_examples: int
    ) -> str:
        """
        Estimate expected improvement from training.

        Args:
            current_score: Current performance score
            num_examples: Number of training examples

        Returns:
            Improvement estimate description
        """
        # Simple heuristic: more examples = more improvement, diminishing returns
        base_improvement = min(num_examples * 0.3, 30)  # Max 30% improvement

        # Lower scores have more room for improvement
        if current_score < 40:
            multiplier = 1.5
        elif current_score < 60:
            multiplier = 1.2
        else:
            multiplier = 0.8

        estimated_improvement = base_improvement * multiplier
        new_score = min(current_score + estimated_improvement, 95)

        return f"+{estimated_improvement:.1f}% (to ~{new_score:.1f}%)"

    def create_training_plan(
        self,
        priority: str = "all",
        include_data_generation: bool = True
    ) -> Dict[str, Any]:
        """
        Create a comprehensive training plan.

        Args:
            priority: Focus on "high", "medium", "low", or "all" priority items
            include_data_generation: Include data generation instructions

        Returns:
            Training plan
        """
        if not self.recommendations:
            self.generate_recommendations()

        # Filter by priority
        if priority.upper() != "ALL":
            filtered_recs = [
                r for r in self.recommendations
                if r['priority'] == priority.upper()
            ]
        else:
            filtered_recs = self.recommendations

        # Calculate totals
        total_examples = sum(r['estimated_examples'] for r in filtered_recs)
        categories = [r['category'] for r in filtered_recs]

        plan = {
            'plan_name': f"Training Plan - Priority: {priority.title()}",
            'num_focus_areas': len(filtered_recs),
            'focus_categories': categories,
            'total_examples_needed': total_examples,
            'recommendations': filtered_recs,
            'execution_steps': self._generate_execution_steps(filtered_recs),
        }

        if include_data_generation:
            plan['data_generation'] = self._generate_data_instructions(filtered_recs)

        return plan

    def _generate_execution_steps(
        self,
        recommendations: List[Dict[str, Any]]
    ) -> List[str]:
        """Generate step-by-step execution plan."""
        steps = [
            "1. Review gap analysis and recommendations",
            "2. Prepare training data:"
        ]

        for i, rec in enumerate(recommendations, 1):
            steps.append(f"   {chr(96+i)}. {rec['category']}: {rec['estimated_examples']} examples")

        steps.extend([
            "3. Generate or collect training examples",
            "4. Validate data quality (score > 60)",
            "5. Execute training session",
            "6. Run targeted benchmark tests",
            "7. Analyze results and compare to previous performance",
            "8. Iterate if needed"
        ])

        return steps

    def _generate_data_instructions(
        self,
        recommendations: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Generate data generation instructions."""
        instructions = {
            'method': 'synthetic_generation',
            'by_category': {}
        }

        for rec in recommendations:
            category = rec['category']
            instructions['by_category'][category] = {
                'num_examples': rec['estimated_examples'],
                'difficulty': 'mixed',
                'topics': rec['suggested_topics'],
                'sample_prompt': f"Generate financial advisory questions about {category}, covering topics like: {', '.join(rec['suggested_topics'][:3])}"
            }

        return instructions

    def generate_report(self) -> str:
        """
        Generate human-readable training recommendations report.

        Returns:
            Formatted report
        """
        if not self.recommendations:
            self.generate_recommendations()

        report = ["=" * 80]
        report.append("TRAINING RECOMMENDATIONS REPORT")
        report.append("=" * 80)
        report.append("")

        if not self.recommendations:
            report.append("No recommendations available. Performance data needed.")
            return "\n".join(report)

        # Summary
        total_examples = sum(r['estimated_examples'] for r in self.recommendations)
        report.append(f"Total Focus Areas: {len(self.recommendations)}")
        report.append(f"Total Training Examples Needed: {total_examples}")
        report.append("")

        # Detailed recommendations
        report.append("RECOMMENDED TRAINING PRIORITIES:")
        report.append("-" * 80)

        for i, rec in enumerate(self.recommendations, 1):
            priority_symbol = {
                'HIGH': '🔴',
                'MEDIUM': '🟡',
                'LOW': '🟢'
            }.get(rec['priority'], '⚪')

            report.append(f"\n{i}. {priority_symbol} {rec['category']} - Priority: {rec['priority']}")
            report.append(f"   Current Score: {rec['current_score']:.1f}%")
            report.append(f"   Trend: {rec['trend']}")
            report.append(f"   Training Intensity: {rec['intensity']}")
            report.append(f"   Recommended Examples: {rec['estimated_examples']}")
            report.append(f"   Expected Improvement: {rec['expected_improvement']}")
            report.append("")
            report.append("   Action Items:")
            for item in rec['action_items']:
                report.append(f"     • {item}")
            report.append("")
            report.append("   Suggested Topics:")
            for topic in rec['suggested_topics'][:5]:  # Top 5 topics
                report.append(f"     - {topic}")

        report.append("")
        report.append("=" * 80)
        report.append("NEXT STEPS:")
        report.append("")
        report.append("1. Generate training data for priority categories")
        report.append("2. Focus on weak/declining areas first")
        report.append("3. Use diverse examples covering suggested topics")
        report.append("4. Run targeted tests after training")
        report.append("5. Track improvement and adjust strategy")
        report.append("=" * 80)

        return "\n".join(report)

    def save_recommendations(self, filepath: str):
        """
        Save recommendations to JSON file.

        Args:
            filepath: Output file path
        """
        if not self.recommendations:
            self.generate_recommendations()

        Path(filepath).parent.mkdir(parents=True, exist_ok=True)

        data = {
            'recommendations': self.recommendations,
            'training_plan': self.create_training_plan(),
            'report': self.generate_report()
        }

        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

        print(f"Recommendations saved to: {filepath}")

    def load_recommendations(self, filepath: str):
        """
        Load recommendations from JSON file.

        Args:
            filepath: Input file path
        """
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)

        self.recommendations = data.get('recommendations', [])

    def get_quick_wins(self) -> List[Dict[str, Any]]:
        """
        Identify quick wins - categories that can improve quickly.

        Returns:
            List of quick win opportunities
        """
        if not self.recommendations:
            self.generate_recommendations()

        # Quick wins: moderate performance, not too many examples needed
        quick_wins = [
            rec for rec in self.recommendations
            if 50 <= rec['current_score'] < 70 and rec['estimated_examples'] <= 50
        ]

        return quick_wins

    def prioritize_by_impact(self) -> List[Dict[str, Any]]:
        """
        Sort recommendations by expected impact.

        Returns:
            Recommendations sorted by impact
        """
        if not self.recommendations:
            self.generate_recommendations()

        # Calculate impact score (combination of priority and potential improvement)
        def impact_score(rec):
            priority_weight = {'HIGH': 3, 'MEDIUM': 2, 'LOW': 1}
            improvement_potential = 100 - rec['current_score']
            return priority_weight.get(rec['priority'], 1) * improvement_potential

        sorted_recs = sorted(
            self.recommendations,
            key=impact_score,
            reverse=True
        )

        return sorted_recs