import os
import numpy as np
from math import log, sqrt

class AcousticCollapseTheoryV4:
    def __init__(self, num_frequencies=4096, shape_threshold=0.35, max_ode_steps=8):
        self.NUM_FREQS = num_frequencies  # WIDEBAND: Prevents acoustic aliasing
        self.SHAPE_THRESHOLD = shape_threshold # Shape overlap required to confirm a match
        self.MAX_ODE_STEPS = max_ode_steps
        self.kb = []
        
        self.idf_weights = np.ones(self.NUM_FREQS) 
        self.total_docs = 0
        
        print(f"[ACT v4] Initialized Wideband AGC Acoustic Cortex ({self.NUM_FREQS} bins).")

    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 _apply_agc(self, signal):
        """Automatic Gain Control: Normalizes RMS energy to 1.0"""
        rms = np.sqrt(np.mean(np.abs(signal)**2))
        if rms > 0:
            return signal / rms
        return signal

    def encode_to_signal(self, text, apply_idf=False):
        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 
            amplitude = self.idf_weights[freq] if apply_idf else 1.0
            signal[freq] += amplitude * np.exp(1j * phase)
            
        # Apply AGC so sentence length doesn't bias the energy
        return self._apply_agc(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):
        magnitudes = np.abs(signal)
        phases = np.angle(signal)
        
        # Sharpen dominant pure tones, crush background noise
        sharpened = np.power(magnitudes + 0.01, 2.0) 
        
        phase_drift = (0.15 * time_step * np.arange(self.NUM_FREQS)) / self.NUM_FREQS
        new_phases = phases + phase_drift
        
        new_signal = sharpened * np.exp(1j * new_phases)
        
        # Re-apply AGC after ODE step to keep volume locked at 1.0
        return self._apply_agc(new_signal)

    def signal_shape_match(self, sig1, sig2):
        """
        Pure Phase-Aware Shape Matching (Complex Cosine Similarity).
        Because of AGC, this is purely geometric: 1.0 = identical shape, 0.0 = no overlap.
        """
        active_bins = (np.abs(sig1) > 0.01) | (np.abs(sig2) > 0.01)
        if not np.any(active_bins): return 0.0
        
        s1 = sig1[active_bins]
        s2 = sig2[active_bins]
        
        dot = np.vdot(s1, s2)
        norm = np.linalg.norm(s1) * np.linalg.norm(s2)
        return abs(dot / norm) if norm > 0 else 0.0

    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 = []
        
        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

        self.total_docs = len(raw_sentences)
        
        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
                
        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)} wideband 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)
        
        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] WAVEFORM LOCKED at t={t}.")
                break
                
            prev_signal = current_signal.copy()
            current_signal = self.apply_ode_dynamics(current_signal, t)
            print(f"[ODE-CCT] t={t}: Resonating... (Drift: {change:.4f})")

        best_score = 0.0
        best_response = "... (Noise Rejected - No Shape Match)"

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

        # Gate based on geometric shape overlap, not raw volume
        if best_score < self.SHAPE_THRESHOLD:
            print(f"[RESULT] Shape Overlap: {best_score:.3f} (Below Gate {self.SHAPE_THRESHOLD})")
            print(f"Output:  [REJECTED]\n" + "-"*60)
            return "[REJECTED]"
        else:
            print(f"[RESULT] Shape Overlap: {best_score:.3f} (Passed Gate)")
            print(f"Output:  {best_response}\n" + "-"*60)
            return best_response

if __name__ == "__main__":
    # 4096 bins, threshold 0.35 (35% shape overlap required)
    ai = AcousticCollapseTheoryV4(num_frequencies=4096, shape_threshold=0.35, 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)
