"""Fetch five public-domain literary works and train a fresh RU/EN byte BPE.
Source texts are not redistributed. Check copyright in your jurisdiction.
"""
import hashlib
import json
import re
from pathlib import Path
import numpy as np
import requests
from tokenizers import ByteLevelBPETokenizer
from transformers import GPT2TokenizerFast

ROOT = Path(__file__).resolve().parent
CACHE = ROOT / '.cache'
MODEL = ROOT / 'model'
SOURCES = [
    (14741, 'ru', 'Духовные оды', 'Gavriil Derzhavin'),
    (16527, 'ru', '1001 задача для умственного счета', 'Sergei Rachinskii'),
    (30774, 'ru', 'Московия в представлении иностранцев XVI-XVII в.', 'Pavel Apostol; editor Georgii Lukomskii'),
    (11, 'en', "Alice’s Adventures in Wonderland", 'Lewis Carroll'),
    (1342, 'en', 'Pride and Prejudice', 'Jane Austen'),
]

def main():
    CACHE.mkdir(exist_ok=True)
    MODEL.mkdir(exist_ok=True)
    manifest, texts = [], {'ru': [], 'en': []}
    for book, lang, title, author in SOURCES:
        url = f'https://www.gutenberg.org/cache/epub/{book}/pg{book}.txt'
        local = CACHE / f'pg{book}.txt'
        if not local.exists():
            response = requests.get(url, timeout=90)
            response.raise_for_status()
            local.write_bytes(response.content)
        raw = local.read_bytes()
        text = raw.decode('utf-8-sig').replace('\r\n', '\n')
        start = re.search(r'\*\*\* START OF (?:THE|THIS) PROJECT GUTENBERG EBOOK.*?\*\*\*', text, re.S)
        end = re.search(r'\*\*\* END OF (?:THE|THIS) PROJECT GUTENBERG EBOOK', text)
        if not start or not end:
            raise ValueError(f'Cannot locate Gutenberg boundaries for {book}')
        text = text[start.end():end.start()].strip()
        if len(text) < 10000:
            raise ValueError('Unexpectedly short source')
        # Keep the final 5% of each work out of tokenizer and model training.
        boundary = text.rfind('\n', 0, int(len(text) * .95))
        train, valid = text[:boundary], text[boundary:]
        texts[lang].append((train, valid))
        manifest.append(dict(id=book, language=lang, title=title, author=author,
            url=url, catalog_url=f'https://www.gutenberg.org/ebooks/{book}',
            raw_sha256=hashlib.sha256(raw).hexdigest(),
            train_characters=len(train), validation_characters=len(valid)))
        print(f'{lang}: {title}: {len(text):,} chars', flush=True)
    training_files = []
    for lang in texts:
        p = CACHE / f'train_{lang}.txt'
        p.write_text('\n\n'.join(t[0] for t in texts[lang]), encoding='utf-8')
        training_files.append(str(p))
    bpe = ByteLevelBPETokenizer(add_prefix_space=False)
    bpe.train(files=training_files, vocab_size=23258, min_frequency=1,
              special_tokens=['<|endoftext|>', '<|pad|>'])
    if bpe.get_vocab_size() != 23258:
        raise ValueError(f'Unexpected vocab size: {bpe.get_vocab_size()}')
    bpe.save_model(str(MODEL))
    tok = GPT2TokenizerFast(vocab_file=str(MODEL/'vocab.json'),
        merges_file=str(MODEL/'merges.txt'), bos_token='<|endoftext|>',
        eos_token='<|endoftext|>', unk_token='<|endoftext|>', pad_token='<|pad|>',
        model_max_length=512, clean_up_tokenization_spaces=False)
    tok.save_pretrained(MODEL)
    for lang, docs in texts.items():
        for split, index in [('train', 0), ('valid', 1)]:
            ids = []
            for doc in docs:
                ids.extend(tok.encode(doc[index], add_special_tokens=False, truncation=False))
                ids.append(tok.eos_token_id)
            np.save(CACHE/f'{split}_{lang}.npy', np.array(ids, dtype=np.int32))
            print(f'{split}/{lang}: {len(ids):,} tokens', flush=True)
    (ROOT/'data_sources.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')

if __name__ == '__main__':
    main()