theme_generation/corpus.py

"""Corpus loading and descriptive statistics."""

from __future__ import annotations

import sqlite3
from collections import Counter
from pathlib import Path
from typing import Any

from .common import Symbol, parse_key_root


def load_sequences(db_path: Path, allowed_durations: set[str], min_len: int) -> list[tuple[Symbol, ...]]:
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    themes = conn.execute(
        """
        SELECT id, abc_key
        FROM themes
        WHERE parse_error IS NULL
        ORDER BY id
        """
    ).fetchall()

    sequences = []
    for theme in themes:
        root = parse_key_root(theme["abc_key"])
        if root is None:
            continue
        rows = conn.execute(
            """
            SELECT pitch_class, duration_value
            FROM notes
            WHERE theme_id = ?
            ORDER BY start_tick, note_index
            """,
            (theme["id"],),
        ).fetchall()
        seq = []
        for row in rows:
            duration = row["duration_value"]
            if duration not in allowed_durations:
                continue
            seq.append(Symbol((row["pitch_class"] - root) % 12, duration))
        if len(seq) >= min_len:
            sequences.append(tuple(seq))
    conn.close()
    return sequences


def endpoint_priors(db_path: Path, smoothing: float = 0.5) -> tuple[dict[int, float], dict[int, float]]:
    conn = sqlite3.connect(db_path)
    first_counts = Counter()
    last_counts = Counter()
    for first, last in conn.execute(
        "SELECT salient_degree, last_degree FROM endpoint_analysis "
        "WHERE salient_degree IS NOT NULL AND last_degree IS NOT NULL"
    ):
        first_counts[int(first)] += 1
        last_counts[int(last)] += 1
    conn.close()

    def weights(counter: Counter[int]) -> dict[int, float]:
        total = sum(counter.values()) + smoothing * 12
        probs = {degree: (counter[degree] + smoothing) / total for degree in range(12)}
        mean = sum(probs.values()) / 12
        return {degree: probs[degree] / mean for degree in range(12)}

    return weights(first_counts), weights(last_counts)


def symbol_stats(sequences: list[tuple[Symbol, ...]]) -> dict[str, Any]:
    lengths = [len(seq) for seq in sequences]
    vocab = Counter(symbol for seq in sequences for symbol in seq)
    durations = Counter(symbol.duration for symbol, count in vocab.items() for _ in range(count))
    return {
        "sequence_count": len(sequences),
        "event_count": sum(lengths),
        "vocab_size": len(vocab),
        "length_min": min(lengths),
        "length_max": max(lengths),
        "length_mean": sum(lengths) / len(lengths),
        "top_symbols": vocab.most_common(12),
        "top_durations": durations.most_common(12),
    }