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

class AcousticCollapseTheoryV5:
    def __init__(self, num_frequencies=4096, coherence_threshold=0.25, max_ode_steps=8):
        self.NUM_FREQS = num_frequencies  
        self.COHERENCE_THRESHOLD = coherence_threshold # True phase alignment required
        self.MAX_ODE_STEPS = max_ode_steps
        self.kb = []
        
        self.idf_weights = np.ones(self.NUM_FREQS) 
        self.total_docs = 0
        
        print(f"[ACT v5] Initialized Phase-Coherent 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 mapping based on word order
            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)
            
        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)
        
        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)
        return self._apply_agc(new_signal)

    def signal_phase_coherence(self, sig1, sig2):
        """
        TRUE CONSTRUCTIVE INTERFERENCE.
        Only measures bins where BOTH signals exist.
        Uses Real() to enforce phase direction. Misaligned phases cancel out.
        """
        # Intersection: Only bins where both waves have amplitude
        intersect_bins = (np.abs(sig1) > 0.05) & (np.abs(sig2) > 0.05)
        
        if not np.any(intersect_bins):
            return 0.0
            
        s1 = sig1[intersect_bins]
        s2 = sig2[intersect_bins]
        
        # Complex dot product gives: Sum(Amp1 * Amp2 * cos(phase1 - phase2))
        # Taking np.real() extracts the constructive interference.
        # If phases are opposite (cos = -1), this goes negative.
        constructive_energy = np.real(np.vdot(s1, s2))
        
        # Max possible energy if phases were perfectly aligned
        max_possible = np.linalg.norm(s1) * np.linalg.norm(s2)
        
        if max_possible == 0: return 0.0
        
        coherence = constructive_energy / max_possible
        
        # If coherence is negative, destructive interference occurred. Clamp to 0.
        return max(0.0, coherence)

    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)} phase-locked 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 = "... (Destructive Interference - No Match)"

        for trigger_sig, response_sig, response_text in self.kb:
            # Use the new Phase Coherence metric
            score = self.signal_phase_coherence(current_signal, trigger_sig)
            if score > best_score:
                best_score = score
                best_response = response_text

        if best_score < self.COHERENCE_THRESHOLD:
            print(f"[RESULT] Phase Coherence: {best_score:.3f} (Below Gate {self.COHERENCE_THRESHOLD})")
            print(f"Output:  [REJECTED - Out of Phase]\n" + "-"*60)
            return "[REJECTED]"
        else:
            print(f"[RESULT] Phase Coherence: {best_score:.3f} (Constructive Interference Achieved)")
            print(f"Output:  {best_response}\n" + "-"*60)
            return best_response

if __name__ == "__main__":
    # Lowered threshold slightly to 0.25 because intersection math is stricter
    ai = AcousticCollapseTheoryV5(num_frequencies=4096, coherence_threshold=0.25, 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)
