import os
import re
import glob
import pickle
import numpy as np
from scipy.ndimage import laplace

class VisualCollapseChatbot:
    def __init__(self, model_path="vct_brain.pkl"):
        self.grid_size = 8
        self.model_path = model_path
        self.kb = []             # Knowledge Base: list of (glyph, source_sentence)
        self.inverse_atlas = {}  # Spatial inverse hash lookup: (x, y, color_channel) -> list of words

    def _clean_text(self, text):
        """Standardizes input text by purging markdown structures and normalizing spacing."""
        # Remove markdown link syntax
        text = re.sub(r'\[.*?\]\(.*?\)', '', text) 
        # Safely strip common formatting tokens individually
        for char in ['#', '*', '`', '_', '-', '狂']:
            text = text.replace(char, '')
        # Normalize whitespace
        text = re.sub(r'\s+', ' ', text)           
        return text.strip()

    def encode_to_glyph(self, text):
        """Paints text onto an 8x8x3 RGB matrix using deterministic spatial-semantic hashing."""
        glyph = np.zeros((8, 8, 3), dtype=np.float32)
        words = text.lower().split()
        
        for i, word in enumerate(words):
            # Spatial Mapping via coordinates
            x = hash(word) % 8
            y = hash(word + str(i)) % 8
            
            # Semantic Channel Mapping (Verbs/Nouns/Context modifiers)
            r = (hash(word + "verb") % 100) / 100.0
            g = (hash(word + "noun") % 100) / 100.0
            b = (hash(word + "adj") % 100) / 100.0
            
            # Relationship bleeding through neighbor proximity
            for dx in range(-1, 2):
                for dy in range(-1, 2):
                    nx, ny = (x + dx) % 8, (y + dy) % 8
                    glyph[ny, nx] += np.array([r, g, b], dtype=np.float32) * 0.5
                    
        return np.clip(glyph, 0, 1)

    def train_on_folder(self, folder_path):
        """Processes a folder of markdown files and catalogs spatial-chromatic nodes."""
        print(f"[*] Analyzing ecosystem directory: {folder_path}")
        search_path = os.path.join(folder_path, "**", "*.md")
        md_files = glob.glob(search_path, recursive=True)
        
        if not md_files:
            print(f"[!] Warning: No markdown (.md) documents located in {folder_path}.")
            return

        for file_path in md_files:
            print(f" -> Absorbing text space: {os.path.basename(file_path)}")
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                
            clean_content = self._clean_text(content)
            sentences = re.split(r'(?<=[.!?])\s+', clean_content)
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence.split()) < 3: 
                    continue # Ignore shorthand fragments
                
                # Register concept in raw memory space
                glyph = self.encode_to_glyph(sentence)
                self.kb.append((glyph, sentence))
                
                # Document real coordinates for the Inverse Hash Atlas
                words = sentence.lower().split()
                for i, word in enumerate(words):
                    x = hash(word) % 8
                    y = hash(word + str(i)) % 8
                    
                    # Track dominant concept color channel to map dictionary index
                    r = hash(word + "verb") % 100
                    g = hash(word + "noun") % 100
                    b = hash(word + "adj") % 100
                    channels = [r, g, b]
                    dom_channel = int(np.argmax(channels))
                    
                    atlas_key = (x, y, dom_channel)
                    if atlas_key not in self.inverse_atlas:
                        self.inverse_atlas[atlas_key] = set()
                    self.inverse_atlas[atlas_key].add(word)

        # Convert sets to sorted lists to ensure persistent deterministic serialization
        for key in self.inverse_atlas:
            self.inverse_atlas[key] = sorted(list(self.inverse_atlas[key]))

        print(f"[+] Sync Complete. Ingested {len(self.kb)} concepts.")
        print(f"[+] Inverse Hash Atlas pinned {len(self.inverse_atlas)} active structural indices.")
        self.save_model()

    def save_model(self):
        """Saves weights and linguistic maps to disk."""
        try:
            with open(self.model_path, 'wb') as f:
                pickle.dump({'kb': self.kb, 'inverse_atlas': self.inverse_atlas}, f)
            print(f"[+] Model preserved successfully: '{self.model_path}'")
        except Exception as e:
            print(f"[!] Failed to save weights: {e}")

    def load_model(self):
        """Restores network structures from a previous training phase."""
        if os.path.exists(self.model_path):
            try:
                with open(self.model_path, 'rb') as f:
                    data = pickle.load(f)
                    self.kb = data['kb']
                    self.inverse_atlas = data['inverse_atlas']
                print(f"[+] Weights initialized. Model loaded: '{self.model_path}' ({len(self.kb)} concepts).")
                return True
            except Exception as e:
                print(f"[!] Corrupted model file found: {e}. Defaulting to empty initialization.")
        else:
            print(f"[-] No model file named '{self.model_path}' found. Ready for initial ingestion.")
        return False

    def think_ode_step(self, glyph, t):
        """Simulates 2D Reaction-Diffusion equations to purify concepts."""
        new_glyph = glyph.copy()
        
        # 1. Spatial Diffusion: Colors bleed logically into matching fields
        for c in range(3): 
            new_glyph[:, :, c] += 0.1 * laplace(glyph[:, :, c])
            
        # 2. Reaction: Mixed states are driven down by calculating local purity
        max_c = np.max(new_glyph, axis=2, keepdims=True)
        min_c = np.min(new_glyph, axis=2, keepdims=True)
        purity = (max_c - min_c) / (max_c + 0.001)
        
        new_glyph *= purity
        return np.clip(new_glyph, 0, 1)

    def chromatic_overlap(self, g1, g2):
        """Calculates exact intersecting volumes of two active spatial color grids."""
        p1 = g1.reshape(-1, 3)
        p2 = g2.reshape(-1, 3)
        mask = (np.sum(p1, axis=1) > 0.1) & (np.sum(p2, axis=1) > 0.1)
        
        if not np.any(mask): 
            return 0.0
            
        dot = np.sum(p1[mask] * p2[mask])
        norm = np.linalg.norm(p1[mask]) * np.linalg.norm(p2[mask])
        return dot / norm if norm > 0 else 0.0

    def generative_query(self, prompt, top_n=3, steps=25):
        """Superimposes relevant vectors and uses fluid physics to collapse to a response."""
        prompt_glyph = self.encode_to_glyph(prompt)
        
        if not self.kb:
            return "The canvas is empty. Please run training on markdown text documents.", prompt_glyph

        scores = []
        for kb_glyph, text in self.kb:
            score = self.chromatic_overlap(prompt_glyph, kb_glyph)
            scores.append((score, kb_glyph, text))
            
        scores.sort(key=lambda x: x[0], reverse=True)
        top_candidates = scores[:top_n]
        
        if not top_candidates or top_candidates[0][0] == 0:
            return "The internal matrix remains dark. No corresponding conceptual intersections.", prompt_glyph

        # Phase 1: Superposition
        superposition = np.zeros((8, 8, 3), dtype=np.float32)
        for score, glyph, _ in top_candidates:
            superposition += glyph * (score if score > 0 else 0.1)
            
        # Phase 2: Reaction-Annihilation Engine Loop
        for t in range(1, steps + 1):
            superposition = self.think_ode_step(superposition, t)
            if np.var(superposition) < 0.0005:
                break
                
        # Phase 3: Text Decoding via Inverse Atlas mapping
        generated_text = self.decode_glyph_to_text(superposition)
        return generated_text, superposition

    def decode_glyph_to_text(self, glyph):
        """Converts structural spatial states back to standard readable prose strings."""
        words = []
        temp_glyph = glyph.copy()
        
        for _ in range(15): # Enforce maximum length of generated thought string
            flat_idx = np.argmax(np.sum(temp_glyph, axis=2))
            y, x = divmod(flat_idx, 8)
            
            if np.sum(temp_glyph[y, x]) < 0.15:
                break # Canvas energy expired
                
            color_idx = int(np.argmax(temp_glyph[y, x]))
            atlas_key = (x, y, color_idx)
            
            # Fetch corresponding structural tokens from the real trained corpus map
            if atlas_key in self.inverse_atlas:
                matched_words = self.inverse_atlas[atlas_key]
                # Deterministically pivot through available options to keep output steady
                word = matched_words[hash(str(x) + str(y)) % len(matched_words)]
                words.append(word)
            
            # Local decay loop to ensure sequential progression
            temp_glyph[y, x] *= 0.0
            if x < 7: 
                temp_glyph[y, x + 1] *= 1.2 # Shift attention forward chronologically
                
        if not words:
            return "..."
        return " ".join(words).capitalize() + "."

# --- Conversational Interface Core ---
def main():
    BOOK_FOLDER = "./my_books"
    MODEL_FILE = "vct_brain.pkl"
    
    bot = VisualCollapseChatbot(model_path=MODEL_FILE)
    
    # 1. Attempt Workspace Setup or Recovery
    has_model = bot.load_model()
    
    if not has_model:
        os.makedirs(BOOK_FOLDER, exist_ok=True)
        # Seed example database if entirely vacant
        if not os.listdir(BOOK_FOLDER):
            print(f"[*] Initializing sample training note in '{BOOK_FOLDER}'...")
            with open(os.path.join(BOOK_FOLDER, "ai_theory.md"), "w", encoding='utf-8') as f:
                f.write("# Visual Collapse Paradigm\n")
                f.write("A quiet cat sleeps soundly under the warm wooden table.\n")
                f.write("The active dog runs fast outside across the green grass.\n")
                f.write("An analytical operator programs computers inside the server facility.\n")
                
        bot.train_on_folder(BOOK_FOLDER)

    # 2. Start Live Loop
    print("\n" + "="*60)
    print(" VCT 8x8x3 SPATIAL GEN-AI CHAT INTERFACE ACTIVATED ")
    print(" Commands: '/train' to re-ingest folder | '/exit' to close")
    print("="*60 + "\n")
    
    while True:
        try:
            user_input = input("User > ").strip()
            if not user_input:
                continue
                
            if user_input.lower() == '/exit':
                print("[*] Dissolving network matrix. Goodbye.")
                break
                
            if user_input.lower() == '/train':
                bot.kb = []
                bot.inverse_atlas = {}
                bot.train_on_folder(BOOK_FOLDER)
                continue
            
            # Process prompt through fluid dynamics engine
            reply, _ = bot.generative_query(user_input, top_n=2)
            print(f"VCT > {reply}\n")
            
        except (KeyboardInterrupt, EOFError):
            print("\n[*] Exiting structural runtime cleanly.")
            break

if __name__ == "__main__":
    main()