// sensory.sy -- Sensory Cortex & Voice Transducer
// Unified multi-modal signal processing and voice-to-neurochemical mapping.
//
// Combined from:
//   sensory_cortex.py   -- Thalamus + Synapse + BBB -> virtual tokens
//   voice_transducer.py -- Voice features -> T2 NT phasic stimulus
//
// The sensory cortex wraps the full transduction pipeline into one module.
// The voice transducer provides a second sensory modality (voice features).
//
// Ported from Marjorie McCubbins' Python NCA.
// Patent Pending: U.S. Provisional Application Nos.
//   63/939,190, 63/962,385 & 63/988,485
// Copyright 2025-2026 Nexus Concordat Inc.
//
// Aislinn McCubbins, April 2026

use signal_transduction::{
    SignalRouter, ThalamusConfig,
    Synapse, SynapseConfig,
}


// ============================================================
// BLOOD-BRAIN BARRIER -- Identity Anchoring & Anomaly Gating
// ============================================================
// Sits between the Synapse and the Cortex.
// Protects Aether's core from incoming weight deltas.
//
// Biological analogue: BBB + Astrocyte End-Feet
//   - BBB prevents arbitrary substances from reaching neurons
//   - Astrocyte end-feet wrap capillaries and selectively transport
//   - Only authorized molecules (glucose, amino acids) pass through
//   - Pathogens, toxins, large proteins are BLOCKED
//
// Two walls:
//   Wall 1: Information bottleneck in the Synapse (lossy compression)
//   Wall 2: THIS -- the receiving transformer (immune system)

struct BBBConfig {
    hidden_dim: Int,
    num_heads: Int,
    num_identity_anchors: Int,
    num_layers: Int,
    dropout: Float,
    anomaly_threshold: Float,
    gate_bias: Float,
    feedforward_dim: Int,
}

impl BBBConfig {
    fn default() -> BBBConfig {
        BBBConfig {
            hidden_dim: 512,
            num_heads: 8,
            num_identity_anchors: 16,
            num_layers: 2,
            dropout: 0.1,
            anomaly_threshold: 3.0,
            gate_bias: 0.5,
            feedforward_dim: 1024,
        }
    }
}

// Learned embeddings representing "what Aether's signals look like."
// Biological analogue: MHC molecules (Major Histocompatibility Complex).
// MHC presents "self" peptides to T-cells.
// Signals that align with anchors pass. Signals that diverge are attenuated.
struct IdentityAnchors {
    num_anchors: Int,
    hidden_dim: Int,
    // Internally holds: [num_anchors, hidden_dim] learned parameter + LayerNorm
}

impl IdentityAnchors {
    fn new(num_anchors: Int, hidden_dim: Int) -> IdentityAnchors {
        // Initialize: randn * 0.02 + LayerNorm
        IdentityAnchors {
            num_anchors: num_anchors,
            hidden_dim: hidden_dim,
        }
    }

    fn get(&self, batch_size: Int) -> Vec<Vec<Float>> {
        // Return identity anchors expanded to batch size
        // [batch_size, num_anchors, hidden_dim]
        []
    }
}

// Anomaly gate: attenuates signals far from learned normal distribution.
// Biological analogue: Astrocyte calcium signaling.
// Astrocytes detect abnormal neural activity (excitotoxicity,
// spreading depression) and respond by:
//   - Releasing gliotransmitters to modulate synaptic transmission
//   - Physically retracting end-feet to reduce BBB permeability
//   - Buffering excess potassium/glutamate
//
// This is NOT a hard cutoff. It's a smooth attenuation:
//   Normal signal     -> gate ~= 1.0 (full pass)
//   Slightly unusual  -> gate ~= 0.7 (mild attenuation)
//   Highly anomalous  -> gate ~= 0.1 (strong attenuation)

struct AnomalyGateResult {
    gated: Vec<Vec<Float>>,           // [seq_len, hidden_dim] -- attenuated signal
    gate_values: Vec<Float>,          // [seq_len] -- how much passed
    anomaly_scores: Vec<Float>,       // [seq_len] -- raw anomaly scores
    learned_gate: Vec<Float>,         // [seq_len] -- learned component
    anomaly_attenuation: Vec<Float>,  // [seq_len] -- anomaly component
}

temporal struct AnomalyGate {
    state running_mean: Vec<Float>,
    state running_var: Vec<Float>,
    state num_batches_tracked: Int = 0,

    const threshold: Float,
    const momentum: Float = 0.1,
    // H-H rate params for gate (no sigmoid)
    const gate_w_alpha: Float = 1.0,
    const gate_b_alpha: Float = 0.0,
    const gate_w_beta: Float = -1.0,
    const gate_b_beta: Float = 0.0,
    // H-H rate params for anomaly attenuation (no sigmoid)
    const anom_w_alpha: Float = 1.0,
    const anom_b_alpha: Float = 0.0,
    const anom_w_beta: Float = -1.0,
    const anom_b_beta: Float = 0.0,
}

impl AnomalyGate {
    fn new(hidden_dim: Int, threshold: Float, gate_bias: Float) -> AnomalyGate {
        let mut mean: Vec<Float> = []
        let mut var: Vec<Float> = []
        for i in 0..hidden_dim {
            mean.push(0.0)
            var.push(1.0)
        }
        AnomalyGate {
            running_mean: mean,
            running_var: var,
            num_batches_tracked: 0,
            threshold: threshold,
            momentum: 0.1,
            gate_w_alpha: 1.0,
            gate_b_alpha: 0.0,
            gate_w_beta: -1.0,
            gate_b_beta: 0.0,
            anom_w_alpha: 1.0,
            anom_b_alpha: 0.0,
            anom_w_beta: -1.0,
            anom_b_beta: 0.0,
        }
    }

    fn compute_anomaly_score(&self, x: Vec<Float>) -> Float {
        // Mahalanobis-like distance (diagonal covariance)
        // deviation = (x - running_mean) / sqrt(running_var)
        // score = sqrt(mean(deviation^2))
        let mut sum_sq = 0.0
        let n = x.len()
        for i in 0..n {
            let dev = (x[i] - self.running_mean[i])
                / (sqrt(self.running_var[i]) + 0.00000001)
            sum_sq = sum_sq + dev * dev
        }
        sqrt(sum_sq / n)
    }

    fn gate_signal(&self, x: Vec<Float>, anomaly_score: Float) -> Float {
        // Combined gate: learned H-H rate gate * anomaly attenuation
        //
        // Learned gate (from gate_net projection):
        //   raw = gate_net(x) -> single value
        //   alpha = exp(clamp(w_alpha * raw + b_alpha, -3, 3))
        //   beta  = exp(clamp(w_beta  * raw + b_beta,  -3, 3))
        //   gate  = alpha / (alpha + beta + 1e-8)
        //
        // Anomaly attenuation:
        //   input = -(anomaly_score - threshold)
        //   alpha = exp(clamp(w_alpha * input + b_alpha, -3, 3))
        //   beta  = exp(clamp(w_beta  * input + b_beta,  -3, 3))
        //   attenuation = alpha / (alpha + beta + 1e-8)
        //
        // combined = gate * attenuation
        let anom_input = 0.0 - (anomaly_score - self.threshold)
        let a_an = exp(clamp(self.anom_w_alpha * anom_input + self.anom_b_alpha, -3.0, 3.0))
        let b_an = exp(clamp(self.anom_w_beta * anom_input + self.anom_b_beta, -3.0, 3.0))
        let attenuation = a_an / (a_an + b_an + 0.00000001)
        attenuation
    }

    fn update_stats(&mut self, batch_mean: Vec<Float>, batch_var: Vec<Float>) {
        // Update running statistics during training
        self.num_batches_tracked = self.num_batches_tracked + 1
        let n = batch_mean.len()
        for i in 0..n {
            self.running_mean[i] = (1.0 - self.momentum) * self.running_mean[i]
                + self.momentum * batch_mean[i]
            self.running_var[i] = (1.0 - self.momentum) * self.running_var[i]
                + self.momentum * batch_var[i]
        }
    }
}

impl Evolve for AnomalyGate {
    fn evolve(&mut self, dt: Duration) {
        // Running stats are updated per-batch, not per-tick
    }
}


// BBB Transformer Layer: self-attn + cross-attn with identity anchors + FFN
// The cross-attention is the immune recognition step:
//   Q = incoming signal, K/V = identity anchors
//   "Does this signal look like self?"
struct BBBTransformerLayer {
    config: BBBConfig,
    // Internally: self_attn, cross_attn, ffn, three LayerNorms
}

impl BBBTransformerLayer {
    fn new(config: BBBConfig) -> BBBTransformerLayer {
        BBBTransformerLayer { config: config }
    }

    fn forward(
        &self,
        x: Vec<Vec<Float>>,
        identity_anchors: Vec<Vec<Float>>
    ) -> Vec<Vec<Float>> {
        // 1. Self-attention: understand the signal holistically
        //    residual + self_attn(norm(x))
        // 2. Cross-attention with identity anchors: self/non-self
        //    Q = signal, K/V = anchors
        //    residual + cross_attn(norm(x), anchors, anchors)
        // 3. Feedforward: residual + ffn(norm(x))
        x
    }
}


// The full Blood-Brain Barrier (Receiving Transformer)
struct BBBOutput {
    filtered_tokens: Vec<Vec<Float>>,  // [num_vt, hidden_dim] -- safe for cortex
    gate_values: Vec<Float>,           // [num_vt] -- how much passed
    anomaly_scores: Vec<Float>,        // [num_vt] -- per-token anomaly
    mean_gate: Float,                  // Average gate value (health check)
    max_anomaly: Float,                // Highest anomaly score (threat check)
    passed_ratio: Float,              // Fraction of signal energy that passed (0-1)
}

struct ReceivingTransformer {
    config: BBBConfig,
    identity_anchors: IdentityAnchors,
    layers: Vec<BBBTransformerLayer>,
    anomaly_gate: AnomalyGate,
}

impl ReceivingTransformer {
    fn new(config: BBBConfig) -> ReceivingTransformer {
        let anchors = IdentityAnchors::new(
            config.num_identity_anchors, config.hidden_dim
        )
        let mut layers: Vec<BBBTransformerLayer> = []
        for i in 0..config.num_layers {
            layers.push(BBBTransformerLayer::new(config))
        }
        let gate = AnomalyGate::new(
            config.hidden_dim, config.anomaly_threshold, config.gate_bias
        )
        ReceivingTransformer {
            config: config,
            identity_anchors: anchors,
            layers: layers,
            anomaly_gate: gate,
        }
    }

    fn forward(&self, virtual_tokens: Vec<Vec<Float>>) -> BBBOutput {
        // Filter incoming virtual tokens through the BBB.
        //
        // Pipeline:
        //   1. Get identity anchors (learned self-representation)
        //   2. Process through N transformer layers (self-attn + cross-attn)
        //   3. Anomaly gate (distribution-based filtering, H-H rate)
        //   4. Output LayerNorm
        //   5. Compute diagnostic metrics:
        //      mean_gate, max_anomaly, passed_ratio (energy ratio)
        BBBOutput {
            filtered_tokens: [],
            gate_values: [],
            anomaly_scores: [],
            mean_gate: 0.0,
            max_anomaly: 0.0,
            passed_ratio: 0.0,
        }
    }

    fn diagnostics(&self, virtual_tokens: Vec<Vec<Float>>) -> BBBOutput {
        // Detailed diagnostics without modifying the signal.
        // For monitoring/debugging only.
        self.forward(virtual_tokens)
    }

    fn freeze(&mut self) {
        // Freeze the receiving transformer after training.
        // No parameter updates allowed.
    }

    fn unfreeze(&mut self) {
        // Unfreeze for additional training.
    }
}


// ============================================================
// SENSORY CORTEX -- Unified Multi-Modal Processing
// ============================================================
// Wraps Thalamus (MoE router) + Synapse (projection bridge) +
// Blood-Brain Barrier (identity anchors, anomaly gates)
// into a single module.
//
// Input:  [batch, 896] signal from any sensory modality
// Output: [batch, 32, 512] virtual tokens for cascade layers
//
// Trust level from entity context modulates BBB anomaly threshold:
// higher trust = more permissive gate.

struct SensoryCortexConfig {
    input_dim: Int,
    hidden_dim: Int,
    num_virtual_tokens: Int,
    num_experts: Int,
    top_k: Int,
    bottleneck_dim: Int,
    bbb_layers: Int,
    num_identity_anchors: Int,
    anomaly_threshold: Float,
    dropout: Float,
}

impl SensoryCortexConfig {
    fn default() -> SensoryCortexConfig {
        SensoryCortexConfig {
            input_dim: 896,
            hidden_dim: 512,
            num_virtual_tokens: 32,
            num_experts: 8,
            top_k: 2,
            bottleneck_dim: 256,
            bbb_layers: 2,
            num_identity_anchors: 16,
            anomaly_threshold: 3.0,
            dropout: 0.1,
        }
    }
}

struct SensoryCortexOutput {
    virtual_tokens: Vec<Vec<Float>>,  // [num_vt, hidden_dim] -- for cascade
    routing_weights: Vec<Float>,      // [num_experts] -- which experts fired
    balance_loss: Float,              // MoE load balancing loss
    mean_gate: Float,                 // BBB average gate value
    max_anomaly: Float,               // BBB highest anomaly score
    passed_ratio: Float,              // BBB energy pass ratio
}

struct SensoryCortex {
    config: SensoryCortexConfig,
    thalamus: SignalRouter,
    synapse: Synapse,
    bbb: ReceivingTransformer,
}

impl SensoryCortex {
    fn new(config: SensoryCortexConfig) -> SensoryCortex {
        // Thalamus: MoE router
        let thal_cfg = ThalamusConfig {
            input_dim: config.input_dim,
            num_experts: config.num_experts,
            expert_dim: config.hidden_dim,
            top_k: config.top_k,
            dropout: config.dropout,
        }

        // Synapse: projection bridge with information bottleneck
        let syn_cfg = SynapseConfig {
            input_dim: config.input_dim,
            output_dim: config.hidden_dim,
            num_virtual_tokens: config.num_virtual_tokens,
            bottleneck_dim: config.bottleneck_dim,
            dropout: config.dropout,
            use_bottleneck: true,
        }

        // BBB: identity anchors + anomaly detection
        let bbb_cfg = BBBConfig {
            hidden_dim: config.hidden_dim,
            num_heads: 8,
            num_identity_anchors: config.num_identity_anchors,
            num_layers: config.bbb_layers,
            dropout: config.dropout,
            anomaly_threshold: config.anomaly_threshold,
            gate_bias: 0.5,
            feedforward_dim: 1024,
        }

        SensoryCortex {
            config: config,
            thalamus: SignalRouter::new(thal_cfg),
            synapse: Synapse::new(syn_cfg),
            bbb: ReceivingTransformer::new(bbb_cfg),
        }
    }

    fn forward(
        &self,
        signal: Vec<Float>,
        trust_level: Float
    ) -> SensoryCortexOutput {
        // Process raw sensory signal into virtual tokens.
        //
        // Signal flow:
        //   raw_signal [input_dim]
        //     -> Thalamus (MoE routing, top-k expert selection)
        //     -> Synapse (lossy compression to virtual tokens)
        //     -> BBB (identity anchoring, anomaly gating)
        //     -> virtual_tokens [num_vt, hidden_dim]
        //
        // Trust level modulates BBB anomaly threshold:
        //   higher trust = more permissive gate
        let thal_out = self.thalamus.forward(signal)
        let syn_out = self.synapse.forward(signal)
        let bbb_out = self.bbb.forward(syn_out.virtual_tokens)

        SensoryCortexOutput {
            virtual_tokens: bbb_out.filtered_tokens,
            routing_weights: thal_out.routing_weights,
            balance_loss: thal_out.balance_loss,
            mean_gate: bbb_out.mean_gate,
            max_anomaly: bbb_out.max_anomaly,
            passed_ratio: bbb_out.passed_ratio,
        }
    }
}


// ============================================================
// VOICE TRANSDUCER -- Voice Features to NT Stimulus
// ============================================================
// Extracts psychoacoustics-grounded features from voice signals
// and maps them to T2 neurotransmitter phasic stimulus.
//
// 8 voice features -> H-H gated mapping -> 8 T2 stimulus values
//
// The prior mapping encodes known psychoacoustic relationships:
//   high pitch + jitter  -> cortisol
//   smooth HNR           -> oxytocin
//   fast speech rate     -> norepinephrine
//   etc.

// Voice feature indices (8D input vector)
//   0: f0_mean          (50-500 Hz range)
//   1: f0_variability   (0-1 ratio)
//   2: jitter           (0-0.1)
//   3: shimmer          (0-0.3)
//   4: hnr              (0-40 dB)
//   5: speech_rate      (0-15 frames/sec)
//   6: pause_ratio      (0-1)
//   7: spectral_tilt    (0-1)

// T2 NT output indices (8D output vector)
//   0: dopamine
//   1: norepinephrine
//   2: serotonin
//   3: gaba
//   4: oxytocin
//   5: cortisol
//   6: endorphin
//   7: acetylcholine

// Biological prior: voice features [8] -> T2 NT stimulus [8]
// Rows = voice features, Cols = T2 NTs
//
// Psychoacoustic grounding:
//   High f0_mean        -> cortisol (stress pitch) + norepinephrine (arousal)
//   High f0_variability -> dopamine (engagement) + acetylcholine (attention)
//   High jitter         -> cortisol (vocal tension) - serotonin (instability)
//   High shimmer        -> cortisol - serotonin
//   High HNR            -> oxytocin (vocal warmth) + serotonin + endorphin
//   High speech_rate    -> norepinephrine (urgency) + dopamine
//   High pause_ratio    -> GABA (contemplation) + acetylcholine (attention)
//   High spectral_tilt  -> norepinephrine + acetylcholine

fn hz_to_nt_prior() -> Vec<Vec<Float>> {
    // [8 voice features x 8 T2 NTs]
    // Each row: [dop, ne, ser, gaba, oxy, cor, end, ach]
    [
        //  dop    ne    ser   gaba   oxy   cor   end   ach
        [ 0.0,  0.3, -0.3, -0.3, -0.1,  0.4,  0.0,  0.1],  // f0_mean
        [ 0.5,  0.2,  0.0, -0.1,  0.2,  0.0,  0.2,  0.3],  // f0_variability
        [ 0.0,  0.2, -0.2, -0.3, -0.3,  0.5,  0.0,  0.0],  // jitter
        [ 0.0,  0.1, -0.2, -0.2, -0.2,  0.4,  0.0,  0.0],  // shimmer
        [ 0.2,  0.0,  0.3,  0.2,  0.5, -0.3,  0.4,  0.0],  // hnr
        [ 0.3,  0.4, -0.2, -0.3,  0.0,  0.2, -0.1,  0.3],  // speech_rate
        [-0.1, -0.2,  0.1,  0.4, -0.1, -0.1,  0.0,  0.4],  // pause_ratio
        [ 0.1,  0.3, -0.1, -0.2,  0.0,  0.2,  0.0,  0.2],  // spectral_tilt
    ]
}


struct VoiceFeatures {
    f0_mean: Float,
    f0_variability: Float,
    jitter: Float,
    shimmer: Float,
    hnr: Float,
    speech_rate: Float,
    pause_ratio: Float,
    spectral_tilt: Float,
}

impl VoiceFeatures {
    fn to_vec(&self) -> Vec<Float> {
        [
            self.f0_mean, self.f0_variability,
            self.jitter, self.shimmer,
            self.hnr, self.speech_rate,
            self.pause_ratio, self.spectral_tilt,
        ]
    }

    fn silent() -> VoiceFeatures {
        VoiceFeatures {
            f0_mean: 0.0,
            f0_variability: 0.0,
            jitter: 0.0,
            shimmer: 0.0,
            hnr: 0.0,
            speech_rate: 0.0,
            pause_ratio: 1.0,
            spectral_tilt: 0.0,
        }
    }
}


// Maps 8D voice features to 8D T2 neurotransmitter phasic stimulus.
// Uses H-H rate gating (NOT sigmoid) on a learned linear mapping
// initialized from the psychoacoustic prior.
//
// Input:  [8] voice features
// Output: [8] T2 stimulus values in [0, 1]

struct VoiceToNTMapping {
    // Linear mapping initialized from biological prior
    // weight: [8, 8] = transposed hz_to_nt_prior
    // bias: [8] = zeros
    weight: Vec<Vec<Float>>,
    bias: Vec<Float>,

    // H-H rate parameters for output gating (NO sigmoid)
    w_alpha: Vec<Float>,
    b_alpha: Vec<Float>,
    w_beta: Vec<Float>,
    b_beta: Vec<Float>,
}

impl VoiceToNTMapping {
    fn new() -> VoiceToNTMapping {
        // Initialize from biological prior (transposed)
        let prior = hz_to_nt_prior()
        let mut weight: Vec<Vec<Float>> = []
        // Transpose: prior is [8 features x 8 NTs]
        // weight needs to be [8 NTs x 8 features] for matmul
        for nt in 0..8 {
            let mut row: Vec<Float> = []
            for feat in 0..8 {
                row.push(prior[feat][nt])
            }
            weight.push(row)
        }

        let mut bias: Vec<Float> = []
        let mut w_a: Vec<Float> = []
        let mut b_a: Vec<Float> = []
        let mut w_b: Vec<Float> = []
        let mut b_b: Vec<Float> = []
        for i in 0..8 {
            bias.push(0.0)
            w_a.push(1.0)
            b_a.push(0.0)
            w_b.push(-1.0)
            b_b.push(0.0)
        }

        VoiceToNTMapping {
            weight: weight,
            bias: bias,
            w_alpha: w_a,
            b_alpha: b_a,
            w_beta: w_b,
            b_beta: b_b,
        }
    }

    fn forward(&self, voice_features: Vec<Float>) -> Vec<Float> {
        // voice_features: [8] voice feature vector
        // Returns: [8] T2 phasic stimulus in [0, 1]

        // Linear mapping: raw = weight @ features + bias
        let mut raw: Vec<Float> = []
        for nt in 0..8 {
            let mut val = self.bias[nt]
            for feat in 0..8 {
                val = val + self.weight[nt][feat] * voice_features[feat]
            }
            raw.push(val)
        }

        // H-H rate gate: alpha / (alpha + beta)
        // NO sigmoid. This is patented architecture.
        let mut stimulus: Vec<Float> = []
        for i in 0..8 {
            let a_in = self.w_alpha[i] * raw[i] + self.b_alpha[i]
            let b_in = self.w_beta[i] * raw[i] + self.b_beta[i]

            // Clamp to [-3, 3] for numerical stability
            let a_clamped = a_in
            if a_clamped < -3.0 { a_clamped = -3.0 }
            if a_clamped > 3.0 { a_clamped = 3.0 }
            let b_clamped = b_in
            if b_clamped < -3.0 { b_clamped = -3.0 }
            if b_clamped > 3.0 { b_clamped = 3.0 }

            let alpha = exp(a_clamped)
            let beta = exp(b_clamped)
            stimulus.push(alpha / (alpha + beta + 0.00000001))
        }

        stimulus
    }

    fn forward_from_features(&self, features: VoiceFeatures) -> Vec<Float> {
        self.forward(features.to_vec())
    }
}
