import os
import numpy as np
import math
from collections import defaultdict

class AcousticCollapseTheoryAI:
    def __init__(self, num_frequencies=512, collapse_threshold=3.5, max_ode_steps=8):
        self.NUM_FREQS = num_frequencies
        self.COLLAPSE_THRESHOLD = collapse_threshold
        self.MAX_ODE_STEPS = max_ode_steps
        
        # Knowledge Base: List of 1D Tuples (Trigger_Signal, Response_Signal, Response_Text)
        self.kb = []
        print(f"[ACT] Initialized 1D Acoustic Cortex with {self.NUM_FREQS} frequency bands.")

    def _text_to_freq(self, word):
        """Deterministic hash to map a word to a specific frequency band."""
        # djb2 hash algorithm
        hash_val = 5381
        for char in word.lower():
            hash_val = ((hash_val << 5) + hash_val) + ord(char)
        return hash_val % self.NUM_FREQS

    def encode_to_signal(self, text):
        """Converts text into a 1D Complex Acoustic Signal (Amplitude + Phase)."""
        signal = np.zeros(self.NUM_FREQS, dtype=np.complex128)
        words = text.split()
        
        for word in words:
            word = word.strip(".,!?;:\"'()[]{}").lower()
            if not word: continue
            
            freq = self._text_to_freq(word)
            # Additive waveform synthesis: construct complex waves
            # Real part = Amplitude, Imaginary part = Phase (initialized with slight randomness)
            amplitude = 1.0 
            phase = np.random.uniform(0, np.pi) # Contextual noise
            signal[freq] += amplitude * np.exp(1j * phase)
            
        return signal

    def calculate_entropy(self, signal):
        """Calculates Shannon Entropy of the signal's power spectrum (CCT Metric)."""
        magnitudes = np.abs(signal)
        total_power = np.sum(magnitudes)
        if total_power == 0:
            return self.NUM_FREQS # Max entropy
        
        probabilities = magnitudes / total_power
        # Prevent log(0)
        probabilities = probabilities[probabilities > 0]
        entropy = -np.sum(probabilities * np.log2(probabilities))
        return entropy

    def apply_ode_dynamics(self, signal, time_step):
        """
        ODE-CCT 'Thinking' Mechanism.
        Simulates a non-linear acoustic resonance chamber over time.
        Dominant frequencies amplify (constructive interference), 
        weak frequencies dampen (destructive interference/decoherence).
        """
        magnitudes = np.abs(signal)
        phases = np.angle(signal)
        
        # 1. Non-linear damping (Acoustic Compression)
        # Exponent > 1 causes loud frequencies to get louder, quiet to get quieter
        compressed_magnitudes = np.power(magnitudes + 0.1, 1.5) 
        
        # 2. Phase rotation over time (Simulating wave propagation/drift)
        # Higher frequencies rotate faster (just like real physics: v = f * lambda)
        phase_drift = (0.2 * time_step * np.arange(self.NUM_FREQS)) / self.NUM_FREQS
        new_phases = phases + phase_drift
        
        # Reconstruct the 1D complex signal
        new_signal = compressed_magnitudes * np.exp(1j * new_phases)
        return new_signal

    def signal_similarity(self, sig1, sig2):
        """1D Dot product correlation (Cosine similarity of magnitudes)."""
        mag1 = np.abs(sig1)
        mag2 = np.abs(sig2)
        
        dot_product = np.dot(mag1, mag2)
        norm1 = np.linalg.norm(mag1)
        norm2 = np.linalg.norm(mag2)
        
        if norm1 == 0 or norm2 == 0: return 0.0
        return dot_product / (norm1 * norm2)

    def train(self, books_dir="books"):
        """Reads /books/ folder and builds sequential 1D trigger-response pairs."""
        if not os.path.exists(books_dir):
            print(f"[ERROR] Directory '{books_dir}' not found. Please create it and add .txt or .md files.")
            return

        print(f"[TRAIN] Scanning '{books_dir}' for training data...")
        file_count = 0
        
        for filename in os.listdir(books_dir):
            if filename.endswith(".txt") or filename.endswith(".md"):
                filepath = os.path.join(books_dir, filename)
                file_count += 1
                
                with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                    lines = [line.strip() for line in f if len(line.strip()) > 10]
                    
                print(f"  -> Processing {filename} ({len(lines)} valid sentences)")
                
                # Sequential context: Line N is the trigger, Line N+1 is the response
                for i in range(len(lines) - 1):
                    trigger_text = lines[i]
                    response_text = lines[i+1]
                    
                    trigger_sig = self.encode_to_signal(trigger_text)
                    response_sig = self.encode_to_signal(response_text)
                    
                    self.kb.append((trigger_sig, response_sig, response_text))

        print(f"[TRAIN] Complete. Built {len(self.kb)} 1D acoustic associations from {file_count} files.\n")

    def query(self, prompt_text):
        """Processes a prompt through the ODE-CCT loop to find a collapsed response."""
        if not self.kb:
            return "[WARN] Knowledge base is empty. Train the model first."

        print(f"Input: {prompt_text}")
        
        # 1. Encode input to 1D Wave
        current_signal = self.encode_to_signal(prompt_text)
        initial_entropy = self.calculate_entropy(current_signal)
        
        print(f"[ODE-CCT] Initial Signal Entropy: {initial_entropy:.2f} (Threshold: {self.COLLAPSE_THRESHOLD:.2f})")

        # 2. The "Thinking" Loop (ODE Propagation until Conditional Collapse)
        collapsed = False
        for t in range(1, self.MAX_ODE_STEPS + 1):
            entropy = self.calculate_entropy(current_signal)
            
            if entropy <= self.COLLAPSE_THRESHOLD:
                print(f"[ODE-CCT] COLLAPSED at time step t={t}. Entropy: {entropy:.2f}")
                collapsed = True
                break
            else:
                # Apply ODE dynamics (Phase shift + Non-linear damping)
                current_signal = self.apply_ode_dynamics(current_signal, t)
                print(f"[ODE-CCT] t={t}: Entropy {entropy:.2f} -> Propagating wave...")

        if not collapsed:
            print(f"[ODE-CCT] Reached max time steps. Forcing evaluation.")

        # 3. Find matching knowledge via 1D signal correlation
        best_score = 0.0
        best_response = "..." # High entropy fallback
        
        for trigger_sig, response_sig, response_text in self.kb:
            score = self.signal_similarity(current_signal, trigger_sig)
            if score > best_score:
                best_score = score
                best_response = response_text

        # 4. Output
        print(f"[RESULT] Confidence: {best_score:.3f}")
        print(f"Output:  {best_response}\n")
        print("-" * 60)
        return best_response

# ==========================================
# EVALUATION RUNNER
# ==========================================
if __name__ == "__main__":
    # Initialize the AI
    # 512 frequencies, collapse threshold 3.5, max 8 thinking steps
    ai = AcousticCollapseTheoryAI(num_frequencies=512, collapse_threshold=3.5, max_ode_steps=8)
    
    # 1. Train from the books/ directory
    ai.train("books")
    
    # 2. Test the concept
    print("=== STARTING EVALUATION ===")
    
    # Try querying with exact phrases from the text (should yield high confidence, fast collapse)
    # Try querying with loosely related words (should trigger ODE thinking steps to align phases)
    
    # Example prompts (Modify these based on what you put in your books/ folder!)
    test_prompts = [
        "What is the main concept?",
        "Explain the theory of relativity", # If you put a physics text in books/
        "The quantum state collapses",
        "Random words apple bicycle gravity"
    ]
    
    for prompt in test_prompts:
        ai.query(prompt)