import os
import json
from collections import defaultdict

BRAINMAP_FILE = "data/memory/brainmap.json"


def load_brainmap():
    if os.path.exists(BRAINMAP_FILE):
        with open(BRAINMAP_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}


def save_brainmap(map_data):
    with open(BRAINMAP_FILE, "w", encoding="utf-8") as f:
        json.dump(map_data, f, indent=2)


brain_map = load_brainmap()


def update_brainmap(words):
    for i, word in enumerate(words):
        for j in range(i+1, len(words)):
            w1 = word
            w2 = words[j]
            if w1 == w2:
                continue
            if w1 not in brain_map:
                brain_map[w1] = {}
            if w2 not in brain_map[w1]:
                brain_map[w1][w2] = 0
            brain_map[w1][w2] += 1
    save_brainmap(brain_map)


def get_brainmap():
    return brain_map