import csv
import io
from pathlib import Path

import numpy as np
import plotly.graph_objects as go
import streamlit as st
from plotly.colors import qualitative

from embeddings import Model, embedding_name, load_tokenizer, nearest_tokens, search_embeddings
from relationships import relationship_figure
from relationship_ui import render_relationship_controls
from model_picker import show_model_picker
from hub_ui import render_hub_loader
from upload_ui import render_upload_loader
from app_environment import hosted_mode, remove_session_directory, UPLOAD_ROOT, PRIVATE_CACHE_ROOT


st.set_page_config(page_title='Token Visualiser', page_icon='🔬', layout='wide', initial_sidebar_state='collapsed')
hosted = hosted_mode()
if st.session_state.pop('clear_hub_token', False):
    st.session_state.pop('hf_token', None)
if 'pending_chart_view' in st.session_state:
    st.session_state.chart_view = st.session_state.pop('pending_chart_view')
if 'pending_model' in st.session_state:
    pending_model = st.session_state.pop('pending_model')
    if not hosted:
        st.session_state.model_location = pending_model


def clear_model():
    # Keep the last browsing location, but release model data and all analysis state.
    model = st.session_state.get('model')
    if model is not None:
        model.close()
    for key, root in [('upload_directory', UPLOAD_ROOT), ('private_hub_cache', PRIVATE_CACHE_ROOT)]:
        if key in st.session_state:
            try:
                remove_session_directory(st.session_state[key], root)
            except OSError:
                pass  # Files still held by a disconnected worker disappear on restart.
    for key in list(st.session_state):
        if key not in {'picker_directory', 'picker_address'}:
            del st.session_state[key]


def activate_model(location, tokenizer_location=None, display_name=None):
    model = Model(location)
    try:
        p = model.path
        auto = (p if p.is_dir() else p.parent) / 'tokenizer.json'
        selected_tokenizer = tokenizer_location or (auto if auto.exists() else None)
        tok = load_tokenizer(selected_tokenizer) if selected_tokenizer else None
    except Exception:
        model.close()
        raise
    st.session_state.model = model
    st.session_state.tok = tok
    st.session_state.loaded = str(model.path)
    st.session_state.model_display_name = display_name or model.path.name
    st.session_state.nearest_results = {}
    st.session_state.clear_hub_token = True


title_col, action_col = st.columns([6, 1], vertical_alignment='center')
if 'model' in st.session_state:
    title_col.subheader('Token Visualiser')
else:
    title_col.title('Token Visualiser')
if 'model' in st.session_state:
    action_col.button('Clear model', key='clear_model', on_click=clear_model, use_container_width=True)
if 'model' not in st.session_state:
    st.caption('Explore the numbers inside a token. Models run in this Space · No GPU required' if hosted else 'Explore the numbers inside a token. Local model files · No GPU required')

if 'model' not in st.session_state:
    _, loader, _ = st.columns([1, 3, 1])
    with loader, st.container(border=True):
        st.header('Load a model')
        st.write('Choose a suggested model, paste a Hugging Face link, or upload weights. Models load on this Space.'
                 if hosted else 'Open a local model or download weights from Hugging Face to explore token embeddings.')
        if hosted:
            hub_tab, local_tab = st.tabs(['Hugging Face', 'Upload files'])
        else:
            local_tab, hub_tab = st.tabs(['Local files', 'Hugging Face'])
        with local_tab:
            if hosted:
                render_upload_loader(activate_model)
        if not hosted:
            with local_tab:
                if st.button('Open model…', key='open_model', use_container_width=True):
                    st.session_state.picker_open = True
                if st.session_state.get('picker_open', False):
                    show_model_picker()
                location = st.text_input('Model file or folder', placeholder=r'C:\Models\gemma.gguf', key='model_location')
                with st.expander('Tokenizer (optional)'):
                    tokenizer_path = st.text_input('Tokenizer JSON or folder', key='tokenizer_location',
                        help='Use the tokenizer from the exact same model. Enables text tokenization without added special tokens. Detected automatically beside the model when available.')
                load = st.button('Load model', key='load_model', type='primary', disabled=not location.strip(), use_container_width=True)
                st.caption('GGUF · F32/F16/BF16 safetensors · Sharded model folders')
                st.caption('Files stay on this computer.')
                if load:
                    try:
                        with st.spinner('Reading model metadata…'):
                            activate_model(location.strip().strip('"'), tokenizer_path)
                        st.rerun()
                    except Exception as e:
                        st.error(f'Could not load model: {e}')
        with hub_tab:
            render_hub_loader(activate_model)
    st.stop()

model = st.session_state.model
st.markdown('<style>' + Path(__file__).with_name('viewer.css').read_text(encoding='utf-8') + '</style>', unsafe_allow_html=True)
chart_pane = st.container(key='chart_pane', height='stretch', border=True)
chart_slot = chart_pane.empty()
with st.container(key='controls_pane'):
    with st.container(key='control_columns'):
        token_panel, neighbour_panel, calculation_panel, details_panel = st.columns([1.15, 1.6, 1.15, 1.05], gap='medium')

with details_panel:
    st.subheader('Chart & model')
    st.caption(st.session_state.get('model_display_name', Path(st.session_state.loaded).name))
    names = sorted(model.tensors, key=lambda n: (not embedding_name(n), n))
    name = st.selectbox('Embedding tensor', names)
    if not embedding_name(name):
        st.warning('This tensor may not contain input token embeddings.')
    chart_view = st.selectbox('Chart view', ['Dimensions', 'Relationship 2D', 'Relationship 3D'], key='chart_view')
    style = st.radio('Chart style', ['Line', 'Scatter'], horizontal=True)
shape = model.tensors[name][2]
vocab = model.vocab
tok = st.session_state.tok
get_row = lambda i: model.row(name, i)
read_rows = lambda start, stop: model.rows(name, start, stop)
with details_panel:
    a, b = st.columns(2)
    a.metric('Vocabulary rows', f'{shape[0]:,}')
    b.metric('Values per token', f'{shape[1]:,}')
    st.caption('Stored input embeddings, before layer processing. Drag the chart to zoom; double-click to reset. Click a legend entry to hide or show a token.')

with token_panel:
    ids = []
    labels = {}
    comparison_scope = f'{st.session_state.loaded}:{name}'
    calculation_key = f'calculations:{comparison_scope}'
    st.subheader('Compare tokens')
    st.caption('Search, select and add tokens. Each curve has its own colour.')
    if tok is not None:
        text = st.text_input('Type text', 'hello world')
        try:
            encoded = tok.encode(text, add_special_tokens=False)
            st.caption('Tokenized pieces: ' + ' · '.join(f'{piece!r} (ID {i})' for piece, i in zip(encoded.tokens, encoded.ids)))
            choices = list(range(len(encoded.ids)))
            positions = st.multiselect('Select tokens to plot', choices, default=choices[:2],
                                       format_func=lambda p: f'{p + 1}: {encoded.tokens[p]!r} · ID {encoded.ids[p]}')
            for p in positions:
                i = encoded.ids[p]
                if i not in ids:
                    ids.append(i)
                labels[i] = f'{encoded.tokens[p]} · {i}'
        except Exception as e:
            st.error(f'Tokenization failed: {e}')
    elif vocab is not None:
        query = st.text_input('Find a vocabulary token', 'hello', help='Searches stored token spellings, not text tokenization. Spaces may appear as ▁ or Ġ.')
        matches = [i for i, word in enumerate(vocab) if query in word][:200] if query else list(range(min(200, len(vocab))))
        matches.sort(key=lambda i: (vocab[i] != query, len(vocab[i]), i))
        # The chart selection lives independently of changing search-result widgets.
        selection_key = f'chart_tokens:{comparison_scope}'
        if selection_key not in st.session_state:
            st.session_state[selection_key] = matches[:1]
        selected = [i for i in st.session_state[selection_key] if 0 <= i < min(shape[0], len(vocab))]
        result, add = st.columns([2.3, 1.7], vertical_alignment='bottom')
        candidate = result.selectbox('Matching token', matches, key=f'match:{comparison_scope}:{query}',
                                     format_func=lambda i: f'{vocab[i]!r} · ID {i}')
        if add.button('Add to chart', type='primary', disabled=candidate is None or candidate in selected or len(selected) + len(st.session_state.get(calculation_key, [])) >= 20):
            selected.append(candidate)
            st.session_state[selection_key] = selected
        ids = st.multiselect('Tokens on chart', selected, default=selected,
                            key=f'chart_selection:{comparison_scope}:{tuple(selected)}',
                            format_func=lambda i: f'{vocab[i]!r} · ID {i}')
        if ids != selected:
            st.session_state[selection_key] = ids
            st.rerun()
        st.caption('Search for another token and click Add to chart to compare it. Remove a token with the × beside its name.')
        labels = {i: f'{vocab[i]} · {i}' for i in ids}
        if not matches:
            st.info('No matching vocabulary entries. Try a substring or enter an ID below.')
    else:
        st.info('Add the matching tokenizer.json to type text. You can inspect any token by ID below.')

    ids = list(ids)
    manual = st.text_input('Additional token IDs (comma-separated)', placeholder='0, 42, 100', key='manual_ids')
    try:
        for part in manual.split(','):
            if part.strip():
                i = int(part.strip())
                if not 0 <= i < shape[0]:
                    raise ValueError(f'ID {i} is outside this tensor’s row range.')
                if i not in ids:
                    ids.append(i)
                labels.setdefault(i, f'ID {i}')
    except ValueError as e:
        st.error(str(e))
        st.stop()
    def token_label(i):
        if vocab is not None and i < len(vocab):
            return f'{vocab[i]!r} · ID {i}'
        if tok is not None:
            piece = tok.id_to_token(i)
            if piece is not None:
                return f'{piece!r} · ID {i}'
        return f'ID {i}'


    extra_key = f'added_neighbours:{comparison_scope}'
    extra_ids = [i for i in st.session_state.get(extra_key, []) if i not in ids and 0 <= i < shape[0]]
    if extra_ids:
        kept = st.multiselect('Added neighbours on chart', extra_ids, default=extra_ids,
                              format_func=token_label, key=f'added_neighbour_selection:{comparison_scope}:{tuple(extra_ids)}')
        if kept != extra_ids:
            st.session_state[extra_key] = kept
            st.rerun()
        for i in kept:
            ids.append(i)
            labels[i] = token_label(i)
    if not ids and not st.session_state.get(calculation_key):
        chart_slot.info('Select a token in the controls below to see its embedding values.')
        st.stop()
    if len(ids) > 20:
        st.warning('Select at most 20 tokens for a readable comparison.')
        st.stop()


    def add_neighbour(i):
        added = list(st.session_state.get(extra_key, []))
        if i not in added:
            added.append(i)
        st.session_state[extra_key] = added


try:
    vectors = {i: get_row(i) for i in ids}
except Exception as e:
    chart_slot.error(f'Could not read embedding: {e}')
    st.stop()

calculations = st.session_state.setdefault(calculation_key, [])
for calculation in calculations:
    vectors[calculation['id']] = calculation['values']
    labels[calculation['id']] = calculation['label']


def remove_calculation(calculation_id):
    st.session_state[calculation_key] = [c for c in st.session_state[calculation_key] if c['id'] != calculation_id]


def operand_label(key):
    return labels[key] if isinstance(key, str) else token_label(key)


def expression_label(key):
    if isinstance(key, str):
        return '(' + labels[key].removeprefix('calc: ') + ')'
    if vocab is not None and key < len(vocab):
        return repr(vocab[key])
    if tok is not None and tok.id_to_token(key) is not None:
        return repr(tok.id_to_token(key))
    return f'ID {key}'


def save_calculation(values, label):
    label = label if label.startswith('calc: ') else 'calc: ' + label
    serial = st.session_state.get('calculation_serial', 0) + 1
    while label in labels.values():
        label += f' [{serial}]'
    st.session_state.calculation_serial = serial
    calculations.append({'id': f'calc:{serial}', 'label': label, 'values': np.asarray(values).copy()})
    st.rerun()


with calculation_panel:
    st.subheader('Vector tools')
    arithmetic_tab, relationship_tab = st.tabs(['Calculations', 'Relationships'])
with arithmetic_tab:
    st.caption('Calculate A − B or A + B, dimension by dimension. Results become derived curves labelled calc:.')
    operands = list(vectors)
    left_operand = st.selectbox('Vector A', operands, format_func=operand_label, key=f'calc_a:{comparison_scope}')
    operation = st.selectbox('Operation', ['Subtract (A − B)', 'Add (A + B)'], key=f'calc_operation:{comparison_scope}')
    right_operand = st.selectbox('Vector B', operands, index=min(1, len(operands) - 1),
                                 format_func=operand_label, key=f'calc_b:{comparison_scope}')
    result_name = st.text_input('Result name (optional)', placeholder='e.g. hello minus world', key=f'calc_name:{comparison_scope}')
    if st.button('Calculate & add to chart', key='add_calculation', type='primary',
                 disabled=len(vectors) >= 20, use_container_width=True):
        subtract = operation.startswith('Subtract')
        values = vectors[left_operand] - vectors[right_operand] if subtract else vectors[left_operand] + vectors[right_operand]
        symbol = '−' if subtract else '+'
        expression = f'{expression_label(left_operand)} {symbol} {expression_label(right_operand)}'
        label = 'calc: ' + (result_name.strip() or expression)
        save_calculation(values, label)
    if calculations:
        st.caption('Calculated curves stay available when you change the selected tokens. They are included in CSV exports.')
        for calculation in calculations:
            label_col, remove_col = st.columns([3, 1.2], vertical_alignment='center')
            label_col.text(calculation['label'])
            remove_col.button('Remove', key=f'remove_calculation:{calculation["id"]}',
                              on_click=remove_calculation, args=(calculation['id'],))
    if len(vectors) >= 20:
        st.caption('The chart supports up to 20 token and calculated curves in total.')

with relationship_tab:
    analysis = render_relationship_controls(vectors, operand_label, comparison_scope, save_calculation)

if len(vectors) > 20:
    token_panel.warning('Remove a token or calculated curve to keep the chart within its 20-curve limit.')
    st.stop()
chart_ids = list(vectors)
colour_key = f'token_colours:{comparison_scope}'
colours = st.session_state.setdefault(colour_key, {})
# Retain colours across searches, keeping every currently selected token distinct.
used = set()
for i in chart_ids:
    if i not in colours or colours[i] in used:
        colours[i] = next(c for c in qualitative.Light24 if c not in used and
                         c not in {colours.get(other) for other in chart_ids if other != i})
    used.add(colours[i])
details_panel.caption(f'Tokens: {len(ids)} · Calculations: {len(calculations)}')
if chart_view == 'Dimensions':
    fig = go.Figure()
    for i, values in vectors.items():
        fig.add_trace(go.Scattergl(x=np.arange(len(values)), y=values,
                                  mode='lines' if style == 'Line' else 'markers', name=labels[i],
                                  line=dict(color=colours[i], width=2 if isinstance(i, str) else 1.5,
                                            dash='dash' if isinstance(i, str) else 'solid'),
                                  marker=dict(color=colours[i], size=4),
                                  hovertemplate='Dimension %{x}<br>Value %{y:.7f}<extra>%{fullData.name}</extra>'))
    fig.update_layout(autosize=True, xaxis_title='Embedding dimension (zero-based)', yaxis_title='Stored embedding value',
                      template='plotly_dark', hovermode='x', margin=dict(l=20, r=20, t=20, b=30),
                      legend=dict(orientation='h', y=1.02, yanchor='bottom', title_text='Tokens'), showlegend=True,
                      uirevision=f'{comparison_scope}:{style}')
    fig.add_hline(y=0, line_width=1, line_color='gray')
    chart_slot.plotly_chart(fig, width='stretch', height='stretch', config={'scrollZoom': True}, key='embedding_chart')

elif analysis is None:
    chart_slot.info('Open Relationships in Vector tools, choose your inputs, then click Show relationship.')
else:
    fig = relationship_figure(analysis, 2 if chart_view == 'Relationship 2D' else 3, vectors, labels)
    chart_slot.plotly_chart(fig, width='stretch', height='stretch', key='embedding_chart', config={'scrollZoom': True})

with neighbour_panel:
    st.subheader('Closest tokens')
    search_mode = st.selectbox('Search mode', ['Nearest tokens', 'Analogy answers'], key=f'neighbour_mode:{comparison_scope}')
    results = st.session_state.setdefault('nearest_results', {})
    eligible = True
    report_id = None
    if search_mode == 'Nearest tokens':
        source_id = st.selectbox('Find neighbours of', chart_ids, format_func=operand_label,
                                 key=f'neighbour_source:{comparison_scope}')
        query_vector = vectors[source_id]
        excluded = [source_id] if isinstance(source_id, int) else []
        result_key = ('ranked', comparison_scope, source_id)
        row_key = str(source_id)
        search_label = 'Find closest tokens'
        st.caption('Cosine similarity uses the full vector, including calculated queries.')
    elif analysis is None:
        st.info('Show a relationship in Vector tools before searching for analogy answers.')
        eligible = False
    else:
        query_vector = analysis['metrics']['prediction']
        excluded = analysis['exclude_ids']
        expected = analysis['expected']
        report_id = expected['key'] if expected and isinstance(expected['key'], int) else None
        result_key = ('analogy', comparison_scope, analysis['id'])
        row_key = f'analogy:{analysis["id"]}'
        search_label = 'Refresh answers'
        st.caption('Search near C + mean(B − A), excluding the input tokens. Results refer to the last shown relationship.')
    if eligible:
        count_col, button_col = st.columns([1, 2], vertical_alignment='bottom')
        neighbour_count = count_col.number_input('Results', min_value=1, max_value=50, value=10, step=1)
        search = button_col.button(search_label, type='primary', key='find_neighbours')
        requested = st.session_state.pop(f'run_analogy:{comparison_scope}', False) if search_mode == 'Analogy answers' else False
        previous_search = any(isinstance(key, tuple) and len(key) == 3 and key[:2] == (comparison_scope, source_id)
                              for key in results) if search_mode == 'Nearest tokens' else False
        if search or requested or (result_key not in results and previous_search):
            bar = st.progress(0, text='Scanning embedding vectors…')
            try:
                found = search_embeddings(read_rows, shape[0], query_vector, count=50, exclude_ids=excluded, report_id=report_id,
                    progress=lambda done, total: bar.progress(done / total, text=f'Scanning embeddings: {done:,} / {total:,}'))
                results[result_key] = found if search_mode == 'Analogy answers' else found['neighbours']
            except Exception as e:
                st.error(f'Could not find neighbours: {e}')
            finally:
                bar.empty()
        if result_key in results:
            found = results[result_key]
            if search_mode == 'Analogy answers':
                neighbours = found['neighbours'][:int(neighbour_count)]
                if report_id is not None:
                    if found['expected_rank'] is not None:
                        st.caption(f'Expected {token_label(report_id)}: rank {found["expected_rank"]:,} of {found["valid_count"]:,}, cosine {found["expected_score"]:.6f}.')
                    else:
                        st.caption('Expected D is excluded as an input token, or has a zero/nonfinite embedding.')
            else:
                neighbours = found[:int(neighbour_count)]
            if neighbours:
                st.caption(f'Showing {len(neighbours)} neighbours. Change Results to show more or fewer without searching again.')
                with st.container(border=False):
                    header = st.columns([.35, 1.7, 1, 1.8])
                    for column, title in zip(header, ['#', 'Token', 'Cosine', '']):
                        if title:
                            column.markdown(f'**{title}**')
                    for rank, (i, score) in enumerate(neighbours, 1):
                        rank_col, token_col, score_col, action_col = st.columns([.35, 1.7, 1, 1.8], vertical_alignment='center')
                        rank_col.text(str(rank))
                        token_col.text(token_label(i))
                        score_col.write(f'{score:.6f}')
                        action_col.button('On chart' if i in ids else 'Add to chart',
                            key=f'add_neighbour:{comparison_scope}:{row_key}:{i}',
                            disabled=i in ids or len(vectors) >= 20, on_click=add_neighbour, args=(i,), use_container_width=True)
                if len(vectors) >= 20:
                    st.caption('The chart has reached its 20-curve limit. Remove a curve before adding another.')
                st.caption('Cosine similarity ranges from −1 to 1. Nearby tokens can reflect spelling, boundaries or usage as well as meaning.')
            else:
                st.info('No other eligible tokens with finite, nonzero embeddings were found.')
with details_panel:
    st.dataframe([{'Token': labels[i], 'Min': float(v.min()), 'Max': float(v.max()),
                   'Mean': float(v.mean()), 'Std dev': float(v.std()), 'L2 norm': float(np.linalg.norm(v))}
                  for i, v in vectors.items()], use_container_width=True, hide_index=True)
    output = io.StringIO()
    writer = csv.writer(output)
    writer.writerow(['dimension'] + [labels[i] for i in chart_ids])
    writer.writerows([d] + [float(vectors[i][d]) for i in chart_ids] for d in range(shape[1]))
    st.download_button('Download values as CSV', output.getvalue(), 'token_embeddings.csv', 'text/csv')
    with st.expander('Inspect exact values'):
        st.dataframe({'dimension': np.arange(shape[1]), **{labels[i]: vectors[i] for i in chart_ids}}, use_container_width=True)