import os
import numpy as np
from math import log

class AcousticCollapseTheoryV3:
    def __init__(self, num_frequencies=512, energy_threshold=0.5, max_ode_steps=8):
        self.NUM_FREQS = num_frequencies
        self.ENERGY_THRESHOLD = energy_threshold # Minimum acoustic energy to accept a match
        self.MAX_ODE_STEPS = max_ode_steps
        self.kb = []
        
        # Acoustic Q-Factor weights (Inverse Document Frequency)
        self.idf_weights = np.ones(self.NUM_FREQS) 
        self.total_docs = 0
        
        print(f"[ACT v3] Initialized Q-Factor Filtered Acoustic Cortex.")

    def _text_to_freq(self, bigram):
        hash_val = 5381
        for char in bigram.lower():
            hash_val = ((hash_val << 5) + hash_val) + ord(char)
        return hash_val % self.NUM_FREQS

    def encode_to_signal(self, text, apply_idf=False):
        """Encodes text into a 1D wave, weighted by informational purity (IDF)."""
        signal = np.zeros(self.NUM_FREQS, dtype=np.complex128)
        words = text.lower().split()
        if len(words) < 2: return signal
        
        bigrams = [f"{words[i]} {words[i+1]}" for i in range(len(words) - 1)]
        num_bigrams = len(bigrams)
        
        for i, bigram in enumerate(bigrams):
            freq = self._text_to_freq(bigram)
            phase = (i / num_bigrams) * 2 * np.pi 
            
            # Acoustic Weighting: Rare words = Loud, Common words = Quiet
            amplitude = self.idf_weights[freq] if apply_idf else 1.0
            signal[freq] += amplitude * np.exp(1j * phase)
            
        return signal

    def calculate_wave_stability(self, sig1, sig2):
        mag1 = np.abs(sig1)
        mag2 = np.abs(sig2)
        active_bins = (mag1 > 0.01) | (mag2 > 0.01)
        if not np.any(active_bins): return 1.0
        
        s1_active = sig1[active_bins]
        s2_active = sig2[active_bins]
        
        dot = np.vdot(s1_active, s2_active)
        norm = np.linalg.norm(s1_active) * np.linalg.norm(s2_active)
        return abs(dot / norm) if norm > 0 else 1.0

    def apply_ode_dynamics(self, signal, time_step):
        """Non-linear resonant cavity with aggressive clipping (squashes background noise)."""
        magnitudes = np.abs(signal)
        phases = np.angle(signal)
        
        # Aggressive power-law sharpening (simulates a physical limiter/compressor)
        sharpened = np.power(magnitudes + 0.01, 2.5) 
        
        phase_drift = (0.15 * time_step * np.arange(self.NUM_FREQS)) / self.NUM_FREQS
        new_phases = phases + phase_drift
        
        return sharpened * np.exp(1j * new_phases)

    def signal_match_score(self, sig1, sig2):
        """
        Acoustic Energy Matching. No normalization tricks.
        Returns raw energy transferred. If they don't share specific pure tones, it's 0.
        """
        # Complex dot product captures both Amplitude overlap AND Phase alignment
        raw_energy_transfer = np.abs(np.vdot(sig1, sig2))
        return raw_energy_transfer

    def train(self, books_dir="books"):
        if not os.path.exists(books_dir):
            print(f"[ERROR] '{books_dir}' not found."); return

        print(f"[TRAIN] Scanning '{books_dir}'...")
        raw_sentences = []
        
        # Pass 1: Count document frequencies for Q-Factor calculation
        for filename in os.listdir(books_dir):
            if filename.endswith(".txt") or filename.endswith(".md"):
                with open(os.path.join(books_dir, filename), 'r', encoding='utf-8', errors='ignore') as f:
                    lines = [line.strip() for line in f if len(line.strip()) > 15]
                    raw_sentences.extend(lines)
                    
                    for line in lines:
                        words = line.lower().split()
                        bigrams = set(f"{words[i]} {words[i+1]}" for i in range(len(words) - 1))
                        for bigram in bigrams:
                            freq = self._text_to_freq(bigram)
                            self.idf_weights[freq] += 1 # Count how many times bin is hit

        self.total_docs = len(raw_sentences)
        
        # Calculate IDF: log(Total_Docs / Doc_Freq) -> High if rare, Low if common
        for i in range(self.NUM_FREQS):
            if self.idf_weights[i] > 0:
                self.idf_weights[i] = log(self.total_docs / self.idf_weights[i]) + 1.0
            else:
                self.idf_weights[i] = 0.0 # Never seen, mute it
                
        # Pass 2: Build KB with IDF-weighted signals
        for i in range(len(raw_sentences) - 1):
            trigger_sig = self.encode_to_signal(raw_sentences[i], apply_idf=True)
            response_sig = self.encode_to_signal(raw_sentences[i+1], apply_idf=True)
            self.kb.append((trigger_sig, response_sig, raw_sentences[i+1]))
            
        print(f"[TRAIN] Complete. {len(self.kb)} high-fidelity associations built.\n")

    def query(self, prompt_text):
        if not self.kb: return "Knowledge base empty."

        print(f"Input:  {prompt_text}")
        current_signal = self.encode_to_signal(prompt_text, apply_idf=True)
        prev_signal = np.zeros_like(current_signal)
        
        # ODE Loop
        for t in range(1, self.MAX_ODE_STEPS + 1):
            stability = self.calculate_wave_stability(current_signal, prev_signal)
            change = 1.0 - stability
            
            if change < 0.02 and t > 1:
                print(f"[ODE-CCT] RESONANCE COLLAPSED at t={t}.")
                break
                
            prev_signal = current_signal.copy()
            current_signal = self.apply_ode_dynamics(current_signal, t)
            print(f"[ODE-CCT] t={t}: Filtering noise... (Shape change: {change:.4f})")

        # Energy-Gated Matching
        best_score = 0.0
        best_response = "... (Acoustic Noise Rejected - No Match)"

        for trigger_sig, response_sig, response_text in self.kb:
            score = self.signal_match_score(current_signal, trigger_sig)
            if score > best_score:
                best_score = score
                best_response = response_text

        # Apply Energy Gate
        if best_score < self.ENERGY_THRESHOLD:
            print(f"[RESULT] Acoustic Energy: {best_score:.3f} (Below Gate {self.ENERGY_THRESHOLD})")
            print(f"Output:  [REJECTED - Insufficient harmonic overlap]\n" + "-"*60)
            return "[REJECTED]"
        else:
            print(f"[RESULT] Acoustic Energy: {best_score:.3f} (Passed Gate)")
            print(f"Output:  {best_response}\n" + "-"*60)
            return best_response

if __name__ == "__main__":
    # Energy threshold might need slight tweaking based on your specific text lengths
    ai = AcousticCollapseTheoryV3(num_frequencies=512, energy_threshold=1.5, max_ode_steps=8)
    ai.train("books")
    
    print("=== STARTING EVALUATION ===")
    test_prompts = [
        "What is the main concept?",
        "Explain the theory of relativity",
        "The quantum state collapses",
        "Poirot looked at the murder scene", 
        "Random words apple bicycle gravity"
    ]
    
    for prompt in test_prompts:
        ai.query(prompt)
