"""Session 17: five list items, fixed arithmetic families, cross-item validation.
Python + NumPy. Local inputs only. No parameter selection uses items 4 or 5.
"""
import hashlib
import itertools
import json
from word_core import *

STREAMS = ['prime_minus_one', 'integer_totient', 'fibonacci']
SCHEDULES = ['continuous_list', 'restart_item', 'label_start']
PARAMS = list(itertools.product(range(33), (1, -1), (1, -1), range(29)))
PAR = np.array(PARAMS)


def extract():
    section = '%'.join(lp[:56]).split('$')[7]
    markers = list(re.finditer(r'([1-5])-', section))
    assert [m.group(1) for m in markers] == list('12345')
    # Source image 55.jpg shows a red paragraph separator before this word.
    # The master transcription does not encode its color or paragraph role.
    stop = section.index('ᚫᛞᛟᚫᛟᛗ-', markers[-1].end())
    items = []
    for i, m in enumerate(markers):
        end = markers[i + 1].start() if i < 4 else stop
        source = section[m.end():end]
        words = tokenize(source)
        start_offset = len(vals(section[:m.end()]))
        count = len(vals(source))
        items.append(dict(label=i + 1, source=source,
            page_start=33 + section[:m.end()].count('%'),
            page_end=33 + section[:end].rstrip('/\n% ').count('%'),
            section_start=start_offset, section_end_exclusive=start_offset + count,
            rune_count=count, word_count=len(words),
            word_lengths=[len(w['values']) for w in words],
            first_word=words[0]['runes']))
    assert [x['rune_count'] for x in items] == [139, 91, 98, 96, 129]
    assert [x['word_count'] for x in items] == [30, 20, 25, 19, 28]
    assert items[4]['source'].count('.') == 2
    assert items[3]['source'].count('%') == 1
    assert items[3]['section_start'] < 1121 < items[3]['section_end_exclusive']
    return section, items


def indices(sources, schedule):
    out = []; at = 0
    for i, s in enumerate(sources):
        n = len(vals(s))
        start = at if schedule == 'continuous_list' else i if schedule == 'label_start' else 0
        out.append(np.arange(n) + start)
        at += n
    return out


def decode(sources, stream, schedule, params):
    offset, a, b, shift = params
    return [(a * vals(s) + b * SEQS[stream][ix + offset] + shift) % 29
            for s, ix in zip(sources, indices(sources, schedule))]


def one_score(v, s):
    return float(FULL_MODEL.score(v[None, :], [len(w['values']) for w in tokenize(s)])[0])


def search(sources):
    winners = []
    for stream in STREAMS:
        for schedule in SCHEDULES:
            ii = indices(sources, schedule)
            summed = np.zeros(len(PARAMS)); weight = 0
            # Model and parameters see only list items 1, 2, 3 here.
            for s, ix in zip(sources[:3], ii[:3]):
                lengths = [len(w['values']) for w in tokenize(s)]
                q = SEQS[stream][PAR[:, 0, None] + ix[None, :]]
                decoded = (PAR[:, 1, None] * vals(s)[None, :] +
                    PAR[:, 2, None] * q + PAR[:, 3, None]) % 29
                w = len(vals(s)) + len(lengths)
                summed += FULL_MODEL.score(decoded, lengths) * w
                weight += w
            scores = summed / weight
            pick = int(np.argmax(scores))
            params = PARAMS[pick]
            d = decode(sources, stream, schedule, params)
            winners.append(dict(stream=stream, schedule=schedule,
                parameters=dict(zip(['offset', 'cipher_sign', 'sequence_sign', 'shift'], params)),
                parameter_tuple=params, selection_score=float(scores[pick]),
                item_scores=[one_score(v, s) for v, s in zip(d, sources)],
                decoded_runes=[''.join(R[int(x)] for x in v) for v in d],
                decoded_text=[render(s, v) for s, v in zip(sources, d)]))
    # Selection still uses only items 1-3. Scores for 4-5 never affect this order.
    return sorted(winners, key=lambda x: x['selection_score'], reverse=True)


def controls():
    # Distinct, non-overlapping word blocks of known LP2 56+57 plaintext.
    # Both pages are outside the LP1 scoring-model training corpus.
    v = vals(lp[56]); ix = np.arange(len(v)); ix[57:] -= 1
    d = (v - SEQS['prime_minus_one'][ix]) % 29; d[56] = v[56]
    known = replace_values(lp[56], d) + '%' + lp[57]
    words = tokenize(known)
    blocks = np.array_split(np.arange(len(words)), 5)
    plain = ['-'.join(words[int(j)]['runes'] for j in block) for block in blocks]
    results = []
    for test, (stream, schedule) in enumerate(itertools.product(STREAMS, SCHEDULES)):
        offset = 7; a = 1 if test % 2 else -1; b = -1 if test % 3 else 1; shift = 11
        cipher = []; at = 0
        # Scalar encryption, separately from vectorized decoder/index routines.
        for item, s in enumerate(plain):
            pv = vals(s); cv = []
            for j, p in enumerate(pv):
                start = at if schedule == 'continuous_list' else item if schedule == 'label_start' else 0
                q = int(SEQS[stream][offset + start + j])
                cv.append((a * (int(p) - b * q - shift)) % 29)
            cipher.append(replace_values(s, cv)); at += len(pv)
        recovered = decode(cipher, stream, schedule, (offset, a, b, shift))
        replay = all(np.array_equal(x, vals(y)) for x, y in zip(recovered, plain))
        assert replay
        picked = search(cipher)[0]
        exact = all(x == ''.join(c for c in s if c in R) for x, s in zip(picked['decoded_runes'], plain))
        results.append(dict(encryption_stream=stream, encryption_schedule=schedule,
            replay_exact=replay, selection_on_first_three_recovers_all_five=exact,
            selected_stream=picked['stream'], selected_schedule=picked['schedule'],
            selected_parameters=picked['parameters'], item_rune_counts=[len(vals(s)) for s in plain]))
    return results


def repetitions(items):
    ws = [tokenize(x['source']) for x in items]; shared = []
    for i, j in itertools.combinations(range(5), 2):
        common = sorted(set(w['runes'] for w in ws[i]) & set(w['runes'] for w in ws[j]))
        aligned = [dict(word_index=k, runes=a['runes'])
            for k, (a, b) in enumerate(zip(ws[i], ws[j])) if a['runes'] == b['runes']]
        wl_prefix = 0
        for a, b in zip(ws[i], ws[j]):
            if len(a['values']) != len(b['values']): break
            wl_prefix += 1
        pairs_i = {tuple(w['runes'] for w in ws[i][k:k+2]) for k in range(len(ws[i])-1)}
        pairs_j = {tuple(w['runes'] for w in ws[j][k:k+2]) for k in range(len(ws[j])-1)}
        shared.append(dict(items=[i+1,j+1], common_words=common, aligned_words=aligned,
            shared_adjacent_word_pairs=sorted(pairs_i & pairs_j),
            matching_word_length_prefix=wl_prefix))
    return shared


def main():
    # Hand-specified indices exercise all three counter policies.
    fixture = ['ᚠ-ᚢ', 'ᚦᚩᚱ', 'ᚳ']
    assert [x.tolist() for x in indices(fixture, 'continuous_list')] == [[0,1],[2,3,4],[5]]
    assert [x.tolist() for x in indices(fixture, 'restart_item')] == [[0,1],[0,1,2],[0]]
    assert [x.tolist() for x in indices(fixture, 'label_start')] == [[0,1],[1,2,3],[2]]
    section, items = extract(); sources = [x['source'] for x in items]
    plan = dict(streams=STREAMS, schedules=SCHEDULES, offsets=list(range(33)),
        cipher_signs=[1,-1], sequence_signs=[1,-1], shifts=list(range(29)),
        candidate_count=len(STREAMS)*len(SCHEDULES)*len(PARAMS),
        selection_items=[1,2,3], validation_items=[4,5],
        item5_primary='Through second sentence, before visible red paragraph break',
        item5_sensitivity='First sentence only; reuse all selected parameters',
        skips=False, digits_or_punctuation_consume_positions=False)
    # Record the fixed configuration before reading candidate outputs.
    (ROOT/'numbered_list_plan.json').write_text(json.dumps(plan, indent=2))
    checks = controls()
    winners = search(sources)
    short = sources[4].split('.',1)[0]
    short_n = len(vals(short)); extra = sources[4].split('.',1)[1]
    for w in winners:
        dv = decode(sources, w['stream'], w['schedule'], w['parameter_tuple'])[4]
        w['item5_first_sentence_score'] = one_score(dv[:short_n], short)
        w['item5_second_sentence_score'] = one_score(dv[short_n:], extra)
    baselines = []
    for name, a, shift in [('direct',1,0),('atbash',-1,28),('shifted_atbash',-1,2)]:
        vs = [(a*vals(s)+shift)%29 for s in sources]
        baselines.append(dict(mapping=name, item_scores=[one_score(v,s) for v,s in zip(vs,sources)],
            decoded_text=[render(s,v) for s,v in zip(sources,vs)]))
    result = dict(plan=plan, items=items, total_runes=sum(x['rune_count'] for x in items),
        total_words=sum(x['word_count'] for x in items), item5_short_runes=short_n,
        transcription_sha256=hashlib.sha256((ROOT/'liber-primus__transcription--master.txt').read_bytes()).hexdigest(),
        image_sha256={n:hashlib.sha256((ROOT/n).read_bytes()).hexdigest() for n in ['53.jpg','54.jpg','55.jpg']},
        first_runes=[x['first_word'][0] for x in items], repetitions=repetitions(items),
        fixture_checks=3, controls=checks, winners=winners, baselines=baselines)
    (ROOT/'numbered_list_results.json').write_text(json.dumps(result, indent=2, ensure_ascii=False))
    print('List:', result['total_runes'], 'runes;', result['total_words'], 'words')
    print('Candidates:', plan['candidate_count'], '; family winners:',len(winners))
    print('Exact search recovery controls:', sum(x['selection_on_first_three_recovers_all_five'] for x in checks), '/', len(checks))
    for w in winners:
        print(w['stream'], w['schedule'], w['parameters'], round(w['selection_score'],4),
            'validation', [round(x,4) for x in w['item_scores'][3:]])


if __name__ == '__main__':
    main()