#!/usr/bin/env python3
"""
CaDaR Benchmark Script
Evaluate CaDaR performance on Darija dataset from Hugging Face
"""

import pandas as pd
import numpy as np
from datasets import load_dataset
import cadar
from tqdm import tqdm
import re
import os
import logging
from typing import List, Dict, Tuple, Optional

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

# Import jiwer directly for WER calculation
try:
    import jiwer
except ImportError:
    logger.error("jiwer is required. Install with: pip install jiwer")
    raise


class CaDaRBenchmark:
    def __init__(self, hf_token: str):
        self.hf_token = hf_token
        self.cadar = cadar
        self.results = []

    def load_darija_dataset(
        self, split: str = "train", max_samples: int = 1000
    ) -> pd.DataFrame:
        """Load Darija dataset from Hugging Face"""
        logger.info(f"Loading arbml/darija dataset ({split} split)...")

        try:
            dataset = load_dataset("arbml/darija", split=split)

            # Convert to DataFrame - ensure it's a proper DataFrame
            if hasattr(dataset, "to_pandas"):
                df = dataset.to_pandas()
            else:
                # Handle case where dataset is already a DataFrame
                df = pd.DataFrame(dataset)

            # Limit samples if specified
            if max_samples and len(df) > max_samples:
                df = df.sample(n=max_samples, random_state=42)
                logger.info(f"Using {max_samples} random samples")

            logger.info(f"Loaded {len(df)} samples")
            return df

        except Exception as e:
            logger.error(f"Error loading dataset: {e}")
            raise

    def preprocess_text(self, text: str) -> str:
        """Basic text preprocessing"""
        if pd.isna(text):
            return ""

        # Remove extra whitespace
        text = str(text).strip()

        # Remove multiple spaces
        text = re.sub(r"\s+", " ", text)

        return text

    def process_sample(self, original_arabizi: str) -> Optional[Dict]:
        """Process a single sample through CaDaR pipeline"""
        try:
            # Preprocess original text
            if original_arabizi is None:
                return None

            original_arabizi_clean = self.preprocess_text(original_arabizi)

            if not original_arabizi_clean:
                return None

            # Step 1: Arabizi to Arabic (bizi2ara)
            arabic_text = self.cadar.bizi2ara(original_arabizi_clean)

            # Step 2: Arabic back to Arabizi (ara2bizi)
            roundtrip_arabizi = self.cadar.ara2bizi(arabic_text)

            # Calculate WER between original and roundtrip arabizi
            wer_score = self.calculate_wer(original_arabizi_clean, roundtrip_arabizi)

            return {
                "original_arabizi": original_arabizi_clean,
                "arabic": arabic_text,
                "roundtrip_arabizi": roundtrip_arabizi,
                "wer_score": wer_score,
                "original_length": len(original_arabizi_clean.split()),
                "arabic_length": len(arabic_text.split()),
                "roundtrip_length": len(roundtrip_arabizi.split()),
            }

        except Exception as e:
            logger.error(f"Error processing sample '{original_arabizi[:50]}...': {e}")
            return None

    def calculate_wer(self, reference: str, hypothesis: str) -> float:
        """Calculate Word Error Rate using jiwer"""
        # Handle empty cases
        if not reference.strip() and not hypothesis.strip():
            return 0.0
        if not reference.strip():
            return 1.0
        if not hypothesis.strip():
            return 1.0

        # Calculate WER using jiwer
        return jiwer.wer(reference, hypothesis)

    def run_benchmark(
        self, dataset_path: Optional[str] = None, max_samples: int = 1000
    ) -> pd.DataFrame:
        """Run the complete benchmark"""
        logger.info("Starting CaDaR Benchmark...")

        # Load dataset
        if dataset_path and os.path.exists(dataset_path):
            logger.info(f"Loading dataset from local file: {dataset_path}")
            df = pd.read_csv(dataset_path)
        else:
            df = self.load_darija_dataset(max_samples=max_samples)

        # Ensure we have a DataFrame with the 'darija' column
        if "darija" not in df.columns:
            logger.error("Dataset must contain 'darija' column")
            raise ValueError("Dataset must contain 'darija' column")

        # Process samples with progress bar
        logger.info(f"Processing {len(df)} samples...")

        processed_samples = []
        for _, row in tqdm(df.iterrows(), total=len(df), desc="Processing"):
            darija_text = str(row.get("darija", "")).strip()
            if darija_text:  # Only process non-empty text
                result = self.process_sample(darija_text)
                if result:
                    processed_samples.append(result)

        # Create results DataFrame
        results_df = pd.DataFrame(processed_samples)

        # Calculate statistics
        self.calculate_statistics(results_df)

        return results_df

    def calculate_statistics(self, df: pd.DataFrame):
        """Calculate benchmark statistics"""
        logger.info("\n" + "=" * 50)
        logger.info("BENCHMARK STATISTICS")
        logger.info("=" * 50)

        total_samples = len(df)
        avg_wer = df["wer_score"].mean()
        perfect_samples = len(df[df["wer_score"] == 0])

        logger.info(f"Total samples processed: {total_samples}")
        logger.info(f"Average WER: {avg_wer:.4f}")
        logger.info(
            f"Perfect conversions (WER=0): {perfect_samples} ({perfect_samples / total_samples * 100:.1f}%)"
        )

        # WER distribution
        logger.info(f"\nWER Distribution:")
        logger.info(f"  WER = 0: {len(df[df['wer_score'] == 0])} samples")
        logger.info(f"  WER < 0.1: {len(df[df['wer_score'] < 0.1])} samples")
        logger.info(f"  WER < 0.2: {len(df[df['wer_score'] < 0.2])} samples")
        logger.info(f"  WER >= 0.5: {len(df[df['wer_score'] >= 0.5])} samples")

        # Length analysis
        logger.info(f"\nLength Analysis:")
        logger.info(f"  Original avg length: {df['original_length'].mean():.1f} words")
        logger.info(f"  Arabic avg length: {df['arabic_length'].mean():.1f} words")
        logger.info(
            f"  Roundtrip avg length: {df['roundtrip_length'].mean():.1f} words"
        )

        # Worst cases
        worst_cases = df.nlargest(5, "wer_score")
        logger.info(f"\nTop 5 Worst Cases (highest WER):")
        for _, row in worst_cases.iterrows():
            logger.info(f"  WER: {row['wer_score']:.4f}")
            logger.info(f"    Original: {row['original_arabizi']}")
            logger.info(f"    Arabic: {row['arabic']}")
            logger.info(f"    Roundtrip: {row['roundtrip_arabizi']}")
            logger.info("-" * 30)

    def save_results(self, df: pd.DataFrame, output_path: str):
        """Save results to CSV"""
        df.to_csv(output_path, index=False)
        logger.info(f"Results saved to: {output_path}")

    def generate_analysis_report(self, df: pd.DataFrame, report_path: str):
        """Generate detailed analysis report"""
        report = []
        report.append("CaDaR Benchmark Analysis Report")
        report.append("=" * 50)
        report.append(f"Dataset: arbml/darija")
        report.append(f"Total samples: {len(df)}")
        report.append(f"Average WER: {df['wer_score'].mean():.4f}")
        report.append(
            f"Perfect conversions: {len(df[df['wer_score'] == 0])} ({len(df[df['wer_score'] == 0]) / len(df) * 100:.1f}%)"
        )
        report.append("")

        # Analysis by WER ranges
        report.append("Performance by WER Range:")
        report.append("-" * 30)
        for wer_range in [0, 0.1, 0.2, 0.5]:
            count = (
                len(df[df["wer_score"] < wer_range])
                if wer_range > 0
                else len(df[df["wer_score"] == 0])
            )
            report.append(
                f"WER < {wer_range}: {count} samples ({count / len(df) * 100:.1f}%)"
            )

        report.append("")
        report.append("Sample Results:")
        report.append("-" * 30)

        # Show good cases
        good_cases = df[df["wer_score"] == 0].head(3)
        report.append("Good Cases (WER = 0):")
        for _, row in good_cases.iterrows():
            report.append(f"Original: {row['original_arabizi']}")
            report.append(f"Arabic: {row['arabic']}")
            report.append("")

        # Show medium cases
        medium_cases = df[(df["wer_score"] > 0) & (df["wer_score"] < 0.2)].head(3)
        if len(medium_cases) > 0:
            report.append("Medium Cases (0 < WER < 0.2):")
            for _, row in medium_cases.iterrows():
                report.append(f"WER: {row['wer_score']:.4f}")
                report.append(f"Original: {row['original_arabizi']}")
                report.append(f"Arabic: {row['arabic']}")
                report.append(f"Roundtrip: {row['roundtrip_arabizi']}")
                report.append("")

        # Show worst cases
        worst_cases = df.nlargest(3, "wer_score")
        report.append("Worst Cases (highest WER):")
        for _, row in worst_cases.iterrows():
            report.append(f"WER: {row['wer_score']:.4f}")
            report.append(f"Original: {row['original_arabizi']}")
            report.append(f"Arabic: {row['arabic']}")
            report.append(f"Roundtrip: {row['roundtrip_arabizi']}")
            report.append("")

        # Write report
        with open(report_path, "w", encoding="utf-8") as f:
            f.write("\n".join(report))

        logger.info(f"Analysis report saved to: {report_path}")


def main():
    """Main function to run the benchmark"""
    # Hugging Face token
    HF_TOKEN = "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"

    # Initialize benchmark
    benchmark = CaDaRBenchmark(HF_TOKEN)

    # Run benchmark with larger sample for better analysis
    logger.info("Running benchmark with 500 samples for analysis...")
    results_df = benchmark.run_benchmark(max_samples=500)

    # Save results
    benchmark.save_results(results_df, "cadar_benchmark_results.csv")

    # Generate analysis report
    benchmark.generate_analysis_report(results_df, "cadar_benchmark_report.txt")

    logger.info("Benchmark completed successfully!")


if __name__ == "__main__":
    main()
