#!/usr/bin/env python3
"""
Data analysis script for metas_v6_tr.jsonl
Analyzes field distribution and vocal data count in v6 metadata
"""

import json
import random
import argparse
from pathlib import Path
from collections import defaultdict, Counter
import matplotlib.pyplot as plt
from typing import Dict, Any, List
from tqdm import tqdm


def sample_jsonl_lines(
    filepath: Path,
    sample_size: int = None,
    random_seed: int = 42,
    sampling_mode: str = "reservoir",
    skip_lines: int = 0,
) -> List[Dict[str, Any]]:
    """
    Sample lines from a large JSONL file

    Args:
        filepath: Path to JSONL file
        sample_size: Number of samples to collect (None = process all lines)
        random_seed: Random seed for reproducibility
        sampling_mode: "reservoir" (memory efficient), "random" (truly random lines), "all" (process everything), or "skip" (skip lines)
        skip_lines: For skip mode, skip N lines between each processed line (999 = process every 1000th line)
    """
    random.seed(random_seed)
    sampled_lines = []

    if sampling_mode == "all" or sample_size is None:
        # Process all lines without sampling
        with open(filepath, "r") as f:
            for line in tqdm(f, desc="Reading all lines", unit=" lines"):
                sampled_lines.append(json.loads(line.strip()))
    elif sampling_mode == "random":
        # First pass: count total lines
        print("Counting total lines...")
        with open(filepath, "r") as f:
            total_lines = sum(1 for _ in tqdm(f, desc="Counting lines", unit=" lines"))

        # Generate random line numbers to sample
        if sample_size >= total_lines:
            sample_indices = set(range(total_lines))
        else:
            sample_indices = set(random.sample(range(total_lines), sample_size))

        # Second pass: collect sampled lines
        with open(filepath, "r") as f:
            for i, line in enumerate(tqdm(f, desc="Sampling lines", unit=" lines", total=total_lines)):
                if i in sample_indices:
                    sampled_lines.append(json.loads(line.strip()))
                    if len(sampled_lines) >= sample_size:
                        break
    elif sampling_mode == "skip":
        # Skip-line sampling: process every (skip_lines + 1)th line
        with open(filepath, "r") as f:
            line_count = 0
            for line in tqdm(f, desc="Skip-line sampling", unit=" lines"):
                if line_count % (skip_lines + 1) == 0:
                    sampled_lines.append(json.loads(line.strip()))
                    if sample_size and len(sampled_lines) >= sample_size:
                        break
                line_count += 1
    else:
        # Reservoir sampling (original method)
        with open(filepath, "r") as f:
            for i, line in enumerate(tqdm(f, desc="Reading JSONL file", unit=" lines")):
                if len(sampled_lines) < sample_size:
                    sampled_lines.append(json.loads(line.strip()))
                else:
                    # Reservoir sampling: randomly replace elements
                    j = random.randint(0, i)
                    if j < sample_size:
                        sampled_lines[j] = json.loads(line.strip())

    return sampled_lines


def analyze_fields(data: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Analyze all fields found in the data sample
    """
    field_counts = Counter()
    field_types = defaultdict(set)
    field_examples = {}
    vocal_count = 0
    has_stems_count = 0

    for record in tqdm(data, desc="Analyzing records", unit=" records"):
        for field, value in record.items():
            field_counts[field] += 1
            field_types[field].add(type(value).__name__)

            # Store example values (truncate if too long)
            if field not in field_examples:
                if isinstance(value, str) and len(value) > 100:
                    field_examples[field] = value[:100] + "..."
                else:
                    field_examples[field] = value

        # Check for vocal data indicators
        if "stems" in record:
            has_stems_count += 1
            stems = record["stems"]
            if isinstance(stems, dict) and any("vocal" in k.lower() for k in stems.keys()):
                vocal_count += 1

        # Also check tags for vocal indicators
        if "tags" in record and isinstance(record["tags"], list):
            if any("vocal" in str(tag).lower() for tag in record["tags"]):
                vocal_count += 1

    return {
        "field_counts": dict(field_counts),
        "field_types": {k: list(v) for k, v in field_types.items()},
        "field_examples": field_examples,
        "vocal_count": vocal_count,
        "has_stems_count": has_stems_count,
        "total_samples": len(data),
    }


def analyze_stem_names(data: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Analyze all unique stem names found in the stems field
    """
    stem_name_counts = Counter()
    stem_examples = {}
    records_with_stems = 0
    total_stems = 0

    for record in tqdm(data, desc="Analyzing stem names", unit=" records"):
        if "stems" in record and isinstance(record["stems"], dict):
            records_with_stems += 1
            stems = record["stems"]
            total_stems += len(stems)

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

                # Store example record ID for each stem name (first occurrence)
                if stem_name not in stem_examples:
                    stem_examples[stem_name] = {
                        "record_id": record.get("id", "N/A"),
                        "stem_path": stems[stem_name],
                    }

    return {
        "stem_name_counts": dict(stem_name_counts),
        "stem_examples": stem_examples,
        "records_with_stems": records_with_stems,
        "total_stems": total_stems,
        "unique_stem_names": len(stem_name_counts),
        "total_samples": len(data),
    }


def create_field_distribution_chart(field_counts: Dict[str, int], output_path: Path):
    """
    Create bar chart showing field distribution
    """
    fields = list(field_counts.keys())
    counts = list(field_counts.values())

    plt.figure(figsize=(12, 8))
    bars = plt.bar(fields, counts)
    plt.title("Field Distribution in metas_v6_tr.jsonl Sample")
    plt.xlabel("Fields")
    plt.ylabel("Frequency")
    plt.xticks(rotation=45, ha="right")

    # Add count labels on bars
    for bar, count in zip(bars, counts):
        plt.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + max(counts) * 0.01,
            str(count),
            ha="center",
            va="bottom",
        )

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close()


def create_stem_names_chart(stem_name_counts: Dict[str, int], output_path: Path, top_n: int = 20):
    """
    Create bar chart showing most common stem names
    """
    # Get top N most common stem names
    sorted_stems = sorted(stem_name_counts.items(), key=lambda x: x[1], reverse=True)[:top_n]
    stem_names = [item[0] for item in sorted_stems]
    counts = [item[1] for item in sorted_stems]

    plt.figure(figsize=(15, 8))
    bars = plt.bar(stem_names, counts)
    plt.title(f"Top {len(stem_names)} Most Common Stem Names in metas_v6_tr.jsonl Sample")
    plt.xlabel("Stem Names")
    plt.ylabel("Frequency")
    plt.xticks(rotation=45, ha="right")

    # Add count labels on bars
    for bar, count in zip(bars, counts):
        plt.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + max(counts) * 0.01,
            str(count),
            ha="center",
            va="bottom",
        )

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close()


def main():
    parser = argparse.ArgumentParser(description="Analyze v6 metadata fields")
    parser.add_argument(
        "--data-path",
        type=str,
        default="/app2/suno/data/auk_v0/metas_v6_tr.jsonl",
        help="Path to the JSONL file",
    )
    parser.add_argument("--sample-size", type=int, default=10000, help="Number of samples to analyze")
    parser.add_argument("--output-dir", type=str, default=".", help="Output directory for results")
    parser.add_argument("--random-seed", type=int, default=42, help="Random seed for sampling")
    parser.add_argument(
        "--sampling-mode",
        type=str,
        default="reservoir",
        choices=["reservoir", "random", "all", "skip"],
        help="Sampling method: reservoir (efficient), random (truly random lines), all (process everything), or skip (skip lines)",
    )
    parser.add_argument(
        "--skip-lines",
        type=int,
        default=999,
        help="For skip sampling: skip N lines between each processed line (999 = every 1000th line)",
    )
    parser.add_argument(
        "--analyze-stems",
        action="store_true",
        help="Perform stem name analysis instead of general field analysis",
    )

    args = parser.parse_args()

    data_path = Path(args.data_path)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Set default sampling for stem analysis
    if args.analyze_stems and args.sampling_mode == "reservoir":
        print("Stem analysis mode: switching to skip sampling (1 in 1000)")
        args.sampling_mode = "skip"

    if args.sampling_mode == "all":
        print(f"Processing all lines from {data_path}")
        sampled_data = sample_jsonl_lines(
            data_path, None, args.random_seed, args.sampling_mode, args.skip_lines
        )
    elif args.sampling_mode == "skip":
        print(f"Skip sampling from {data_path} (processing every {args.skip_lines + 1} lines)")
        sampled_data = sample_jsonl_lines(
            data_path, args.sample_size, args.random_seed, args.sampling_mode, args.skip_lines
        )
    else:
        print(f"Sampling {args.sample_size} lines from {data_path} using {args.sampling_mode} sampling")
        sampled_data = sample_jsonl_lines(
            data_path, args.sample_size, args.random_seed, args.sampling_mode, args.skip_lines
        )
    print(f"Successfully sampled {len(sampled_data)} records")

    if args.analyze_stems:
        print("Analyzing stem names...")
        stem_results = analyze_stem_names(sampled_data)

        # Save stem analysis results to JSON
        results_path = output_dir / "stem_names_frequency.json"
        with open(results_path, "w") as f:
            json.dump(stem_results, f, indent=2)
        print(f"Stem analysis results saved to {results_path}")

        # Create stem names chart
        chart_path = output_dir / "stem_names_chart.png"
        create_stem_names_chart(stem_results["stem_name_counts"], chart_path)
        print(f"Stem names chart saved to {chart_path}")

        # Print stem analysis summary
        print("\n=== STEM ANALYSIS SUMMARY ===")
        print(f"Total samples analyzed: {stem_results['total_samples']}")
        print(f"Records with stems: {stem_results['records_with_stems']}")
        print(f"Total stems found: {stem_results['total_stems']}")
        print(f"Unique stem names: {stem_results['unique_stem_names']}")
        print(
            f"Average stems per record: {stem_results['total_stems'] / stem_results['records_with_stems']:.2f}"
        )

        # Print all stem names sorted by frequency
        print(f"\nAll {stem_results['unique_stem_names']} unique stem names (sorted by frequency):")
        sorted_stems = sorted(stem_results["stem_name_counts"].items(), key=lambda x: x[1], reverse=True)
        for stem_name, count in sorted_stems:
            print(f"  {stem_name}: {count}")

    else:
        print("Analyzing fields...")
        analysis_results = analyze_fields(sampled_data)

        # Save analysis results to JSON
        results_path = output_dir / "field_analysis_results.json"
        with open(results_path, "w") as f:
            json.dump(analysis_results, f, indent=2)
        print(f"Analysis results saved to {results_path}")

        # Create bar chart
        chart_path = output_dir / "field_distribution_chart.png"
        create_field_distribution_chart(analysis_results["field_counts"], chart_path)
        print(f"Field distribution chart saved to {chart_path}")

        # Print summary
        print("\n=== ANALYSIS SUMMARY ===")
        print(f"Total samples analyzed: {analysis_results['total_samples']}")
        print(f"Records with stems: {analysis_results['has_stems_count']}")
        print(f"Records with vocal indicators: {analysis_results['vocal_count']}")
        print(
            f"Vocal percentage: {analysis_results['vocal_count'] / analysis_results['total_samples'] * 100:.2f}%"
        )

        print("\nFields found:")
        for field, count in sorted(analysis_results["field_counts"].items()):
            types = ", ".join(analysis_results["field_types"][field])
            print(f"  {field}: {count} ({types})")


if __name__ == "__main__":
    main()