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

class VisualCollapseAI:
    def __init__(self):
        self.grid_size = 8
        self.kb = []             
        self.inverse_atlas = {}  

    def _clean_text(self, text):
        text = re.sub(r'\[.*?\]\(.*?\)', '', text) 
        text = re.sub(r'[#*`_\-狂]', '', text)     
        text = re.sub(r'\s+', ' ', text)           
        return text.strip()

    def encode_to_glyph(self, text):
        glyph = np.zeros((8, 8, 3), dtype=np.float32)
        words = text.lower().split()
        
        for i, word in enumerate(words):
            x = hash(word) % 8
            y = hash(word + str(i)) % 8
            
            r = (hash(word + "verb") % 100) / 100.0
            g = (hash(word + "noun") % 100) / 100.0
            b = (hash(word + "adj") % 100) / 100.0
            
            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):
        print(f"[*] Commencing training on ecosystem: {folder_path}")
        search_path = os.path.join(folder_path, "**", "*.md")
        md_files = glob.glob(search_path, recursive=True)
        
        if not md_files:
            print(f"[!] No markdown (.md) files found in {folder_path}!")
            return

        for file_path in md_files:
            print(f" -> Ingesting: {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 
                
                glyph = self.encode_to_glyph(sentence)
                self.kb.append((glyph, sentence))
                
                words = sentence.lower().split()
                for i, word in enumerate(words):
                    x = hash(word) % 8
                    y = hash(word + str(i)) % 8
                    
                    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)

        print(f"[+] Ingestion complete. KB contains {len(self.kb)} concepts.")
        print(f"[+] Inverse Atlas mapped {len(self.inverse_atlas)} unique spatial coordinates.")

    def think_ode_step(self, glyph, t):
        new_glyph = glyph.copy()
        
        for c in range(3):
            new_glyph[:, :, c] += 0.1 * laplace(glyph[:, :, c])
            
        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):
        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):
        prompt_glyph = self.encode_to_glyph(prompt)
        
        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 fluid remains dark. No matching concepts found.", prompt_glyph

        superposition = np.zeros((8, 8, 3), dtype=np.float32)
        for score, glyph, _ in top_candidates:
            superposition += glyph * (score if score > 0 else 0.1)
            
        for t in range(1, steps + 1):
            superposition = self.think_ode_step(superposition, t)
            if np.var(superposition) < 0.0005:
                break
                
        generated_text = self.decode_glyph_to_text(superposition)
        return generated_text, superposition

    def decode_glyph_to_text(self, glyph):
        words = []
        temp_glyph = glyph.copy()
        
        for _ in range(12): 
            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 
                
            color_idx = int(np.argmax(temp_glyph[y, x]))
            atlas_key = (x, y, color_idx)
            
            if atlas_key in self.inverse_atlas:
                matched_words = list(self.inverse_atlas[atlas_key])
                word = matched_words[hash(str(x)+str(y)) % len(matched_words)]
                words.append(word)
            
            temp_glyph[y, x] *= 0.0
            if x < 7: 
                temp_glyph[y, x + 1] *= 1.15 
                
        return " ".join(words).capitalize() + "."

if __name__ == "__main__":
    BOOK_FOLDER = "./my_books"
    os.makedirs(BOOK_FOLDER, exist_ok=True)
    
    if not os.listdir(BOOK_FOLDER):
        with open(os.path.join(BOOK_FOLDER, "sample.md"), "w") as f:
            f.write("# Chronicles of the Lab\n")
            f.write("The quick brown fox jumps over the lazy dog. ")
            f.write("A quiet cat sleeps soundly under the wooden table. ")
            f.write("The aggressive dog barked loudly at the passing stranger.")
            
    vct_brain = VisualCollapseAI()
    vct_brain.train_on_folder(BOOK_FOLDER)
    
    print("-" * 50)
    query = "Is there a dog near the table?"
    print(f"Prompt: '{query}'")
    
    response, final_canvas = vct_brain.generative_query(query, top_n=2)
    print(f"Generated VCT Reason: {response}")