#!/usr/bin/env python3
"""
CaDaR Improvement Analysis Script
Analyze benchmark results and provide specific improvement recommendations
"""

import pandas as pd
import numpy as np
import re
from typing import Dict, List, Tuple, Set
import logging

# Set up logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class CaDaRImprovementAnalyzer:
    def __init__(self, results_csv_path: str):
        self.df = pd.read_csv(results_csv_path)
        self.improvement_areas = {}
        self.specific_issues = {}

    def analyze_wer_distribution(self):
        """Analyze WER distribution to identify improvement areas"""
        logger.info("Analyzing WER distribution...")

        # Detailed WER analysis
        wer_ranges = [
            (0, 0, "Perfect"),
            (0.001, 0.1, "Excellent"),
            (0.1, 0.3, "Good"),
            (0.3, 0.5, "Fair"),
            (0.5, 1.0, "Poor"),
            (1.0, float("inf"), "Critical"),
        ]

        total_samples = len(self.df)
        results = []

        for lower, upper, category in wer_ranges:
            if upper == float("inf"):
                count = len(self.df[self.df["wer_score"] >= lower])
            else:
                count = len(
                    self.df[
                        (self.df["wer_score"] >= lower) & (self.df["wer_score"] < upper)
                    ]
                )

            percentage = (count / total_samples) * 100
            results.append(
                {
                    "category": category,
                    "wer_range": f"[{lower}, {upper if upper != float('inf') else '∞'})",
                    "count": count,
                    "percentage": percentage,
                }
            )

        wer_df = pd.DataFrame(results)
        logger.info("\nWER Distribution Analysis:")
        logger.info(wer_df.to_string(index=False))

        self.improvement_areas["wer_distribution"] = wer_df
        return wer_df

    def identify_failure_patterns(self):
        """Identify specific patterns causing failures"""
        logger.info("Identifying failure patterns...")

        # Analyze perfect cases
        perfect_cases = self.df[self.df["wer_score"] == 0]
        logger.info(f"\nPerfect cases analysis ({len(perfect_cases)} samples):")
        for _, row in perfect_cases.head(10).iterrows():
            logger.info(f"  {row['original_arabizi']} -> {row['arabic']}")

        # Analyze critical failures (WER >= 1.0)
        critical_cases = self.df[self.df["wer_score"] >= 1.0]
        logger.info(f"\nCritical failures analysis ({len(critical_cases)} samples):")

        critical_patterns = []
        for _, row in critical_cases.head(10).iterrows():
            original = row["original_arabizi"]
            arabic = row["arabic"]
            roundtrip = row["roundtrip_arabizi"]

            # Identify specific issues
            issues = []

            # Check character mapping issues
            if len(str(original).split()) != len(str(roundtrip).split()):
                issues.append("Length mismatch")

            # Check for specific problematic characters
            if "7" in original and "ح" not in arabic:
                issues.append("ḥ (7) mapping issue")
            if "3" in original and "ع" not in arabic:
                issues.append("ʿayn (3) mapping issue")
            if "9" in original and "ق" not in arabic:
                issues.append("qāf (9) mapping issue")
            if "8" in original and "ح" not in arabic and "خ" not in arabic:
                issues.append("ḵāf (8) mapping issue")

            critical_patterns.append(
                {
                    "original": original,
                    "arabic": arabic,
                    "roundtrip": roundtrip,
                    "wer": row["wer_score"],
                    "issues": issues,
                }
            )

            logger.info(f"  Original: {original}")
            logger.info(f"  Arabic: {arabic}")
            logger.info(f"  Roundtrip: {roundtrip}")
            logger.info(f"  Issues: {', '.join(issues) if issues else 'Unknown'}")
            logger.info("-" * 50)

        self.improvement_areas["critical_failures"] = critical_patterns
        return critical_patterns

    def analyze_character_mapping_accuracy(self):
        """Analyze specific character mapping issues"""
        logger.info("Analyzing character mapping accuracy...")

        # Common problematic character mappings
        problematic_chars = {
            "7": ["ح", "ḥ", "ھ"],
            "3": ["ع", "ʿ"],
            "9": ["ق", "ٯ"],
            "8": ["ح", "خ", "ھ"],
            "5": ["خ"],
            "6": ["ط"],
            "2": ["ء", "أ", "إ"],
            "(": ["("],
            ")": [")"],
            "!": ["!"],
            "?": ["؟"],
        }

        mapping_issues = []

        for _, row in self.df.iterrows():
            original = row["original_arabizi"]
            arabic = row["arabic"]
            roundtrip = row["roundtrip_arabizi"]

            if row["wer_score"] > 0.5:  # Focus on poor performance cases
                for char, expected_arabic in problematic_chars.items():
                    if char in original:
                        found_in_arabic = any(
                            expected in arabic
                            for expected in expected_arabic
                            if isinstance(expected_arabic, list)
                        )
                        found_in_roundtrip = any(
                            expected in roundtrip
                            for expected in expected_arabic
                            if isinstance(expected_arabic, list)
                        )

                        if not found_in_arabic or not found_in_roundtrip:
                            mapping_issues.append(
                                {
                                    "char": char,
                                    "expected": expected_arabic,
                                    "original": original,
                                    "found_in_arabic": found_in_arabic,
                                    "found_in_roundtrip": found_in_roundtrip,
                                    "wer": row["wer_score"],
                                }
                            )

        logger.info(f"\nCharacter mapping issues found: {len(mapping_issues)}")
        logger.info("Most common mapping problems:")

        # Count frequency of each issue
        char_issue_counts = {}
        for issue in mapping_issues:
            char = issue["char"]
            if char not in char_issue_counts:
                char_issue_counts[char] = 0
            char_issue_counts[char] += 1

        for char, count in sorted(
            char_issue_counts.items(), key=lambda x: x[1], reverse=True
        ):
            logger.info(f"  {char}: {count} issues")

        self.improvement_areas["character_mapping"] = mapping_issues
        return mapping_issues

    def generate_improvement_recommendations(self):
        """Generate specific improvement recommendations based on analysis"""
        logger.info("Generating improvement recommendations...")

        recommendations = []

        # WER-based recommendations
        perfect_percentage = (
            len(self.df[self.df["wer_score"] == 0]) / len(self.df) * 100
        )
        if perfect_percentage < 5:
            recommendations.append(
                {
                    "priority": "HIGH",
                    "area": "Overall Performance",
                    "issue": f"Only {perfect_percentage:.1f}% perfect conversions (target: >20%)",
                    "recommendation": "Improve core mapping algorithms and expand linguistic rules",
                    "estimated_impact": "Could improve perfect conversions to 15-25%",
                }
            )

        # Character mapping recommendations
        char_mapping_issues = self.improvement_areas.get("character_mapping", [])
        if len(char_mapping_issues) > 50:
            recommendations.append(
                {
                    "priority": "HIGH",
                    "area": "Character Mapping",
                    "issue": "Multiple character mapping inconsistencies",
                    "recommendation": "Implement context-aware character mapping and validate all common Darija character patterns",
                    "estimated_impact": "Could reduce WER by 15-20%",
                }
            )

        # Word-level recommendations
        poor_performance = self.df[self.df["wer_score"] >= 0.5]
        if len(poor_performance) > len(self.df) * 0.8:
            recommendations.append(
                {
                    "priority": "HIGH",
                    "area": "Word-level Processing",
                    "issue": "80%+ of conversions have poor performance",
                    "recommendation": "Implement better word boundary detection and handle Darija-specific contractions",
                    "estimated_impact": "Could improve WER by 0.2-0.3 points",
                }
            )

        # Vowel handling recommendations
        vowel_issues = self.df[
            (self.df["wer_score"] > 0.3) & (self.df["wer_score"] < 0.8)
        ]
        if len(vowel_issues) > 0:
            recommendations.append(
                {
                    "priority": "MEDIUM",
                    "area": "Vowel Handling",
                    "issue": "Medium performance cases likely have vowel-related issues",
                    "recommendation": "Implement better vowel preservation rules for Darija phonology",
                    "estimated_impact": "Could improve medium cases by 30-50%",
                }
            )

        # Punctuation handling
        punctuation_issues = self.df[
            self.df["roundtrip_arabizi"].str.contains(r"[!?.]", na=False)
        ]
        if len(punctuation_issues) > 0:
            recommendations.append(
                {
                    "priority": "MEDIUM",
                    "area": "Punctuation Handling",
                    "issue": "Punctuation inconsistencies in roundtrip",
                    "recommendation": "Improve bidirectional punctuation mapping",
                    "estimated_impact": "Minor improvement (1-2% WER reduction)",
                }
            )

        # High priority specific improvements
        recommendations.extend(
            [
                {
                    "priority": "HIGH",
                    "area": "Emphatic Consonant Recognition",
                    "issue": "Misidentification of emphatic vs regular consonants",
                    "recommendation": "Expand emphatic lexicon and improve root-based detection",
                    "estimated_impact": "Could handle 20-30% more cases correctly",
                },
                {
                    "priority": "HIGH",
                    "area": "Number Pattern Recognition",
                    "issue": "Darija number notation (7na, 3la, etc.) not handled properly",
                    "recommendation": "Expand number pattern lexicon with common Darija expressions",
                    "estimated_impact": "Could improve 10-15% of cases with numbers",
                },
                {
                    "priority": "HIGH",
                    "area": "Clitic Handling",
                    "issue": "Function word prefixes (و, ف, ب, ل) not properly separated",
                    "recommendation": "Implement better clitic splitting algorithm",
                    "estimated_impact": "Could improve word boundary accuracy by 25%",
                },
                {
                    "priority": "MEDIUM",
                    "area": "Validation Stage",
                    "issue": "Overly strict validation causing failures",
                    "recommendation": "Relax validation rules for common Darija patterns",
                    "estimated_impact": "Could recover 5-10% of failed conversions",
                },
                {
                    "priority": "MEDIUM",
                    "area": "Error Recovery",
                    "issue": "No graceful handling of conversion failures",
                    "recommendation": "Implement fallback mechanisms and partial conversion",
                    "estimated_impact": "Could provide usable output for additional 15-20% of cases",
                },
                {
                    "priority": "LOW",
                    "area": "Vowel Mode Configuration",
                    "issue": "Vowel mode hardcoded to Darija",
                    "recommendation": "Make vowel mode configurable for different use cases",
                    "estimated_impact": "Provides flexibility for different user needs",
                },
            ]
        )

        # Sort by priority
        priority_order = {"HIGH": 0, "MEDIUM": 1, "LOW": 2}
        recommendations.sort(key=lambda x: priority_order[x["priority"]])

        self.improvement_areas["recommendations"] = recommendations

        logger.info("\n=== IMPROVEMENT RECOMMENDATIONS ===")
        for i, rec in enumerate(recommendations, 1):
            logger.info(f"\n{i}. [{rec['priority']}] {rec['area']}")
            logger.info(f"   Issue: {rec['issue']}")
            logger.info(f"   Recommendation: {rec['recommendation']}")
            logger.info(f"   Estimated Impact: {rec['estimated_impact']}")

        return recommendations

    def create_priority_improvement_plan(self):
        """Create a structured improvement plan with implementation steps"""
        logger.info("Creating priority improvement plan...")

        recommendations = self.improvement_areas.get("recommendations", [])

        improvement_plan = {
            "phase_1_immediate": {
                "title": "Immediate Critical Fixes",
                "timeline": "1-2 weeks",
                "priority": "HIGH",
                "tasks": [
                    {
                        "task": "Fix character mapping consistency",
                        "description": "Ensure all common Darija characters (7, 3, 9, 8, etc.) map correctly",
                        "files_to_modify": ["src/stages/script_generation.rs"],
                        "estimated_effort": "3-5 days",
                    },
                    {
                        "task": "Improve validation stage tolerance",
                        "description": "Relax validation rules for common Darija patterns",
                        "files_to_modify": ["src/stages/validation.rs"],
                        "estimated_effort": "2-3 days",
                    },
                    {
                        "task": "Expand emphatic consonant lexicon",
                        "description": "Add more emphatic words and improve detection algorithms",
                        "files_to_modify": ["src/linguistic/emphatic.rs"],
                        "estimated_effort": "3-4 days",
                    },
                ],
            },
            "phase_2_short_term": {
                "title": "Short-term Improvements",
                "timeline": "3-4 weeks",
                "priority": "MEDIUM",
                "tasks": [
                    {
                        "task": "Implement better clitic handling",
                        "description": "Improve function word prefix separation",
                        "files_to_modify": ["src/stages/tokenization.rs"],
                        "estimated_effort": "4-5 days",
                    },
                    {
                        "task": "Expand Darija number patterns",
                        "description": "Add common Darija number expressions (7na, 3la, etc.)",
                        "files_to_modify": ["src/linguistic/patterns.rs"],
                        "estimated_effort": "3-4 days",
                    },
                    {
                        "task": "Implement error recovery mechanisms",
                        "description": "Add graceful fallback for failed conversions",
                        "files_to_modify": ["src/stages/mod.rs", "src/lib.rs"],
                        "estimated_effort": "5-6 days",
                    },
                ],
            },
            "phase_3_medium_term": {
                "title": "Medium-term Enhancements",
                "timeline": "2-3 months",
                "priority": "MEDIUM",
                "tasks": [
                    {
                        "task": "Implement configurable vowel modes",
                        "description": "Allow user selection of vowel processing mode",
                        "files_to_modify": [
                            "src/stages/script_generation.rs",
                            "src/lib.rs",
                            "src/python_bindings.rs",
                        ],
                        "estimated_effort": "2-3 weeks",
                    },
                    {
                        "task": "Add context-aware mapping",
                        "description": "Implement position-based and context-aware character mapping",
                        "files_to_modify": ["src/stages/icr.rs"],
                        "estimated_effort": "3-4 weeks",
                    },
                    {
                        "task": "Create comprehensive test suite",
                        "description": "Add extensive testing for edge cases and linguistic patterns",
                        "files_to_modify": ["tests/"],
                        "estimated_effort": "2-3 weeks",
                    },
                ],
            },
            "phase_4_long_term": {
                "title": "Long-term Architecture Improvements",
                "timeline": "3-6 months",
                "priority": "LOW",
                "tasks": [
                    {
                        "task": "Implement plugin architecture for dialects",
                        "description": "Allow easy addition of regional Darija variations",
                        "files_to_modify": ["src/", "src/types/dialect.rs"],
                        "estimated_effort": "2-3 months",
                    },
                    {
                        "task": "Add machine learning component",
                        "description": "Integrate ML model for better context understanding",
                        "files_to_modify": ["src/", "Cargo.toml"],
                        "estimated_effort": "3-4 months",
                    },
                    {
                        "task": "Performance optimization",
                        "description": "Optimize for speed and memory usage",
                        "files_to_modify": ["src/"],
                        "estimated_effort": "1-2 months",
                    },
                ],
            },
        }

        self.improvement_areas["improvement_plan"] = improvement_plan

        logger.info("\n=== IMPROVEMENT PLAN ===")
        for phase, phase_data in improvement_plan.items():
            logger.info(f"\n{phase_data['title']} ({phase_data['timeline']})")
            logger.info(f"Priority: {phase_data['priority']}")

            for i, task in enumerate(phase_data["tasks"], 1):
                logger.info(f"\n  {i}. {task['task']}")
                logger.info(f"     Description: {task['description']}")
                logger.info(f"     Files: {', '.join(task['files_to_modify'])}")
                logger.info(f"     Effort: {task['estimated_effort']}")

        return improvement_plan


def main():
    """Main function to run the improvement analysis"""
    analyzer = CaDaRImprovementAnalyzer("cadar_benchmark_results.csv")

    # Run all analyses
    analyzer.analyze_wer_distribution()
    analyzer.identify_failure_patterns()
    analyzer.analyze_character_mapping_accuracy()
    recommendations = analyzer.generate_improvement_recommendations()
    improvement_plan = analyzer.create_priority_improvement_plan()

    # Save analysis results
    import json

    with open("cadar_improvement_analysis.json", "w", encoding="utf-8") as f:
        json.dump(
            analyzer.improvement_areas, f, ensure_ascii=False, indent=2, default=str
        )

    logger.info(
        f"\nAnalysis complete! Results saved to cadar_improvement_analysis.json"
    )

    # Summary statistics
    total_recommendations = len(recommendations)
    high_priority = len([r for r in recommendations if r["priority"] == "HIGH"])

    logger.info(f"\n=== SUMMARY ===")
    logger.info(f"Total recommendations: {total_recommendations}")
    logger.info(f"High priority: {high_priority}")
    logger.info(f"Current average WER: {analyzer.df['wer_score'].mean():.4f}")
    logger.info(f"Target WER: <0.3 (75%+ improvement)")


if __name__ == "__main__":
    main()
