"""Analyzer for stems and stems_captions fields."""

from collections import Counter, defaultdict
from typing import List, Dict, Any
import logging

from .base_analyzer import BaseAnalyzer

logger = logging.getLogger(__name__)


class StemsAnalyzer(BaseAnalyzer):
    """Analyze stems and stems_captions distribution."""

    def __init__(self, output_dir):
        super().__init__("stems", output_dir)

    def process_chunk(self, records: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Process a chunk of records for stems analysis."""
        stem_counts = Counter()
        stem_names = Counter()
        stems_captions_counts = defaultdict(int)
        records_with_stems = 0
        records_with_stems_captions = 0
        stem_count_distribution = []

        for record in records:
            # Analyze stems
            if "stems" in record and record["stems"]:
                records_with_stems += 1
                stems = record["stems"]
                stem_count_distribution.append(len(stems))

                for stem_name in stems.keys():
                    stem_names[stem_name] += 1

            # Analyze stems_captions
            if "stems_captions" in record and record["stems_captions"]:
                records_with_stems_captions += 1
                stems_captions = record["stems_captions"]

                for stem_name, captions in stems_captions.items():
                    stems_captions_counts[stem_name] += (
                        len(captions) if isinstance(captions, list) else 1
                    )

        return {
            "stem_names": dict(stem_names),
            "stems_captions_counts": dict(stems_captions_counts),
            "records_with_stems": records_with_stems,
            "records_with_stems_captions": records_with_stems_captions,
            "stem_count_distribution": stem_count_distribution,
            "total_records": len(records),
        }

    def aggregate(self, chunk_results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Aggregate results from all chunks."""
        from ..utils.stats_utils import aggregate_counters

        all_stem_names = []
        all_stems_captions = []
        all_stem_counts = []
        total_with_stems = 0
        total_with_captions = 0
        total_records = 0

        for chunk in chunk_results:
            all_stem_names.append(chunk.get("stem_names", {}))
            all_stems_captions.append(chunk.get("stems_captions_counts", {}))
            all_stem_counts.extend(chunk.get("stem_count_distribution", []))
            total_with_stems += chunk.get("records_with_stems", 0)
            total_with_captions += chunk.get("records_with_stems_captions", 0)
            total_records += chunk.get("total_records", 0)

        # Aggregate
        total_stem_names = aggregate_counters(all_stem_names)
        total_stems_captions = aggregate_counters(all_stems_captions)

        # Sort by frequency and keep top 200
        sorted_stem_names = sorted(total_stem_names.items(), key=lambda x: x[1], reverse=True)[:200]
        sorted_captions = sorted(total_stems_captions.items(), key=lambda x: x[1], reverse=True)[:200]

        # Calculate distribution
        stem_count_dist = self.calculate_distribution(all_stem_counts) if all_stem_counts else {}

        return {
            "summary": {
                "total_records": total_records,
                "records_with_stems": total_with_stems,
                "records_with_stems_captions": total_with_captions,
                "unique_stem_names": len(total_stem_names),
                "unique_stems_captions": len(total_stems_captions),
            },
            "stem_names_frequency": sorted_stem_names,
            "stems_captions_frequency": sorted_captions,
            "stem_count_distribution": stem_count_dist,
        }