#!/usr/bin/env python3
"""Reproduce the published orthography atlas from two local official ZIP files.

Read-only inputs; no network requests. Requires Python >=3.10 and pyxlsb==1.0.10.
Example:
  python scripts/audit-atlas.py --infra /path/lexique-infra.zip \
    --lexique /path/Lexique383.zip --output /tmp/atlas-reproduction.json
The output contains aggregates and source provenance, never the full corpora.
"""
import argparse
import collections
import csv
import hashlib
import io
import json
import statistics
import tempfile
import zipfile
from pathlib import Path

SOURCES = {
    'infra': {
        'name': 'Lexique-Infra 1.11',
        'url': 'https://osf.io/xenas/download',
        'metadataUrl': 'https://api.osf.io/v2/files/xenas/',
        'sha256': 'd3cf11506ff572918045d8ab8c72bedcdc33c91a07c317cee05ab3f6688e5c22',
    },
    'lexique': {
        'name': 'Lexique 3.83',
        'url': 'https://lexique.org/databases/Lexique383/Lexique383.zip',
        'sha256': 'e181d132b3b0d3d87efc98d376968441b517353933011cebea5366321a6024e6',
    },
}
PROTOCOL = {
    'unit': 'One orthographic form with exactly one row in Lexique-Infra',
    'filter': 'Letters only (Python str.isalpha), 4–20 source characters, freqlivres >=10; join Lexique383 on ortho and cgram, requiring exactly one match',
    'normalization': 'Compute on original source spelling; replace oe with œ only in the displayed top 100',
    'components': {
        'mean': 'Negative FreqTok_PG',
        'minimum': 'Negative minfreqgraphTok_PG',
        'irregularShare': 'countregTo_PG / number of dot-separated graphemes',
        'complexity': 'complexmoygraph',
    },
    'percentiles': 'Average one-based rank for ties, divided by selected corpus size',
    'weights': {'mean': .35, 'minimum': .25, 'irregularShare': .25, 'complexity': .15},
    'score': '100 * weighted sum of component percentiles; use integer percentage weights before division to preserve ties; round only for display',
    'order': 'Descending unrounded score, then descending book frequency, then source spelling',
    'correlation': 'Pearson correlation of average ranks (Spearman), computed on all selected forms',
    'silentForms': 'Count each grapheme at most once per form; phoneme # marks silence',
    'soundProbabilities': 'Proba PG sheet, whole-word token probabilities (column G); token frequencies derive from Lexique film subtitles, not book frequency',
}


def average_ranks(values):
    """One-based average ranks, including ties; do not round input values."""
    ranks = [0.0] * len(values)
    ordered = sorted(range(len(values)), key=values.__getitem__)
    start = 0
    while start < len(ordered):
        end = start + 1
        while end < len(ordered) and values[ordered[end]] == values[ordered[start]]:
            end += 1
        rank = (start + 1 + end) / 2
        for index in ordered[start:end]:
            ranks[index] = rank
        start = end
    return ranks


def select_forms(infra, lexique):
    counts = collections.Counter(row['item'] for row in infra)
    lookup = collections.defaultdict(list)
    for row in lexique:
        lookup[(row['ortho'], row['cgram'])].append(row)
    selected = []
    for row in infra:
        form = row['item']
        if counts[form] != 1 or not form.isalpha() or not 4 <= len(form) <= 20:
            continue
        matches = lookup[(form, row['cgram'])]
        if len(matches) != 1 or float(matches[0]['freqlivres']) < 10:
            continue
        selected.append({
            'form': form, 'category': row['cgram'],
            'frequency': float(matches[0]['freqlivres']), 'length': len(form),
            'mean': -float(row['FreqTok_PG']), 'minimum': -float(row['minfreqgraphTok_PG']),
            'irregularShare': float(row['countregTo_PG']) / len(row['grapheme'].split('.')),
            'complexity': float(row['complexmoygraph']), 'associations': row['assoc'],
        })
    return selected


def score_forms(forms):
    n = len(forms)
    if n < 2:
        raise ValueError('At least two selected forms are required')
    ranks = {key: average_ranks([row[key] for row in forms]) for key in PROTOCOL['weights']}
    for i, row in enumerate(forms):
        row['score'] = sum(round(100 * PROTOCOL['weights'][key]) * ranks[key][i] for key in ranks) / n
        row['equalScore'] = 25 * sum(ranks[key][i] for key in ranks) / n
    return sorted(forms, key=lambda row: (-row['score'], -row['frequency'], row['form']))


def silent_counts(forms):
    occurrences, affected, final = collections.Counter(), collections.Counter(), collections.Counter()
    any_silent = 0
    repeated = []
    for row in forms:
        parts = row['associations'].split('.')
        silent = [part[:-2] for part in parts if part.endswith('-#')]
        any_silent += bool(silent)
        occurrences.update(silent)
        affected.update(set(silent))
        if parts[-1].endswith('-#'):
            final.update([parts[-1][:-2]])
        if len(silent) != len(set(silent)):
            repeated.append(row['form'])
    return {
        'formsWithAnySilentGrapheme': any_silent,
        'percentWithAnySilentGrapheme': 100 * any_silent / len(forms),
        'graphemes': [{'grapheme': g, 'forms': n, 'occurrences': occurrences[g], 'finalForms': final[g]}
                      for g, n in sorted(affected.items(), key=lambda item: (-item[1], item[0]))],
        'formsWithRepeatedSilentGrapheme': sorted(repeated),
    }


def verify_archive(path, key):
    actual = hashlib.sha256(Path(path).read_bytes()).hexdigest()
    if actual != SOURCES[key]['sha256']:
        raise ValueError(f'{key} archive SHA256 mismatch: {actual}; expected {SOURCES[key]["sha256"]}')


def load_sources(infra_path, lexique_path):
    from pyxlsb import open_workbook
    verify_archive(infra_path, 'infra')
    verify_archive(lexique_path, 'lexique')
    with zipfile.ZipFile(lexique_path) as archive:
        lexique = list(csv.DictReader(io.StringIO(archive.read('Lexique383.tsv').decode('utf-8')), delimiter='\t'))
    with tempfile.TemporaryDirectory(prefix='atlas-') as directory:
        with zipfile.ZipFile(infra_path) as archive:
            filename = next(name for name in archive.namelist() if 'Corresp' in name and name.endswith('.xlsb'))
            workbook_path = Path(directory) / 'correspondances.xlsb'
            workbook_path.write_bytes(archive.read(filename))
        with open_workbook(str(workbook_path)) as workbook:
            with workbook.get_sheet('Mots') as sheet:
                infra, headers = [], None
                for cells in sheet.rows():
                    values = [cell.v for cell in cells][:42]
                    if headers is None:
                        if values[:3] == ['item', 'phono', 'cgram']:
                            headers = values
                        continue
                    if values and values[0]:
                        infra.append(dict(zip(headers, values)))
            sounds = collections.defaultdict(list)
            with workbook.get_sheet('Proba PG') as sheet:
                for cells in sheet.rows():
                    values = [cell.v for cell in cells]
                    if len(values) < 7 or not isinstance(values[6], (int, float)):
                        continue
                    sound = str(int(values[1])) if isinstance(values[1], float) and values[1].is_integer() else str(values[1])
                    sounds[sound].append({'grapheme': values[0], 'probability': values[6]})
    if not infra:
        raise ValueError('No source records found after the expected Mots headers')
    return infra, lexique, sounds


def audit(infra, lexique, sounds):
    counts = collections.Counter(row['item'] for row in infra)
    forms = score_forms(select_forms(infra, lexique))
    scores = [row['score'] for row in forms]
    lengths = [row['length'] for row in forms]
    equal = sorted(forms, key=lambda row: (-row['equalScore'], -row['frequency'], row['form']))
    lengths_grouped = collections.defaultdict(list)
    for row in forms:
        lengths_grouped[row['length']].append(row['score'])
    sound_codes = {'E': 'ɛ', 'e': 'e', '5': 'ɛ̃', '9': 'œ', 's': 's', 'j': 'j'}
    return {
        'sources': SOURCES, 'protocol': PROTOCOL,
        'sourceRows': len(infra), 'distinctForms': len(counts),
        'multiEntryForms': sum(count > 1 for count in counts.values()),
        'selectedForms': len(forms), **silent_counts(forms),
        'lengthScoreSpearman': statistics.correlation(average_ranks(lengths), average_ranks(scores)),
        'equalWeightsSpearman': statistics.correlation(average_ranks(scores), average_ranks([row['equalScore'] for row in forms])),
        'top20Overlap': len({row['form'] for row in forms[:20]} & {row['form'] for row in equal[:20]}),
        'lengthMedians': [{'letters': length, 'forms': len(values), 'median': statistics.median(values)}
                          for length, values in sorted(lengths_grouped.items())],
        'sounds': [{'phoneme': sound_codes[code], 'attested': len(sounds[code]),
                    'atLeastHalfPercent': sum(row['probability'] >= .005 for row in sounds[code]),
                    'graphemes': sorted(sounds[code], key=lambda row: -row['probability'])} for code in sound_codes],
        'top100': [{'rank': i + 1, 'form': row['form'].replace('oe', 'œ'), 'category': row['category'],
                    'frequencyBooksPerMillion': row['frequency'], 'structuralIndex': round(row['score'], 1)}
                   for i, row in enumerate(forms[:100])],
    }


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--infra', required=True, type=Path, help='Official Lexique-Infra 1.11 ZIP')
    parser.add_argument('--lexique', required=True, type=Path, help='Official Lexique383 ZIP')
    parser.add_argument('--output', type=Path, help='Write aggregate JSON here instead of stdout')
    args = parser.parse_args()
    result = audit(*load_sources(args.infra, args.lexique))
    text = json.dumps(result, ensure_ascii=False, indent=2) + '\n'
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(text, encoding='utf-8')
    else:
        print(text, end='')
