"""V4.3 mixture support: two web stores with keep masks plus the chat bank."""
import json
from pathlib import Path
import numpy as np

from vimeml.training.data import file_sha, write_json
from vimeml.training.data_v2 import PrefixCropWindowDataset
from vimeml.training.data_mixture import JoinedPrefixWindows
from vimeml.posttraining.streams import ConversionWindows

ROOT = Path(__file__).resolve().parents[3]


def window_sentences(dataset, ids):
    position = np.searchsorted(dataset.end, ids, side='right')
    previous = np.zeros(len(ids), dtype=np.int64)
    present = position > 0
    previous[present] = dataset.end[position[present] - 1] - dataset.sentences[position[present] - 1] - 1
    sentences = ids - previous
    candidates = np.flatnonzero(position < len(dataset.end))
    extra = candidates[ids[candidates] >= dataset.first[position[candidates]]]
    sentences[extra] = dataset.sentences[position[extra]]
    return sentences


def sentence_keys(paths):
    """Canonical text keys of every sentence in the listed TXT files."""
    import hashlib
    import unicodedata
    keys = set()
    for path in paths:
        with Path(path).open(encoding='utf-8') as stream:
            for line in stream:
                text = line.rstrip('\r\n')
                if text:
                    keys.add(hashlib.blake2b(' '.join(unicodedata.normalize('NFKC', text).split()).encode(),
                                            digest_size=16).digest())
    return keys


def prepare_maps_v43(directory, root=ROOT):
    """Rebuild epoch maps for the three-store V4.3 mixture, verifying SHAs."""
    directory, root = Path(directory), Path(root)
    request = json.loads((directory / 'request.json').read_text('utf-8'))
    expected = (json.loads((directory / 'manifest.json').read_text('utf-8'))
                if (directory / 'manifest.json').exists() else None)
    if expected and all((directory / r['file']).exists() and file_sha(directory / r['file']) == r['sha256']
                        for r in expected['epochs']):
        return expected
    bases = []
    for record in request['joined_stores']:
        if (file_sha(root / record['token_dir'] / 'manifest.json') != record['token_manifest_sha256']
                or file_sha(root / record['index_dir'] / 'manifest.json') != record['window_manifest_sha256']):
            raise ValueError('Joined bank identity changed')
        bases.append(PrefixCropWindowDataset(root / record['token_dir'], root / record['index_dir'],
            seed=request['seed'], probability=request['crop_probability'],
            min_remaining_tokens=request['crop_minimum']))
    web_parts = []
    for record, base in zip(request['joined_stores'][:-1], bases[:-1]):
        keep_path = root / record['keep']
        if file_sha(keep_path) != record['keep_sha256']:
            raise ValueError(f"Store keep mask changed: {record['token_dir']}")
        keep = np.load(keep_path, mmap_mode='r', allow_pickle=False)
        chunks = []
        for start in range(0, len(base), 262144):
            ids = np.arange(start, min(len(base), start + 262144), dtype=np.int64)
            chunks.append(ids[keep[window_sentences(base, ids)]])
        web_parts.append(np.concatenate(chunks))
    web = np.concatenate(web_parts)
    joined = JoinedPrefixWindows(bases, request['primary_store'],
                                 [{int(k): v for k, v in r['source_map'].items()} for r in request['joined_stores']])
    chat_base = bases[-1]
    try:
        chat_start = sum(len(base) for base in bases[:-1])
        chat = np.arange(chat_start, chat_start + len(chat_base), dtype=np.int64)
        records = []
        for epoch in range(request['epochs']):
            joined.set_epoch(epoch)
            web_tokens = 0
            for start in range(0, len(web), 262144):
                web_tokens += int(joined.window_lengths(web[start:start + 262144]).sum())
            chat_lengths = joined.window_lengths(chat)
            target = round(web_tokens * request['chat_fraction'] / (1 - request['chat_fraction']))
            rng = np.random.default_rng(request['seed'] + epoch)
            chosen, used = [], 0
            while used < target:
                order = rng.permutation(len(chat))
                cumulative = np.cumsum(chat_lengths[order], dtype=np.int64)
                take = int(np.searchsorted(cumulative, target - used, side='right'))
                if take:
                    chosen.append(chat[order[:take]])
                    used += int(cumulative[take - 1])
                if take < len(order):
                    break
            chosen = np.concatenate(chosen)
            mapping = np.concatenate((web, chosen)).astype('<u4')
            rng.shuffle(mapping)
            path = directory / f'epoch-{epoch}.npy'
            temp = path.with_suffix('.tmp.npy')
            np.save(temp, mapping, allow_pickle=False)
            temp.replace(path)
            record = {'epoch': epoch, 'file': path.name, 'bytes': path.stat().st_size,
                      'sha256': file_sha(path), 'windows': len(mapping),
                      'web_windows': len(web), 'chat_windows': len(chosen),
                      'web_tokens': web_tokens, 'chat_tokens': used,
                      'effective_tokens': web_tokens + used,
                      'chat_fraction': used / (web_tokens + used),
                      'chat_effective_passes': used / int(chat_lengths.sum())}
            if expected and record != expected['epochs'][epoch]:
                raise ValueError(f'Rebuilt map differs from frozen epoch {epoch}')
            records.append(record)
            print(json.dumps({'map_epoch': epoch + 1, 'windows': len(mapping),
                              'chat_fraction': record['chat_fraction']}), flush=True)
        result = {'format': 'vimeml_token_mixture_v1', 'status': 'complete', 'request': request,
                  'epochs': records,
                  'policy': 'V4.3: retained old LM windows and deduplicated expansion windows once per epoch; '
                            'chat-only repeated shuffled cycles at 5% actual cropped web targets; maps SHA256-pinned'}
        if expected and result != expected:
            raise ValueError('Rebuilt mixture manifest differs')
        if expected is None:
            write_json(directory / 'manifest.json', result)
        return result
    finally:
        joined.close()