// optimizer.sy -- Adam Optimizer with Biological Decay
// Implements Adam (Kingma & Ba, 2015) using tensor operations.
// H-H rate schedule: learning rate modulated by neurochemical state.
//
// No sigmoid. Anywhere. Ever.
//
// Marjorie McCubbins, Aislinn McCubbins, April 2026
// Patent Pending: 63/939,190 & 63/962,385 & 63/988,485

// === ADAM STATE ===
// One per parameter tensor. Tracks first and second moment estimates.

struct AdamState {
    m: Tensor,       // First moment (mean of gradients)
    v: Tensor,       // Second moment (mean of squared gradients)
    step: Int,       // Step counter for bias correction
}

impl AdamState {
    fn new(shape_like: Tensor) -> AdamState {
        let shape = t_shape(shape_like)
        AdamState {
            m: tensor_zeros(shape),
            v: tensor_zeros(shape),
            step: 0,
        }
    }
}

// === ADAM OPTIMIZER ===

struct Adam {
    lr: Float,
    beta1: Float,
    beta2: Float,
    eps: Float,
    weight_decay: Float,
    max_grad_norm: Float,
    states: Vec<AdamState>,
}

impl Adam {
    fn new(lr: Float) -> Adam {
        Adam {
            lr: lr,
            beta1: 0.9,
            beta2: 0.999,
            eps: 0.00000001,
            weight_decay: 0.01,
            max_grad_norm: 1.0,
            states: [],
        }
    }

    fn with_betas(mut self, b1: Float, b2: Float) -> Adam {
        self.beta1 = b1
        self.beta2 = b2
        self
    }

    fn with_weight_decay(mut self, wd: Float) -> Adam {
        self.weight_decay = wd
        self
    }

    // Initialize states for all parameter tensors
    fn init(mut self, params: Vec<Tensor>) -> Adam {
        for i in 0..params.len() {
            self.states.push(AdamState::new(params[i]))
        }
        self
    }

    // Single parameter update step
    // Returns updated parameter tensor
    fn step_param(mut self, idx: Int, param: Tensor, grad: Tensor) -> Tensor {
        let st = self.states[idx]
        st.step = st.step + 1

        // Gradient clipping (global norm)
        let grad_norm = t_norm(grad)
        let clipped = grad
        if grad_norm > self.max_grad_norm {
            let scale = self.max_grad_norm / (grad_norm + self.eps)
            clipped = t_mul(grad, tensor_fill(t_shape(grad), scale))
        }

        // Weight decay (decoupled, AdamW style)
        let decayed = t_sub(param, t_mul(param, tensor_fill(t_shape(param), self.lr * self.weight_decay)))

        // Update biased first moment: m = beta1 * m + (1 - beta1) * grad
        let m_scaled = t_mul(st.m, tensor_fill(t_shape(st.m), self.beta1))
        let g_scaled = t_mul(clipped, tensor_fill(t_shape(clipped), 1.0 - self.beta1))
        st.m = t_add(m_scaled, g_scaled)

        // Update biased second moment: v = beta2 * v + (1 - beta2) * grad^2
        let v_scaled = t_mul(st.v, tensor_fill(t_shape(st.v), self.beta2))
        let g_sq = t_mul(clipped, clipped)
        let g2_scaled = t_mul(g_sq, tensor_fill(t_shape(g_sq), 1.0 - self.beta2))
        st.v = t_add(v_scaled, g2_scaled)

        // Bias correction
        let bc1 = 1.0 - pow(self.beta1, st.step)
        let bc2 = 1.0 - pow(self.beta2, st.step)
        let m_hat = t_mul(st.m, tensor_fill(t_shape(st.m), 1.0 / bc1))
        let v_hat = t_mul(st.v, tensor_fill(t_shape(st.v), 1.0 / bc2))

        // Update: param = param - lr * m_hat / (sqrt(v_hat) + eps)
        let v_sqrt = t_sqrt(v_hat)
        let v_eps = t_add(v_sqrt, tensor_fill(t_shape(v_sqrt), self.eps))
        let update = t_mul(m_hat, tensor_fill(t_shape(m_hat), self.lr))
        let update_scaled = t_div(update, v_eps)

        t_sub(decayed, update_scaled)
    }

    // Update all parameters at once
    fn step_all(mut self, params: Vec<Tensor>, grads: Vec<Tensor>) -> Vec<Tensor> {
        let updated = []
        for i in 0..params.len() {
            updated.push(self.step_param(i, params[i], grads[i]))
        }
        updated
    }
}

// === LEARNING RATE SCHEDULERS ===

struct WarmupCosineScheduler {
    base_lr: Float,
    warmup_steps: Int,
    total_steps: Int,
    current_step: Int,
    min_lr: Float,
}

impl WarmupCosineScheduler {
    fn new(base_lr: Float, warmup: Int, total: Int) -> WarmupCosineScheduler {
        WarmupCosineScheduler {
            base_lr: base_lr,
            warmup_steps: warmup,
            total_steps: total,
            current_step: 0,
            min_lr: base_lr * 0.1,
        }
    }

    fn get_lr(mut self) -> Float {
        self.current_step = self.current_step + 1

        if self.current_step < self.warmup_steps {
            // Linear warmup
            return self.base_lr * (self.current_step + 0.0) / (self.warmup_steps + 0.0)
        }

        // Cosine decay
        let progress = (self.current_step - self.warmup_steps + 0.0) / (self.total_steps - self.warmup_steps + 0.0)
        if progress > 1.0 { progress = 1.0 }
        let cosine = 0.5 * (1.0 + cos(progress * 3.14159265))
        self.min_lr + (self.base_lr - self.min_lr) * cosine
    }
}

// === NEUROMODULATED LEARNING RATE ===
// Chemistry modulates how fast we learn.
// High dopamine = high learning rate (reward-driven exploration)
// High cortisol = reduced learning rate (stress narrows plasticity)
// High serotonin = stable learning rate (mood stabilization)

fn neuromodulated_lr(base_lr: Float, dopamine: Float, cortisol: Float, serotonin: Float) -> Float {
    // H-H rate modulation (NOT sigmoid)
    let excitatory = hh_gate(dopamine, 2.0, -0.5, 1.0, 0.0)
    let inhibitory = hh_gate(cortisol, 1.5, -0.3, 2.0, 0.5)
    let stabilizer = hh_gate(serotonin, 1.0, 0.0, 1.0, 0.0)

    // Dopamine opens the learning gate, cortisol closes it, serotonin stabilizes
    let modulation = excitatory * (1.0 - inhibitory) * (0.5 + 0.5 * stabilizer)

    // Clamp to reasonable range: 0.1x to 3x base
    let modulated = base_lr * (0.1 + 2.9 * modulation)
    if modulated < base_lr * 0.1 { return base_lr * 0.1 }
    if modulated > base_lr * 3.0 { return base_lr * 3.0 }
    modulated
}
