#!/usr/bin/env python3
"""
Analyze validation metadata and generate summary statistics.

This script analyzes the filtered validation metadata to extract:
- Duration distribution
- Tag distribution (number of tags per entry, top tags, excluding tags with ":")
- Language distribution
- Text/lyrics availability and alignment information
"""

import json
import os
from collections import Counter, defaultdict
from pathlib import Path
from typing import Dict, List, Tuple

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(iterable, *args, **kwargs):
        return iterable


def analyze_metadata(jsonl_path: str) -> Dict:
    """
    Analyze metadata JSONL file and collect statistics.
    
    Args:
        jsonl_path: Path to the metadata JSONL file
        
    Returns:
        Dictionary containing all collected statistics
    """
    print(f"Analyzing metadata from: {jsonl_path}")
    
    # Statistics collectors
    durations = []
    all_tags = []
    tags_per_entry = []
    languages = Counter()
    
    # Text/lyrics statistics
    has_text = 0
    no_text = 0
    has_text_aligned = 0
    has_text_no_aligned = 0
    
    # Count total lines
    print("Counting lines...")
    with open(jsonl_path, "r") as f:
        total_lines = sum(1 for _ in f)
    print(f"Total entries: {total_lines:,}")
    
    # Process each entry
    print("Processing entries...")
    with open(jsonl_path, "r") as f:
        for line in tqdm(f, total=total_lines, desc="Analyzing"):
            line = line.strip()
            if not line:
                continue
            
            try:
                meta = json.loads(line)
                
                # Duration
                duration = meta.get("duration_s")
                if duration is not None:
                    durations.append(duration)
                
                # Tags - filter out tags containing ":"
                tags = meta.get("tags", [])
                if tags:
                    filtered_tags = [tag for tag in tags if ":" not in str(tag)]
                    all_tags.extend(filtered_tags)
                    tags_per_entry.append(len(filtered_tags))
                else:
                    tags_per_entry.append(0)
                
                # Language
                lang = meta.get("lang")
                if lang:
                    languages[lang] += 1
                
                # Text/lyrics analysis
                text = meta.get("text")
                text_aligned = meta.get("text_aligned")
                
                if text:
                    has_text += 1
                    if text_aligned:
                        has_text_aligned += 1
                    else:
                        has_text_no_aligned += 1
                else:
                    no_text += 1
                    
            except json.JSONDecodeError as e:
                print(f"Warning: Skipping malformed JSON line: {e}")
                continue
    
    # Compile statistics
    stats = {
        "total_entries": total_lines,
        "durations": durations,
        "all_tags": all_tags,
        "tags_per_entry": tags_per_entry,
        "languages": languages,
        "text_stats": {
            "has_text": has_text,
            "no_text": no_text,
            "has_text_aligned": has_text_aligned,
            "has_text_no_aligned": has_text_no_aligned,
        }
    }
    
    return stats


def calculate_duration_stats(durations: List[float]) -> Dict:
    """Calculate duration distribution statistics."""
    if not durations:
        return {}
    
    sorted_durations = sorted(durations)
    n = len(sorted_durations)
    
    return {
        "count": n,
        "min": min(durations),
        "max": max(durations),
        "mean": sum(durations) / n,
        "median": sorted_durations[n // 2],
        "p25": sorted_durations[int(n * 0.25)],
        "p75": sorted_durations[int(n * 0.75)],
        "p95": sorted_durations[int(n * 0.95)],
        "p99": sorted_durations[int(n * 0.99)],
    }


def calculate_tag_stats(all_tags: List[str], tags_per_entry: List[int]) -> Dict:
    """Calculate tag distribution statistics."""
    tag_counts = Counter(all_tags)
    top_10 = tag_counts.most_common(10)
    
    tags_per_entry_dist = Counter(tags_per_entry)
    
    avg_tags = sum(tags_per_entry) / len(tags_per_entry) if tags_per_entry else 0
    
    return {
        "total_tags": len(all_tags),
        "unique_tags": len(tag_counts),
        "avg_tags_per_entry": avg_tags,
        "top_10": top_10,
        "tags_per_entry_distribution": dict(sorted(tags_per_entry_dist.items())),
    }


def write_summary(stats: Dict, output_path: str) -> None:
    """
    Write analysis summary to file.
    
    Args:
        stats: Statistics dictionary
        output_path: Path to write summary file
    """
    print(f"\nWriting summary to: {output_path}")
    
    with open(output_path, "w") as f:
        f.write("=" * 80 + "\n")
        f.write("VALIDATION METADATA ANALYSIS SUMMARY\n")
        f.write("=" * 80 + "\n\n")
        
        # Overall statistics
        f.write(f"Total Entries: {stats['total_entries']:,}\n\n")
        
        # Duration statistics
        f.write("-" * 80 + "\n")
        f.write("DURATION STATISTICS (seconds)\n")
        f.write("-" * 80 + "\n")
        duration_stats = calculate_duration_stats(stats['durations'])
        if duration_stats:
            f.write(f"Count:   {duration_stats['count']:,}\n")
            f.write(f"Min:     {duration_stats['min']:.2f}s\n")
            f.write(f"Max:     {duration_stats['max']:.2f}s\n")
            f.write(f"Mean:    {duration_stats['mean']:.2f}s\n")
            f.write(f"Median:  {duration_stats['median']:.2f}s\n")
            f.write(f"P25:     {duration_stats['p25']:.2f}s\n")
            f.write(f"P75:     {duration_stats['p75']:.2f}s\n")
            f.write(f"P95:     {duration_stats['p95']:.2f}s\n")
            f.write(f"P99:     {duration_stats['p99']:.2f}s\n")
        f.write("\n")
        
        # Tag statistics
        f.write("-" * 80 + "\n")
        f.write("TAG STATISTICS (excluding tags with ':')\n")
        f.write("-" * 80 + "\n")
        tag_stats = calculate_tag_stats(stats['all_tags'], stats['tags_per_entry'])
        f.write(f"Total tags (all entries):    {tag_stats['total_tags']:,}\n")
        f.write(f"Unique tags:                 {tag_stats['unique_tags']:,}\n")
        f.write(f"Average tags per entry:      {tag_stats['avg_tags_per_entry']:.2f}\n\n")
        
        f.write("Top 10 Tags:\n")
        for i, (tag, count) in enumerate(tag_stats['top_10'], 1):
            f.write(f"  {i:2d}. {tag:<40s} {count:>10,} occurrences\n")
        
        f.write("\nTags Per Entry Distribution:\n")
        for num_tags, count in sorted(tag_stats['tags_per_entry_distribution'].items())[:20]:
            pct = 100 * count / stats['total_entries']
            f.write(f"  {num_tags:3d} tags: {count:>8,} entries ({pct:>5.2f}%)\n")
        f.write("\n")
        
        # Language statistics
        f.write("-" * 80 + "\n")
        f.write("LANGUAGE DISTRIBUTION\n")
        f.write("-" * 80 + "\n")
        total_with_lang = sum(stats['languages'].values())
        f.write(f"Entries with language field: {total_with_lang:,}\n\n")
        
        for lang, count in stats['languages'].most_common():
            pct = 100 * count / total_with_lang if total_with_lang > 0 else 0
            f.write(f"  {lang:<10s} {count:>10,} entries ({pct:>5.2f}%)\n")
        f.write("\n")
        
        # Text/lyrics statistics
        f.write("-" * 80 + "\n")
        f.write("TEXT/LYRICS STATISTICS\n")
        f.write("-" * 80 + "\n")
        text_stats = stats['text_stats']
        total = text_stats['has_text'] + text_stats['no_text']
        
        f.write(f"Has text (lyrics):           {text_stats['has_text']:>10,} "
                f"({100 * text_stats['has_text'] / total:>5.2f}%)\n")
        f.write(f"No text:                     {text_stats['no_text']:>10,} "
                f"({100 * text_stats['no_text'] / total:>5.2f}%)\n\n")
        
        if text_stats['has_text'] > 0:
            f.write("Of entries with text:\n")
            f.write(f"  Has text_aligned:          {text_stats['has_text_aligned']:>10,} "
                    f"({100 * text_stats['has_text_aligned'] / text_stats['has_text']:>5.2f}%)\n")
            f.write(f"  No text_aligned:           {text_stats['has_text_no_aligned']:>10,} "
                    f"({100 * text_stats['has_text_no_aligned'] / text_stats['has_text']:>5.2f}%)\n")
        
        f.write("\n")
        f.write("=" * 80 + "\n")
    
    print("Summary written successfully!")


def main():
    """Main function to run analysis."""
    # Configuration
    metadata_path = "/home/tony/Work/tony/RealGen/metas_v5_val_filtered_clean.jsonl"
    output_dir = "/home/tony/Work/tony/RealGen"
    output_filename = "validation_metadata_clean_analysis_summary.txt"
    output_path = os.path.join(output_dir, output_filename)
    
    print("🎵 Validation Metadata Analysis Script")
    print(f"Input:  {metadata_path}")
    print(f"Output: {output_path}")
    print()
    
    # Check if input file exists
    if not os.path.exists(metadata_path):
        raise FileNotFoundError(f"Metadata file not found: {metadata_path}")
    
    # Analyze metadata
    stats = analyze_metadata(metadata_path)
    
    # Write summary
    write_summary(stats, output_path)
    
    # Also print summary to console
    print("\n" + "=" * 80)
    print("QUICK SUMMARY")
    print("=" * 80)
    print(f"Total entries:         {stats['total_entries']:,}")
    print(f"Entries with duration: {len(stats['durations']):,}")
    print(f"Unique tags:           {len(Counter(stats['all_tags'])):,}")
    print(f"Languages found:       {len(stats['languages'])}")
    print(f"Entries with text:     {stats['text_stats']['has_text']:,}")
    print("=" * 80)
    
    print("\n✅ Analysis completed successfully!")


if __name__ == "__main__":
    main()