import numpy as np
from flask import Blueprint, render_template, jsonify, request

bp = Blueprint(
    'brain_map',
    __name__,
    template_folder='templates',
    static_folder='static',
)

# Will be injected from body.py
system = None


@bp.route('/graph')
def graph():
    return render_template('graph.html')


@bp.route('/data')
def data():
    if system is None:
        return jsonify({"nodes": [], "edges": []})

    # 1) embeddings → cosine sims
    emb = system.brain.token_emb.weight.detach().cpu().numpy()
    N = emb.shape[0]
    norms = np.linalg.norm(emb, axis=1, keepdims=True)
    emb_norm = emb / (norms + 1e-8)
    sim = emb_norm.dot(emb_norm.T)

    # 2) filters
    min_degree = int(request.args.get('min_degree', 1))
    max_nodes  = int(request.args.get('max_nodes', 200))

    # 3) valid tokens
    items = [(tok, idx) for tok, idx in system.sensory.stoi.items() if 0 <= idx < N]

    # 4) build undirected unique pairs & degree counts
    deg = {int(idx): 0 for _, idx in items}
    unique_pairs = set()
    for _, raw_i in items:
        i = int(raw_i)
        count = 0
        for raw_j in np.argsort(-sim[i]):
            j = int(raw_j)
            if j == i or j not in deg:
                continue
            pair = (min(i, j), max(i, j))
            if pair in unique_pairs:
                continue
            unique_pairs.add(pair)
            deg[i] += 1
            deg[j] += 1
            count += 1
            if count >= 3:
                break

    # 5) filter & cap nodes
    filtered = [(tok, idx) for tok, idx in items if deg[int(idx)] >= min_degree]
    filtered.sort(key=lambda x: (-deg[int(x[1])], int(x[1])))
    subset = filtered[:max_nodes]
    subset_ids = {int(idx) for _, idx in subset}

    # 6) build nodes with HSL coloring
    max_deg = max((deg[idx] for _, idx in subset), default=1)
    nodes = []
    for tok, raw_idx in subset:
        idx = int(raw_idx)
        d = deg[idx]
        hue = int((1 - d / max_deg) * 240)
        nodes.append({
            'id': idx,
            'label': tok,
            'color': {
                'background': f'hsl({hue},80%,40%)',
                'border':     f'hsl({hue},60%,30%)',
                'highlight': {
                    'background': f'hsl({hue},100%,50%)',
                    'border':     f'hsl({hue},80%,40%)'
                }
            }
        })

    # 7) build edges
    edges = [
        {'from': a, 'to': b}
        for (a, b) in unique_pairs
        if a in subset_ids and b in subset_ids
    ]

    return jsonify({'nodes': nodes, 'edges': edges})