"""
Vocab Minification Tool for NeuroFlow NFv1.1

Scans training corpus to identify which tokens in the 128K vocabulary
actually appear. Generates:
1. active_vocab_mask.json  - binary mask (1=active, 0=cold) for gradient gating
2. vocab_coverage.json     - coverage statistics

Usage:
    python vocab_minification.py --corpus <corpus_file> --tokenizer <tokenizer.json> --output <output_dir>
    python vocab_minification.py --corpus data/train.tok1 --tokenizer configs/tokenizer_cn_013.json --output configs/
"""

import json
import argparse
import os
import sys
from collections import Counter


def load_tokenizer(tokenizer_path):
    with open(tokenizer_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def scan_tok1_corpus(corpus_path, vocab_size):
    token_counts = Counter()
    total_tokens = 0
    line_count = 0

    with open(corpus_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                tokens = [int(x) for x in line.split()]
                for t in tokens:
                    if 0 <= t < vocab_size:
                        token_counts[t] += 1
                        total_tokens += 1
            except ValueError:
                continue
            line_count += 1
            if line_count % 10000 == 0:
                print(f"  Scanned {line_count} lines, {total_tokens} tokens, {len(token_counts)} unique", flush=True)

    return token_counts, total_tokens, line_count


def scan_text_corpus(corpus_path, tokenizer_data):
    from preprocess_corpus import BPETokenizer
    tokenizer = BPETokenizer(tokenizer_data)
    vocab_size = tokenizer_data.get('vocab_size', 128000)

    token_counts = Counter()
    total_tokens = 0
    line_count = 0

    with open(corpus_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            tokens = tokenizer.encode(line)
            for t in tokens:
                if 0 <= t < vocab_size:
                    token_counts[t] += 1
                    total_tokens += 1
            line_count += 1
            if line_count % 10000 == 0:
                print(f"  Scanned {line_count} lines, {total_tokens} tokens, {len(token_counts)} unique", flush=True)

    return token_counts, total_tokens, line_count


def generate_mask(token_counts, vocab_size, min_freq=1):
    mask = [0] * vocab_size
    active_count = 0
    for token_id in range(vocab_size):
        if token_counts.get(token_id, 0) >= min_freq:
            mask[token_id] = 1
            active_count += 1
    return mask, active_count


def main():
    parser = argparse.ArgumentParser(description='NeuroFlow Vocab Minification')
    parser.add_argument('--corpus', required=True, help='Path to training corpus')
    parser.add_argument('--tokenizer', default='configs/tokenizer_cn_013.json', help='Tokenizer config')
    parser.add_argument('--output', default='configs/', help='Output directory')
    parser.add_argument('--min-freq', type=int, default=1, help='Minimum token frequency to be active')
    parser.add_argument('--format', choices=['tok1', 'text'], default='tok1', help='Corpus format')
    args = parser.parse_args()

    tokenizer_data = load_tokenizer(args.tokenizer)
    vocab_size = tokenizer_data.get('vocab_size', 128000)
    print(f"Vocab size: {vocab_size}")
    print(f"Scanning corpus: {args.corpus} (format: {args.format})")

    if args.format == 'tok1':
        token_counts, total_tokens, line_count = scan_tok1_corpus(args.corpus, vocab_size)
    else:
        token_counts, total_tokens, line_count = scan_text_corpus(args.corpus, tokenizer_data)

    print(f"\nScan complete:")
    print(f"  Lines: {line_count}")
    print(f"  Total tokens: {total_tokens}")
    print(f"  Unique tokens: {len(token_counts)}")
    print(f"  Coverage: {len(token_counts)/vocab_size*100:.1f}%")

    mask, active_count = generate_mask(token_counts, vocab_size, args.min_freq)
    cold_count = vocab_size - active_count

    print(f"\nMask generated (min_freq={args.min_freq}):")
    print(f"  Active tokens: {active_count} ({active_count/vocab_size*100:.1f}%)")
    print(f"  Cold tokens:   {cold_count} ({cold_count/vocab_size*100:.1f}%)")

    os.makedirs(args.output, exist_ok=True)

    mask_path = os.path.join(args.output, 'active_vocab_mask.json')
    with open(mask_path, 'w', encoding='utf-8') as f:
        json.dump({
            'vocab_size': vocab_size,
            'active_count': active_count,
            'cold_count': cold_count,
            'min_freq': args.min_freq,
            'mask': mask
        }, f)
    print(f"  Saved: {mask_path}")

    coverage_path = os.path.join(args.output, 'vocab_coverage.json')
    top_tokens = token_counts.most_common(100)
    with open(coverage_path, 'w', encoding='utf-8') as f:
        json.dump({
            'vocab_size': vocab_size,
            'total_tokens': total_tokens,
            'unique_tokens': len(token_counts),
            'coverage_pct': len(token_counts) / vocab_size * 100,
            'active_count': active_count,
            'cold_count': cold_count,
            'top_100': [[str(t), c] for t, c in top_tokens],
            'freq_distribution': {
                '1_occurrence': sum(1 for c in token_counts.values() if c == 1),
                '2_10_occurrences': sum(1 for c in token_counts.values() if 2 <= c <= 10),
                '11_100_occurrences': sum(1 for c in token_counts.values() if 11 <= c <= 100),
                '100_plus': sum(1 for c in token_counts.values() if c > 100),
            }
        }, f, indent=2)
    print(f"  Saved: {coverage_path}")

    print(f"\nMemory savings estimate:")
    embed_params = vocab_size * 512
    cold_params = cold_count * 512
    print(f"  Embedding params: {embed_params:,} ({embed_params*4/1024/1024:.1f} MB)")
    print(f"  Cold token params: {cold_params:,} ({cold_params*4/1024/1024:.1f} MB)")
    print(f"  Gradient zeroed: {cold_params:,} params (saves {cold_params*4/1024/1024:.1f} MB gradient memory)")


if __name__ == '__main__':
    main()