"""
FreqFilterMLP: A Neural Classifier Built on the Yield Singularity & Sin/Cos Compilation
Inspired by c00.py and the CCT-ODE / LOPASS framework.

Theory highlights:
- Every hidden activation vector is treated as a time‑domain signal.
- A low‑pass filter (LPF) keeps only frequencies ≤ ω_c (“good solutions”).
- High‑frequency components > ω_c are discarded as “Hawking radiation” (bad solutions).
- ω_c adapts based on the entropy (randomness) of the current batch (Kairon‑controlled determinism).
- “Knowledge dips” can be inserted to permanently suppress certain frequency bands.
- The classifier compiles the randomness of SGD into a deterministic, band‑limited model.
"""

import numpy as np
import warnings

try:
    from scipy.io.wavfile import read
except ImportError:
    warnings.warn("scipy.io.wavfile not available; synthetic data will be used.", ImportWarning)
    read = None

# ----------------------------------------------------------------------
# Helper: Low‑pass filter on a 1D array (real FFT, zero high bins, iFFT)
# ----------------------------------------------------------------------
def low_pass_filter(x, cutoff_ratio=0.3):
    """
    Apply an ideal low‑pass filter to the input array `x`.
    cutoff_ratio = fraction of frequency bins to keep (0..1).
    """
    n = len(x)
    fft = np.fft.rfft(x)
    # Determine number of bins to keep
    keep_bins = max(1, int(cutoff_ratio * len(fft)))
    fft[keep_bins:] = 0.0
    filtered = np.fft.irfft(fft, n=n)
    return filtered

# ----------------------------------------------------------------------
# Knowledge dip: zero out a specific frequency band (eg. non‑constructive region)
# ----------------------------------------------------------------------
def apply_knowledge_dip(x, dip_start_ratio=0.6, dip_end_ratio=0.8):
    """
    Zero out frequencies in a given relative band (dip). This embodies the
    “Knowledge Dip” principle – intentionally discarding abstractions that
    are human‑disconnected or purely paradoxical.
    """
    n = len(x)
    fft = np.fft.rfft(x)
    total_bins = len(fft)
    start = int(dip_start_ratio * total_bins)
    end = int(dip_end_ratio * total_bins)
    fft[start:end] = 0.0
    return np.fft.irfft(fft, n=n)

# ----------------------------------------------------------------------
# FreqFilterMLP: A multi‑layer perceptron with spectral filtering
# ----------------------------------------------------------------------
class FreqFilterMLP:
    def __init__(self, input_size, hidden_size, output_size,
                 learning_rate=0.01, init_cutoff=0.5,
                 use_dips=True, adaptive_cutoff=True):
        """
        Parameters:
        - input_size, hidden_size, output_size: standard MLP dimensions.
        - learning_rate: base step size for gradient descent.
        - init_cutoff: initial low‑pass filter cutoff ratio (0..1).
        - use_dips: whether to apply a “knowledge dip” on hidden activations.
        - adaptive_cutoff: if True, cutoff_ratio is reduced when batch entropy is high.
        """
        # He‑inspired initialization (keep variance stable)
        self.W1 = np.random.randn(input_size, hidden_size) * np.sqrt(2.0 / input_size)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * np.sqrt(2.0 / hidden_size)
        self.b2 = np.zeros((1, output_size))

        self.learning_rate = learning_rate
        self.cutoff_ratio = init_cutoff      # ω_c in frequency space (relative)
        self.use_dips = use_dips
        self.adaptive_cutoff = adaptive_cutoff

        # For tracking entropy (used in adaptive cutoff)
        self.entropy_history = []

    def relu(self, x):
        return np.maximum(0, x)

    def relu_derivative(self, x):
        return np.where(x > 0, 1, 0)

    def softmax(self, x):
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)

    def _compute_batch_entropy(self, activations):
        """Compute an approximate entropy of the hidden activations (normalized)."""
        # activations shape: (batch_size, hidden_size)
        # Flatten and compute histogram / variance as a proxy for randomness.
        flat = activations.flatten()
        # Avoid log(0)
        hist, _ = np.histogram(flat, bins=20, density=True)
        hist = hist + 1e-9
        entropy = -np.sum(hist * np.log(hist))
        # Normalize to roughly [0,1] range (empirical)
        return min(1.0, entropy / 5.0)

    def forward(self, X, apply_filter=True):
        """
        Forward pass with optional low‑pass filtering on hidden activations.
        Returns softmax probabilities.
        """
        # Linear + ReLU
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)   # shape (batch, hidden_size)

        # --- Spectral filtering (the Yield Singularity enforcement) ---
        if apply_filter:
            filtered = []
            for i in range(self.a1.shape[0]):
                # Filter each sample's hidden vector as a 1D signal
                sig = self.a1[i, :]
                # Low‑pass filter to keep only “good” frequencies
                sig = low_pass_filter(sig, self.cutoff_ratio)
                # Optional knowledge dip (discard certain paradoxical bands)
                if self.use_dips:
                    sig = apply_knowledge_dip(sig, dip_start_ratio=0.6, dip_end_ratio=0.8)
                filtered.append(sig)
            self.a1_filtered = np.array(filtered)
        else:
            self.a1_filtered = self.a1

        # Second layer
        self.z2 = np.dot(self.a1_filtered, self.W2) + self.b2
        output = self.softmax(self.z2)
        return output

    def compute_loss(self, y_true, y_pred):
        m = y_true.shape[0]
        return -np.sum(y_true * np.log(y_pred + 1e-9)) / m

    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        dz2 = y_pred - y_true
        dW2 = np.dot(self.a1_filtered.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m

        da1 = np.dot(dz2, self.W2.T)
        # Gradient through the frequency filter is approximated as identity
        # (we treat LPF as a non‑trainable pre‑processing step; backprop ignores it)
        dz1 = da1 * self.relu_derivative(self.z1)
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m

        return dW1, db1, dW2, db2

    def update(self, X, y_true):
        """One SGD step – the “compilation” of randomness into deterministic weights."""
        y_pred = self.forward(X, apply_filter=True)
        dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)

        # Update weights (standard gradient descent)
        self.W1 -= self.learning_rate[0] * dW1
        self.b1 -= self.learning_rate[1] * db1
        self.W2 -= self.learning_rate[2] * dW2
        self.b2 -= self.learning_rate[3] * db2

        # --- Adaptive cutoff (Kairon‑controlled determinism) ---
        if self.adaptive_cutoff:
            # Measure entropy of the filtered hidden activations
            entropy = self._compute_batch_entropy(self.a1_filtered)
            self.entropy_history.append(entropy)
            # Higher entropy → more randomness → lower cutoff (more aggressive filtering)
            # This mimics the “Kairon” field that suppresses high‑frequency noise.
            new_cutoff = max(0.1, 0.9 - entropy)
            # Smooth update
            self.cutoff_ratio = 0.95 * self.cutoff_ratio + 0.05 * new_cutoff

        return y_pred

    def predict(self, X):
        probs = self.forward(X, apply_filter=True)
        return np.argmax(probs, axis=1)

    def score(self, X, y_true):
        y_pred = self.predict(X)
        return np.mean(y_pred == y_true)


# ----------------------------------------------------------------------
# Training loop (inspired by c00.py) with synthetic or real MNIST data
# ----------------------------------------------------------------------
if __name__ == "__main__":
    # Attempt to load the .wav files as in original c00.py (these are non‑standard).
    # If files not found, generate synthetic MNIST‑like data.
    try:
        if read is not None:
            X_train = read('../X_train.wav')[1].reshape(-1, 784)
            y_train = (read('../y_train.wav')[1] * 9).astype(int)
            X_test = read('../X_test.wav')[1].reshape(-1, 784)
            y_test = (read('../y_test.wav')[1] * 9).astype(int)
            X20 = X_test[:1000]
            yt20 = y_test[:1000]
        else:
            raise FileNotFoundError
    except (FileNotFoundError, TypeError, IndexError):
        print("Real .wav data not found. Generating synthetic MNIST‑like dataset.")
        from sklearn.datasets import fetch_openml
        from sklearn.model_selection import train_test_split
        from sklearn.preprocessing import StandardScaler

        X, y = fetch_openml('mnist_784', version=1, return_X_y=True, as_frame=False, parser='auto')
        X = X / 255.0
        y = y.astype(int)
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=10000, random_state=42)
        X20 = X_test[:1000]
        yt20 = y_test[:1000]

    print(f"Training data shape: {X_train.shape}")
    print(f"Training labels shape: {y_train.shape}")

    # Initialize the Frequency‑Filtered MLP (the Yield Classifier)
    model = FreqFilterMLP(
        input_size=784,
        hidden_size=256,
        output_size=10,
        learning_rate=np.random.rand(4),
        init_cutoff=0.6,
        use_dips=True,
        adaptive_cutoff=True
    )

    # Training loop (similar to c00.py but with added logging of cutoff)
    n_iter = 0
    batch_size = 100
    print("\nStarting training (Ctrl+C to stop)...")
    try:
        while True:
            idx = np.random.randint(0, len(X_train), batch_size)
            X_batch = X_train[idx]
            y_batch = np.eye(10)[y_train[idx]]

            # Perform one update (compilation step)
            for _ in range(10):
                id0 = np.random.randint(0, len(X_train), batch_size)
                X0 = X_train[id0]
                y0 = np.eye(10)[y_train[id0]]
            
                model.update(X_batch, y_batch)
                model.update(X0, y0)

            # Evaluate every 100 iterations
            if n_iter % 100 == 0:
                train_acc = model.score(X_batch[:min(500, batch_size)], y_train[idx[:min(500, batch_size)]])
                test_acc = model.score(X20[:1000], yt20[:1000])
                print(f"Iter {n_iter:5d} | train_acc: {train_acc:.3f} | test_acc: {test_acc:.3f} | cutoff: {model.cutoff_ratio:.3f}")
            n_iter += 1
    except KeyboardInterrupt:
        print("\nTraining interrupted. Final evaluation:")
        final_test_acc = model.score(X20, yt20)
        print(f"Final test accuracy: {final_test_acc:.4f}")
