"""Relationship controls; analysis snapshots are scoped to model and tensor."""
import numpy as np
import streamlit as st

from relationships import analyse_pairs


def render_relationship_controls(vectors, label, scope, save_calculation):
    options = list(vectors)
    active_key = f'active_relationship:{scope}'
    example_key = f'relationship_examples:{scope}'

    def capture(key):
        return {'key': key, 'label': label(key), 'values': np.array(vectors[key], dtype=float, copy=True)}

    def choose(text, suffix, index=0, optional=False):
        choices = [None] + options if optional else options
        return st.selectbox(text, choices, index=0 if optional else min(index, len(choices) - 1),
                            format_func=lambda key: 'None' if key is None else label(key),
                            key=f'rel_{suffix}:{scope}')

    def remove_example(identifier):
        st.session_state[example_key] = [p for p in st.session_state[example_key] if p['id'] != identifier]

    st.caption('A → B is to C → D. Show relationship applies these inputs to the chart.')
    if len(options) < 2:
        st.info('Add at least two tokens or calculated vectors to the chart to choose a relationship.')
        return st.session_state.get(active_key)
    mode = st.selectbox('Relationship task', ['Compare relationships', 'Complete analogy', 'Average example pairs'],
                        key=f'rel_task:{scope}')
    if mode == 'Average example pairs':
        example_a = choose('Example A', 'example_a')
        example_b = choose('Example B', 'example_b', 1)
        pairs = st.session_state.setdefault(example_key, [])
        if st.button('Add example pair', key='add_relationship_example', disabled=len(pairs) >= 10):
            serial = st.session_state.get('relationship_example_serial', 0) + 1
            st.session_state.relationship_example_serial = serial
            pairs.append({'id': serial, 'a': capture(example_a), 'b': capture(example_b)})
        for number, pair in enumerate(pairs, 1):
            text, action = st.columns([3, 1.2], vertical_alignment='center')
            text.text(f'{number}. {pair["a"]["label"]} → {pair["b"]["label"]}')
            action.button('Remove', key=f'remove_relationship_example:{pair["id"]}',
                          on_click=remove_example, args=(pair['id'],))
        st.caption(f'{len(pairs)} / 10 example pairs · offsets are averaged equally')
    else:
        source_a = choose('A', 'a')
        source_b = choose('B', 'b', 1)
        pairs = [{'a': capture(source_a), 'b': capture(source_b)}]
    target = choose('C (apply relationship to)', 'c', 2)
    expected = choose('D (compare against)' if mode == 'Compare relationships' else 'Expected D (optional)',
                       'd' if mode == 'Compare relationships' else 'expected', 3,
                       optional=mode != 'Compare relationships')
    if st.button('Show relationship', key='show_relationship', type='primary', disabled=not pairs, use_container_width=True):
        try:
            target_data = capture(target)
            expected_data = capture(expected) if expected is not None else None
            metrics = analyse_pairs([(p['a']['values'], p['b']['values']) for p in pairs],
                                     target_data['values'], expected_data['values'] if expected_data else None)
            serial = st.session_state.get('relationship_serial', 0) + 1
            st.session_state.relationship_serial = serial
            excluded = {item['key'] for p in pairs for item in [p['a'], p['b']] if isinstance(item['key'], int)}
            if isinstance(target, int):
                excluded.add(target)
            st.session_state[active_key] = {'id': serial, 'mode': mode, 'pairs': list(pairs),
                                           'target': target_data, 'expected': expected_data,
                                           'metrics': metrics, 'exclude_ids': sorted(excluded)}
            st.session_state.pending_chart_view = 'Relationship 3D'
            st.rerun()
        except ValueError as e:
            st.error(str(e))
    analysis = st.session_state.get(active_key)
    if analysis is None:
        return None
    metrics = analysis['metrics']
    st.caption('Full-dimensional scores for the last shown relationship. No projection is used for these calculations.')
    if 'direction_cosine' in metrics:
        left, right = st.columns(2)
        left.metric('Offset cosine', f'{metrics["direction_cosine"]:.4f}' if metrics['direction_cosine'] is not None else 'Undefined')
        right.metric('Angle', f'{metrics["angle_degrees"]:.1f}°' if metrics['angle_degrees'] is not None else 'Undefined')
        left.metric('Length ratio D−C / offset', f'{metrics["length_ratio"]:.4f}' if metrics['length_ratio'] is not None else 'Undefined')
        right.metric('Prediction error (L2)', f'{metrics["prediction_error"]:.5f}')
    else:
        st.metric('Mean offset length', f'{metrics["offset_length"]:.5f}')
    if len(analysis['pairs']) > 1:
        agreement = metrics['example_agreement']
        st.caption(f'Example offsets vs average: mean cosine {agreement:.4f}' if agreement is not None else 'Example direction agreement is undefined.')
    if metrics['offset_length'] == 0 or metrics.get('observed_length') == 0:
        st.info('A zero-length relationship has no direction. Direction cosine and angle are undefined for it.')
    if st.button('Find analogy answers', key='find_analogy_answers', use_container_width=True):
        st.session_state[f'neighbour_mode:{scope}'] = 'Analogy answers'
        st.session_state[f'run_analogy:{scope}'] = True
    if st.button('Add prediction as calc', key='save_relationship_prediction', disabled=len(vectors) >= 20, use_container_width=True):
        save_calculation(metrics['prediction'], f'analogy prediction #{analysis["id"]}')
    if st.button('Add mean offset as calc', key='save_relationship_offset', disabled=len(vectors) >= 20, use_container_width=True):
        save_calculation(metrics['offset'], f'mean relationship offset #{analysis["id"]}')
    st.caption('Prediction = C + mean(B − A). Analogy searches exclude the input vocabulary tokens. Scores use stored, unnormalised input vectors.')
    return analysis