import numpy as np
from sentence_transformers import SentenceTransformer

class ContextGraph:
    def __init__(self):
        self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    def encode(self, text):
        return self.model.encode(text)

    def connect(self, memory_texts, threshold=0.35):
        """
        Create relationships between memory nodes based on vector similarity.
        """
        vectors = [self.encode(t) for t in memory_texts]
        relationships = []

        for i, vec1 in enumerate(vectors):
            for j, vec2 in enumerate(vectors):
                if i != j:
                    similarity = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
                    if similarity >= threshold:
                        relationships.append((memory_texts[i], memory_texts[j], float(similarity)))

        return relationships

    def score_context(self, query, memories):
        """
        Score stored memories against the new query.
        """
        query_vec = self.encode(query)
        scored = []
        for mem, vec in memories:
            score = np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec))
            scored.append((mem, float(score)))

        return sorted(scored, key=lambda x: x[1], reverse=True)