#!/usr/bin/env python3
"""Compute corpus statistics and validation metrics.

Reads corpus files and computes counts, distributions, coverage, and balance warnings.

Usage:
  python scripts/compute_corpus_stats.py
  python scripts/compute_corpus_stats.py --corpus-dir corpus/
"""

import argparse
import csv
import json
import sys
from collections import Counter
from pathlib import Path

SCRIPT_DIR = Path(__file__).parent
PROJECT_DIR = SCRIPT_DIR.parent
DATA_DIR = PROJECT_DIR / "data"


def load_jsonl(path):
    """Load a JSONL file."""
    entries = []
    if not path.exists():
        return entries
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                entries.append(json.loads(line))
    return entries


def classify_input_type(inp):
    """Classify the input framing type of a training pair."""
    if inp.startswith("Tell me something about"):
        return "word_seeded"
    elif inp.startswith("Tell me a saying about"):
        return "category_seeded"
    elif inp.startswith("What would a"):
        return "persona_seeded"
    elif inp.startswith("Give me a") and "proverb" in inp:
        return "template_seeded"
    elif any(inp.startswith(p) for p in [
        "Tell me some folk", "What do they", "Give me a proverb",
        "Share some", "What's a good"
    ]):
        return "open_ended"
    else:
        return "fictional"


def main():
    parser = argparse.ArgumentParser(description="Compute corpus statistics.")
    parser.add_argument("--corpus-dir", default=str(PROJECT_DIR / "corpus"),
                        help="Corpus directory")
    parser.add_argument("--output", default=None,
                        help="Output JSON file (default: corpus_dir/corpus_stats.json)")
    args = parser.parse_args()

    corpus_dir = Path(args.corpus_dir)
    output_path = Path(args.output) if args.output else corpus_dir / "corpus_stats.json"

    # Load all corpus files
    raw = load_jsonl(corpus_dir / "corpus_raw.jsonl")
    polished = load_jsonl(corpus_dir / "corpus_polished.jsonl")
    filtered = load_jsonl(corpus_dir / "corpus_filtered.jsonl")
    training = load_jsonl(corpus_dir / "training_pairs.jsonl")

    # Load vocab for coverage analysis
    vocab_words = set()
    vocab_path = DATA_DIR / "folksy_vocab.csv"
    if vocab_path.exists():
        with open(vocab_path, newline="", encoding="utf-8") as f:
            for row in csv.DictReader(f):
                vocab_words.add(row["word"])

    stats = {}

    # --- Raw corpus stats ---
    stats["raw_count"] = len(raw)
    raw_by_template = Counter(e.get("meta_template", "unknown") for e in raw)
    stats["raw_by_template"] = dict(sorted(raw_by_template.items()))

    # --- Polish stats ---
    polished_entries = [e for e in polished if e.get("status") == "polished"]
    discarded_entries = [e for e in polished if e.get("status") == "discarded"]
    error_entries = [e for e in polished if e.get("status") == "error"]

    stats["polished_count"] = len(polished_entries)
    stats["discarded_during_polish"] = len(discarded_entries)
    stats["errors_during_polish"] = len(error_entries)
    if polished_entries or discarded_entries:
        total_processed = len(polished_entries) + len(discarded_entries)
        stats["polish_discard_rate"] = f"{len(discarded_entries)/total_processed*100:.1f}%"

    polish_by_template = Counter(e.get("meta_template", "unknown") for e in polished_entries)
    stats["polished_by_template"] = dict(sorted(polish_by_template.items()))

    discard_by_template = Counter(e.get("meta_template", "unknown") for e in discarded_entries)
    stats["discarded_by_template"] = dict(sorted(discard_by_template.items()))

    # --- Filter stats ---
    stats["filtered_count"] = len(filtered)

    filter_by_template = Counter(e.get("meta_template", "unknown") for e in filtered)
    stats["filtered_by_template"] = dict(sorted(filter_by_template.items()))

    # Filter discard count
    stats["discarded_during_filter"] = len(polished_entries) - len(filtered)

    # --- Training pairs stats ---
    stats["training_pair_count"] = len(training)

    training_by_template = Counter(e.get("meta_template", "unknown") for e in training)
    stats["training_by_template"] = dict(sorted(training_by_template.items()))

    input_type_counts = Counter(classify_input_type(e.get("input", "")) for e in training)
    stats["training_by_input_type"] = dict(sorted(input_type_counts.items()))

    # --- Coverage analysis ---
    used_words = set()
    for entry in filtered:
        slots = entry.get("slots", {})
        for v in slots.values():
            word = v.lower().replace(" ", "_")
            if word in vocab_words:
                used_words.add(word)

    stats["unique_slot_words_used"] = len(used_words)
    stats["total_vocab_words"] = len(vocab_words)
    stats["vocab_coverage"] = f"{len(used_words)/len(vocab_words)*100:.1f}%" if vocab_words else "N/A"

    never_used = sorted(vocab_words - used_words)
    stats["words_never_used"] = never_used
    stats["words_never_used_count"] = len(never_used)

    # --- Saying length stats ---
    lengths = []
    for entry in filtered:
        text = entry.get("polished_text", "")
        if text:
            lengths.append(len(text.split()))

    if lengths:
        stats["avg_saying_length_words"] = round(sum(lengths) / len(lengths), 1)
        stats["min_saying_length_words"] = min(lengths)
        stats["max_saying_length_words"] = max(lengths)

    # --- Balance warnings ---
    warnings = []
    if filtered:
        total_filtered = len(filtered)
        for template, count in filter_by_template.items():
            pct = count / total_filtered * 100
            if pct < 10:
                warnings.append(
                    f"WARNING: {template} has only {count} entries ({pct:.1f}%) — "
                    f"below 10% threshold. Generate more raw sayings for this family."
                )

    if training:
        total_training = len(training)
        for template, count in training_by_template.items():
            pct = count / total_training * 100
            if pct < 5:
                warnings.append(
                    f"WARNING: {template} has only {count} training pairs ({pct:.1f}%) — very underrepresented."
                )

    stats["balance_warnings"] = warnings

    # --- Write output ---
    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)

    # --- Print summary ---
    print("=" * 60)
    print("CORPUS STATISTICS")
    print("=" * 60)

    print(f"\nRaw sayings:           {stats['raw_count']}")
    print(f"Polished sayings:      {stats['polished_count']}")
    print(f"Discarded (polish):    {stats.get('discarded_during_polish', 0)} ({stats.get('polish_discard_rate', 'N/A')})")
    print(f"Discarded (filter):    {stats.get('discarded_during_filter', 0)}")
    print(f"Final filtered:        {stats['filtered_count']}")
    print(f"Training pairs:        {stats['training_pair_count']}")

    print(f"\nDistribution by meta-template (filtered):")
    for t, c in sorted(filter_by_template.items()):
        pct = c / len(filtered) * 100 if filtered else 0
        print(f"  {t:30s} {c:5d} ({pct:5.1f}%)")

    print(f"\nDistribution by input framing type:")
    for t, c in sorted(input_type_counts.items()):
        print(f"  {t:20s} {c:5d}")

    print(f"\nVocab coverage: {stats['vocab_coverage']} ({stats['unique_slot_words_used']}/{stats['total_vocab_words']})")
    print(f"Average saying length: {stats.get('avg_saying_length_words', 'N/A')} words")

    if warnings:
        print(f"\nBalance warnings:")
        for w in warnings:
            print(f"  {w}")

    print(f"\nFull stats: {output_path}")


if __name__ == "__main__":
    main()
