"""
Data Explorer
=============

This module provides tools for exploring and understanding the Sentiment140 dataset.
"""

import json
import logging
from pathlib import Path
from typing import Any, Dict, Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class DataExplorer:
    """
    Explores and analyzes the Sentiment140 dataset to understand its characteristics.
    """

    LABEL_MAP = {0: "negative", 2: "neutral", 4: "positive"}

    def __init__(self, data_path: Path):
        """
        Initialize the data explorer.

        Args:
            data_path: Path to the dataset CSV file
        """
        self.data_path = data_path
        self.df: Optional[pd.DataFrame] = None

    def load_data(self) -> pd.DataFrame:
        """Load the Sentiment140 dataset."""
        logger.info(f"Loading data from {self.data_path}")

        colnames = ["target", "ids", "date", "flag", "user", "text"]
        self.df = pd.read_csv(self.data_path, encoding="latin-1", names=colnames)
        self.df = self.df[~self.df["text"].isna()].reset_index(drop=True)
        self.df["target"] = self.df["target"].astype(int)
        self.df["label_name"] = self.df["target"].map(self.LABEL_MAP)

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

    def get_basic_statistics(self) -> Dict[str, Any]:
        """Get basic statistics about the dataset."""
        if self.df is None:
            raise ValueError("Data not loaded. Call load_data() first.")

        # Text length statistics
        self.df["text_length"] = self.df["text"].str.len()
        self.df["word_count"] = self.df["text"].str.split().str.len()

        stats = {
            "total_samples": len(self.df),
            "columns": list(self.df.columns),
            "class_distribution": self.df["target"].value_counts().to_dict(),
            "class_distribution_pct": (
                self.df["target"].value_counts(normalize=True) * 100
            ).to_dict(),
            "text_length": {
                "mean": float(self.df["text_length"].mean()),
                "median": float(self.df["text_length"].median()),
                "min": int(self.df["text_length"].min()),
                "max": int(self.df["text_length"].max()),
                "std": float(self.df["text_length"].std()),
            },
            "word_count": {
                "mean": float(self.df["word_count"].mean()),
                "median": float(self.df["word_count"].median()),
                "min": int(self.df["word_count"].min()),
                "max": int(self.df["word_count"].max()),
                "std": float(self.df["word_count"].std()),
            },
            "missing_values": self.df.isnull().sum().to_dict(),
            "unique_users": int(self.df["user"].nunique()),
        }

        return stats

    def get_sample_texts(self, n_per_class: int = 5) -> Dict[str, list]:
        """Get sample texts for each class."""
        if self.df is None:
            raise ValueError("Data not loaded. Call load_data() first.")

        samples = {}
        for target_val, label_name in self.LABEL_MAP.items():
            if target_val in self.df["target"].unique():
                class_samples = (
                    self.df[self.df["target"] == target_val]["text"]
                    .head(n_per_class)
                    .tolist()
                )
                samples[label_name] = class_samples

        return samples

    def check_data_balance(self) -> Dict[str, Any]:
        """Check if the dataset is balanced across classes."""
        if self.df is None:
            raise ValueError("Data not loaded. Call load_data() first.")

        class_counts = self.df["target"].value_counts()
        total = len(self.df)

        balance_info = {
            "is_balanced": False,
            "class_counts": class_counts.to_dict(),
            "class_percentages": (class_counts / total * 100).to_dict(),
            "imbalance_ratio": (
                float(class_counts.max() / class_counts.min())
                if class_counts.min() > 0
                else None
            ),
        }

        # Consider balanced if ratio is less than 1.5
        if balance_info["imbalance_ratio"] and balance_info["imbalance_ratio"] < 1.5:
            balance_info["is_balanced"] = True

        return balance_info

    def generate_report(self, output_path: Optional[Path] = None) -> str:
        """
        Generate a comprehensive data understanding report.

        Args:
            output_path: Optional path to save the report

        Returns:
            Report as a string
        """
        if self.df is None:
            self.load_data()

        stats = self.get_basic_statistics()
        balance = self.check_data_balance()
        samples = self.get_sample_texts(n_per_class=3)

        report = f"""
# Data Understanding Report
## Sentiment140 Dataset

### Dataset Overview
- Total Samples: {stats['total_samples']:,}
- Unique Users: {stats['unique_users']:,}
- Columns: {', '.join(stats['columns'])}

### Class Distribution
"""

        for target, count in stats["class_distribution"].items():
            label = self.LABEL_MAP.get(target, str(target))
            pct = stats["class_distribution_pct"][target]
            report += (
                f"- {label.capitalize()} ({target}): {count:,} samples ({pct:.2f}%)\n"
            )

        # Format imbalance ratio
        imbalance_str = (
            f"{balance['imbalance_ratio']:.2f}" if balance["imbalance_ratio"] else "N/A"
        )

        report += f"""
### Data Balance
- Balanced: {'Yes' if balance['is_balanced'] else 'No'}
- Imbalance Ratio: {imbalance_str}

### Text Statistics
#### Length (characters)
- Mean: {stats['text_length']['mean']:.1f}
- Median: {stats['text_length']['median']:.1f}
- Range: [{stats['text_length']['min']}, {stats['text_length']['max']}]
- Std Dev: {stats['text_length']['std']:.1f}

#### Word Count
- Mean: {stats['word_count']['mean']:.1f}
- Median: {stats['word_count']['median']:.1f}
- Range: [{stats['word_count']['min']}, {stats['word_count']['max']}]
- Std Dev: {stats['word_count']['std']:.1f}

### Sample Texts
"""

        for label, texts in samples.items():
            report += f"\n#### {label.capitalize()} Examples:\n"
            for i, text in enumerate(texts, 1):
                report += f"{i}. {text[:100]}{'...' if len(text) > 100 else ''}\n"

        if output_path:
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(report)
            logger.info(f"Report saved to {output_path}")

        return report

    def save_statistics(self, output_path: Path):
        """Save statistics to JSON file."""
        stats = self.get_basic_statistics()

        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(stats, f, indent=2, ensure_ascii=False)

        logger.info(f"Statistics saved to {output_path}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    # Example usage
    data_path = Path("data/sentiment140/train.csv")
    explorer = DataExplorer(data_path)
    explorer.load_data()

    print(explorer.generate_report())
