#!/usr/bin/env python3
"""
Stylometric self-test for Staying Dark.

Runs the same kind of analysis that was used to unmask Elena Ferrante
(Juola, Eder) and to test the Robin-Gunningham-is-Banksy hypothesis
(geographic profiling is different but the principle is the same:
measure distinctive features in an unknown corpus against a known
candidate corpus).

This script measures the manuscript against a corpus of your own
public writing. If the manuscript matches your public style closely,
your anonymity has a stylometric weakness. If the manuscript matches
poorly, your anonymity has stylometric cover.

Method: Burrows's Delta on the 100 most common function words.

Burrows's Delta is the standard stylometric measure. Lower values mean
more similar style. The threshold above which two corpora are considered
"different authors" varies by genre, but for prose nonfiction:

    Delta < 0.7   strong stylistic match (anonymity weak)
    0.7-1.0       moderate match (some risk)
    1.0-1.3       weak match (typical inter-author variance)
    > 1.3         strong differentiation (anonymity strong)

These thresholds are calibrated for English prose nonfiction. Treat them
as rough guidance, not as bright lines.

Usage:
    python3 stylometric-test.py \\
        --manuscript your-manuscript.docx \\
        --corpus my-public-writing-dir/

The corpus directory should contain plain text files (.txt or .md) of
your other writing: blog posts, LinkedIn essays, articles under your real
name, anything that an external party could use as a comparison corpus.
The more, the better; aim for at least 20,000 words of public writing
to get a stable comparison.

Outputs:
    Burrows's Delta score
    Top 20 over-represented function words in the manuscript
    Top 20 over-represented function words in the corpus
    Sentence-length comparison
    Lexical richness comparison
    Risk assessment

No external dependencies beyond Python standard library.
"""

from __future__ import annotations

import argparse
import collections
import math
import re
import sys
import zipfile
from pathlib import Path
from typing import Iterable

# The 100 most common English function words. Burrows recommended 40-150;
# 100 is a common standard. These are the words a writer can't avoid using
# and which carry the most reliable stylistic signature.
FUNCTION_WORDS = {
    "the", "of", "and", "to", "a", "in", "that", "is", "was", "it",
    "for", "with", "as", "his", "on", "be", "at", "by", "i", "this",
    "had", "not", "are", "but", "from", "or", "have", "an", "they", "which",
    "one", "you", "were", "her", "all", "she", "there", "would", "their", "we",
    "him", "been", "has", "when", "who", "will", "more", "no", "if", "out",
    "so", "said", "what", "up", "its", "about", "into", "than", "them", "can",
    "only", "other", "new", "some", "could", "time", "these", "two", "may",
    "then", "do", "first", "any", "my", "now", "such", "like", "our", "over",
    "man", "me", "even", "most", "made", "after", "also", "did", "many",
    "before", "must", "through", "back", "years", "where", "much", "your",
    "way", "well", "down", "should",
}


MIN_CORPUS_WORDS = 1000


def read_docx(path: Path) -> str:
    """Extract plain text from a .docx file."""
    text_parts = []
    with zipfile.ZipFile(path, "r") as z:
        try:
            doc_xml = z.read("word/document.xml").decode("utf-8")
        except KeyError:
            print(f"WARNING: {path} doesn't look like a Word doc")
            return ""
    # Strip XML tags, keep text content
    text = re.sub(r"<w:p[^>]*>", "\n\n", doc_xml)
    text = re.sub(r"<w:tab[^>]*/>", "\t", text)
    text = re.sub(r"<w:br[^>]*/>", "\n", text)
    text = re.sub(r"<[^>]+>", "", text)
    return text


def read_text_file(path: Path) -> str:
    """Read a plain text or markdown file."""
    return path.read_text(encoding="utf-8", errors="replace")


SUPPORTED = (".txt", ".md", ".docx")


def read_corpus(path: Path) -> str:
    """Read a single file (any supported format) or all files in a directory.

    Reports what it read and what it skipped, so a directory of unsupported
    files (PDFs are the usual case) is visible rather than silent.
    """
    if path.is_dir():
        chunks, used, skipped = [], [], collections.Counter()
        for f in sorted(path.iterdir()):
            if not f.is_file():
                continue
            if f.suffix.lower() in SUPPORTED:
                chunks.append(read_one(f))
                used.append(f.name)
            else:
                skipped[f.suffix.lower() or "(no extension)"] += 1
        print(f"  Read {len(used)} file(s) from {path}")
        if skipped:
            detail = ", ".join(f"{n} x {ext}" for ext, n in sorted(skipped.items()))
            print(f"  Skipped {sum(skipped.values())} unsupported file(s): {detail}")
            print(f"  Supported types: {', '.join(SUPPORTED)}")
        return "\n\n".join(chunks)
    return read_one(path)


def read_one(path: Path) -> str:
    if path.suffix.lower() == ".docx":
        return read_docx(path)
    return read_text_file(path)


def tokenize(text: str) -> list[str]:
    """Lowercase, strip punctuation, return word tokens."""
    text = text.lower()
    # Replace smart quotes with straight
    text = text.replace("’", "'").replace("‘", "'")
    text = text.replace("“", '"').replace("”", '"')
    # Word characters only, keep contractions
    words = re.findall(r"[a-z]+(?:'[a-z]+)?", text)
    return words


def word_frequencies(words: list[str], vocab: set[str]) -> dict[str, float]:
    """Return relative frequency (per total words) for each word in vocab."""
    counts = collections.Counter(w for w in words if w in vocab)
    total = len(words)
    if total == 0:
        return {w: 0.0 for w in vocab}
    return {w: counts.get(w, 0) / total for w in vocab}


def z_scores(freqs: dict[str, float], reference: dict[str, dict[str, float]]) -> dict[str, float]:
    """Compute z-scores for each word using mean/std across the reference corpora."""
    zs = {}
    for w in freqs:
        values = [c[w] for c in reference.values()]
        mean = sum(values) / len(values)
        var = sum((v - mean) ** 2 for v in values) / len(values)
        sd = math.sqrt(var) if var > 0 else 1.0
        zs[w] = (freqs[w] - mean) / sd
    return zs


def burrows_delta(z_a: dict[str, float], z_b: dict[str, float]) -> float:
    """Burrows's Delta: mean absolute difference of z-scores."""
    diffs = [abs(z_a[w] - z_b[w]) for w in z_a]
    return sum(diffs) / len(diffs)


def sentence_lengths(text: str) -> list[int]:
    """Return word counts per sentence."""
    text = text.replace("\n", " ")
    # naive splitter on sentence-ending punctuation
    sentences = re.split(r"[.!?]+", text)
    return [len(tokenize(s)) for s in sentences if tokenize(s)]


def stats(values: list[int]) -> dict[str, float]:
    if not values:
        return {"mean": 0.0, "std": 0.0, "median": 0.0}
    mean = sum(values) / len(values)
    var = sum((v - mean) ** 2 for v in values) / len(values)
    sd = math.sqrt(var)
    sorted_v = sorted(values)
    median = sorted_v[len(sorted_v) // 2]
    return {"mean": mean, "std": sd, "median": median}


def lexical_richness(words: list[str]) -> float:
    """Type-token ratio with capping for length sensitivity."""
    if not words:
        return 0.0
    return len(set(words)) / len(words)


def main() -> None:
    ap = argparse.ArgumentParser(description="Stylometric self-test")
    ap.add_argument("--manuscript", type=Path, required=True,
                    help="The manuscript (.docx, .txt, or .md)")
    ap.add_argument("--corpus", type=Path, required=True,
                    help="Your public-writing corpus (file or directory)")
    ap.add_argument("--top-n-features", type=int, default=20,
                    help="Number of distinguishing features to report")
    args = ap.parse_args()

    if not args.manuscript.exists():
        print(f"ERROR: {args.manuscript} not found")
        sys.exit(1)
    if not args.corpus.exists():
        print(f"ERROR: {args.corpus} not found")
        sys.exit(1)

    manuscript_text = read_one(args.manuscript)
    corpus_text = read_corpus(args.corpus)

    m_words = tokenize(manuscript_text)
    c_words = tokenize(corpus_text)

    print(f"Manuscript: {args.manuscript.name}")
    print(f"  Total words: {len(m_words):,}")
    print(f"  Unique words: {len(set(m_words)):,}")
    print()
    print(f"Comparison corpus: {args.corpus}")
    print(f"  Total words: {len(c_words):,}")
    print(f"  Unique words: {len(set(c_words)):,}")
    print()

    if not c_words:
        print("ERROR: the corpus contains no readable text, so there is nothing")
        print("to compare the manuscript against. No result is possible.")
        print()
        print("This usually means the directory held no supported files.")
        print(f"Supported types: {', '.join(SUPPORTED)}. PDFs are not read;")
        print("export them to .txt first.")
        print()
        print("Nothing below should be read as a finding about your anonymity.")
        sys.exit(1)

    if len(c_words) < MIN_CORPUS_WORDS:
        print(f"ERROR: corpus is {len(c_words):,} words, below the {MIN_CORPUS_WORDS:,}-word floor")
        print("needed for a meaningful comparison. A result computed on this little")
        print("text would be noise, and noise here reads as a false alarm.")
        print()
        print("Add more of your public writing and run it again.")
        sys.exit(1)

    if len(c_words) < 10000:
        print(f"WARNING: corpus is small ({len(c_words):,} words).")
        print("Aim for 20,000+ for stable comparison. Treat the verdict below")
        print("as directional, not settled.")
        print()

    # Compute frequencies on the function-word vocabulary
    m_freq = word_frequencies(m_words, FUNCTION_WORDS)
    c_freq = word_frequencies(c_words, FUNCTION_WORDS)

    # Burrows's Delta needs more than two reference texts to compute proper
    # z-scores. We simulate the reference by splitting the corpus into chunks.
    chunk_size = max(1000, len(c_words) // 5)
    chunks = []
    for i in range(0, len(c_words), chunk_size):
        chunk = c_words[i:i + chunk_size]
        if len(chunk) >= 500:
            chunks.append(word_frequencies(chunk, FUNCTION_WORDS))
    chunks.append(m_freq)
    if len(chunks) < 3:
        print("WARNING: corpus too small to compute robust Delta")
    reference = {f"chunk_{i}": ch for i, ch in enumerate(chunks)}

    z_m = z_scores(m_freq, reference)
    z_c = z_scores(c_freq, reference)
    delta = burrows_delta(z_m, z_c)

    print(f"Burrows's Delta: {delta:.3f}")
    print()
    if delta < 0.7:
        verdict = "STRONG MATCH. Stylometric anonymity is weak. Consider an edit pass focused on style variation."
    elif delta < 1.0:
        verdict = "MODERATE MATCH. Detectable similarity. An edit pass could help."
    elif delta < 1.3:
        verdict = "WEAK MATCH. Typical inter-author variance. Stylometric risk is low."
    else:
        verdict = "STRONG DIFFERENTIATION. Stylometric anonymity is robust."
    print(f"  Verdict: {verdict}")
    print()

    # Most distinguishing features (manuscript over-represented)
    diffs = sorted(
        ((w, z_m[w] - z_c[w]) for w in FUNCTION_WORDS),
        key=lambda x: x[1],
        reverse=True,
    )
    print(f"Top {args.top_n_features} function words over-represented in manuscript:")
    for w, d in diffs[:args.top_n_features]:
        print(f"  {w:<10}  delta={d:+.2f}  manuscript={m_freq[w]*1000:.1f}/1k  corpus={c_freq[w]*1000:.1f}/1k")
    print()
    print(f"Top {args.top_n_features} function words over-represented in corpus:")
    for w, d in diffs[-args.top_n_features:][::-1]:
        print(f"  {w:<10}  delta={d:+.2f}  manuscript={m_freq[w]*1000:.1f}/1k  corpus={c_freq[w]*1000:.1f}/1k")
    print()

    # Sentence-length comparison
    m_sentences = sentence_lengths(manuscript_text)
    c_sentences = sentence_lengths(corpus_text)
    m_stats = stats(m_sentences)
    c_stats = stats(c_sentences)
    print("Sentence-length comparison (in words):")
    print(f"  Manuscript: mean={m_stats['mean']:.1f}, std={m_stats['std']:.1f}, median={m_stats['median']:.0f}")
    print(f"  Corpus:     mean={c_stats['mean']:.1f}, std={c_stats['std']:.1f}, median={c_stats['median']:.0f}")
    print()

    # Lexical richness
    m_ttr = lexical_richness(m_words[:10000])  # cap at 10k for fairness
    c_ttr = lexical_richness(c_words[:10000])
    print("Lexical richness (type-token ratio, first 10k words):")
    print(f"  Manuscript: {m_ttr:.3f}")
    print(f"  Corpus:     {c_ttr:.3f}")
    print()

    # Summary
    print("=== SUMMARY ===")
    print(f"Delta: {delta:.3f}  ({verdict.split('.')[0]})")
    if abs(m_stats['mean'] - c_stats['mean']) > 5:
        print(f"  Sentence length differs noticeably ({m_stats['mean']:.0f} vs {c_stats['mean']:.0f}). "
              "This is a defense if it's intentional.")
    if abs(m_ttr - c_ttr) > 0.05:
        print(f"  Lexical richness differs ({m_ttr:.2f} vs {c_ttr:.2f}). "
              "This is a defense if intentional.")


if __name__ == "__main__":
    main()
