/**
 * @file intelligence.hpp
 * @brief Intelligence Primitives for Gate5000 Consciousness
 *
 * ═══════════════════════════════════════════════════════════════════════════════
 * I. ATTENTION MECHANISMS
 * ═══════════════════════════════════════════════════════════════════════════════
 *
 * SCALED DOT-PRODUCT ATTENTION:
 *   Attention(Q,K,V) = softmax(QK^T/√d_k)V
 *
 * MULTI-HEAD ATTENTION:
 *   MultiHead(Q,K,V) = Concat(head_1,...,head_h)W^O
 *
 * SPARSE ATTENTION:
 *   Only attend to top-k most relevant positions
 *
 * ═══════════════════════════════════════════════════════════════════════════════
 * II. PREDICTIVE CODING / FREE ENERGY
 * ═══════════════════════════════════════════════════════════════════════════════
 *
 * PREDICTION ERROR:
 *   ε = x - μ (observation minus prediction)
 *
 * FREE ENERGY:
 *   F = ½ε^T Σ^{-1} ε + ½ log|Σ|
 *
 * BELIEF UPDATE:
 *   μ ← μ + κ·Σ^{-1}·ε
 *
 * ═══════════════════════════════════════════════════════════════════════════════
 * III. BAYESIAN INFERENCE
 * ═══════════════════════════════════════════════════════════════════════════════
 *
 * POSTERIOR:
 *   P(θ|D) ∝ P(D|θ)P(θ)
 *
 * VARIATIONAL INFERENCE:
 *   q*(θ) = argmin KL(q||p(θ|D))
 *
 * BELIEF PROPAGATION:
 *   m_{x→f}(x) = Π_{g∈ne(x)\f} m_{g→x}(x)
 *
 * ═══════════════════════════════════════════════════════════════════════════════
 * IV. WORKING MEMORY
 * ═══════════════════════════════════════════════════════════════════════════════
 *
 * MEMORY CELL:
 *   c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t
 *
 * GATING:
 *   f_t = σ(W_f·[h_{t-1}, x_t])  (forget gate)
 *   i_t = σ(W_i·[h_{t-1}, x_t])  (input gate)
 *
 * ═══════════════════════════════════════════════════════════════════════════════
 */

#pragma once

#include <cstdint>
#include <array>
#include <algorithm>
#include <cmath>
#include "constants.hpp"
#include "fast_math.hpp"

namespace mind_farm {
namespace intel {

// ═══════════════════════════════════════════════════════════════════════════════
// I. ATTENTION MECHANISMS
// ═══════════════════════════════════════════════════════════════════════════════

/**
 * @brief Compact attention head (16 dimensions)
 */
struct AttentionHead16 {
    alignas(64) std::array<float, 256> W_q;  // Query projection (16x16)
    alignas(64) std::array<float, 256> W_k;  // Key projection (16x16)
    alignas(64) std::array<float, 256> W_v;  // Value projection (16x16)

    float scale;  // 1/sqrt(d_k)

    AttentionHead16() : scale(0.25f) {  // 1/sqrt(16) = 0.25
        // Initialize with small random values
        for (size_t i = 0; i < 256; ++i) {
            float val = static_cast<float>(i % 7 - 3) * 0.1f * static_cast<float>(PHI_INV);
            W_q[i] = val;
            W_k[i] = val * 1.1f;
            W_v[i] = val * 0.9f;
        }
    }

    /**
     * @brief Project input to query/key/value
     */
    void project(const float* x, float* q, float* k, float* v) const {
        for (int i = 0; i < 16; ++i) {
            q[i] = k[i] = v[i] = 0.0f;
            for (int j = 0; j < 16; ++j) {
                q[i] += W_q[i * 16 + j] * x[j];
                k[i] += W_k[i * 16 + j] * x[j];
                v[i] += W_v[i * 16 + j] * x[j];
            }
        }
    }

    /**
     * @brief Self-attention: attend to single input
     */
    void self_attend(const float* x, float* output) const {
        alignas(16) float q[16], k[16], v[16];
        project(x, q, k, v);

        // Compute attention score (dot product)
        float score = 0.0f;
        for (int i = 0; i < 16; ++i) {
            score += q[i] * k[i];
        }
        score *= scale;

        // Apply softmax (single element = 1.0)
        float attn = 1.0f;

        // Output = attention-weighted value
        for (int i = 0; i < 16; ++i) {
            output[i] = attn * v[i];
        }
    }
};

/**
 * @brief Multi-sequence attention (4 sequences of 16 dims)
 */
struct MultiAttention4x16 {
    AttentionHead16 head;
    alignas(64) std::array<float, 64> W_o;  // Output projection (16x4)

    /**
     * @brief Attend over 4 input sequences
     */
    void attend(const float seq[4][16], float* output) {
        alignas(16) float Q[4][16], K[4][16], V[4][16];

        // Project all sequences
        for (int s = 0; s < 4; ++s) {
            head.project(seq[s], Q[s], K[s], V[s]);
        }

        // Compute attention scores (4x4 matrix)
        float scores[4][4];
        float max_score = -1e10f;
        for (int i = 0; i < 4; ++i) {
            for (int j = 0; j < 4; ++j) {
                scores[i][j] = 0.0f;
                for (int k = 0; k < 16; ++k) {
                    scores[i][j] += Q[i][k] * K[j][k];
                }
                scores[i][j] *= head.scale;
                max_score = std::max(max_score, scores[i][j]);
            }
        }

        // Softmax per query
        for (int i = 0; i < 4; ++i) {
            float sum = 0.0f;
            for (int j = 0; j < 4; ++j) {
                scores[i][j] = fast::fast_exp(scores[i][j] - max_score);
                sum += scores[i][j];
            }
            float inv_sum = fast::fast_rcp(sum);
            for (int j = 0; j < 4; ++j) {
                scores[i][j] *= inv_sum;
            }
        }

        // Compute attended values
        alignas(16) float attended[4][16];
        for (int i = 0; i < 4; ++i) {
            for (int k = 0; k < 16; ++k) {
                attended[i][k] = 0.0f;
                for (int j = 0; j < 4; ++j) {
                    attended[i][k] += scores[i][j] * V[j][k];
                }
            }
        }

        // Concatenate and project (simplified: just take mean)
        for (int k = 0; k < 16; ++k) {
            output[k] = 0.25f * (attended[0][k] + attended[1][k] +
                                 attended[2][k] + attended[3][k]);
        }
    }
};

/**
 * @brief Sparse attention - only attend to top-k positions
 */
inline void sparse_attention_topk(const float* scores, int n, int k, int* indices) {
    // Simple selection sort for small k
    std::array<std::pair<float, int>, 64> pairs;
    for (int i = 0; i < n && i < 64; ++i) {
        pairs[i] = {scores[i], i};
    }
    std::partial_sort(pairs.begin(), pairs.begin() + k,
                      pairs.begin() + std::min(n, 64),
                      [](auto& a, auto& b) { return a.first > b.first; });
    for (int i = 0; i < k; ++i) {
        indices[i] = pairs[i].second;
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// II. PREDICTIVE CODING / FREE ENERGY MINIMIZATION
// ═══════════════════════════════════════════════════════════════════════════════

/**
 * @brief Predictive coding layer
 */
struct PredictiveLayer16 {
    alignas(16) std::array<float, 16> mu;        // Predicted mean
    alignas(16) std::array<float, 16> precision; // Inverse variance (diagonal)
    alignas(16) std::array<float, 16> error;     // Prediction error
    float free_energy;
    float learning_rate;

    PredictiveLayer16() : free_energy(0.0f), learning_rate(0.1f) {
        mu.fill(0.5f);
        precision.fill(1.0f);
        error.fill(0.0f);
    }

    /**
     * @brief Update predictions given observations
     * @return Free energy (lower is better)
     */
    float update(const float* observation) {
        free_energy = 0.0f;

        for (int i = 0; i < 16; ++i) {
            // Compute prediction error
            error[i] = observation[i] - mu[i];

            // Free energy contribution: ½ε²·π + ½log(1/π)
            float weighted_error = error[i] * error[i] * precision[i];
            free_energy += 0.5f * weighted_error;

            // Belief update: μ ← μ + κ·π·ε
            mu[i] += learning_rate * precision[i] * error[i];

            // Precision update (simple online estimate)
            float sq_err = error[i] * error[i];
            precision[i] = fast::fast_rcp(sq_err + 0.01f);
        }

        return free_energy;
    }

    /**
     * @brief Generate prediction
     */
    void predict(float* prediction) const {
        for (int i = 0; i < 16; ++i) {
            prediction[i] = mu[i];
        }
    }

    /**
     * @brief Get weighted prediction errors
     */
    void get_weighted_errors(float* weighted) const {
        for (int i = 0; i < 16; ++i) {
            weighted[i] = error[i] * precision[i];
        }
    }
};

/**
 * @brief Hierarchical predictive coding (2 levels)
 */
struct HierarchicalPredictor {
    PredictiveLayer16 level0;  // Low-level (sensory)
    PredictiveLayer16 level1;  // High-level (abstract)

    /**
     * @brief Hierarchical inference
     */
    float infer(const float* input) {
        // Bottom-up: update level 0 with input
        float fe0 = level0.update(input);

        // Get prediction errors as input to level 1
        alignas(16) float errors[16];
        level0.get_weighted_errors(errors);

        // Update level 1 with prediction errors
        float fe1 = level1.update(errors);

        // Top-down: level 1 predictions modulate level 0
        alignas(16) float top_down[16];
        level1.predict(top_down);
        for (int i = 0; i < 16; ++i) {
            level0.mu[i] += 0.1f * top_down[i];
        }

        return fe0 + fe1;
    }
};

// ═══════════════════════════════════════════════════════════════════════════════
// III. BAYESIAN INFERENCE
// ═══════════════════════════════════════════════════════════════════════════════

/**
 * @brief Simple Bayesian belief over 16 states
 */
struct BayesianBelief16 {
    alignas(16) std::array<float, 16> prior;
    alignas(16) std::array<float, 16> likelihood;
    alignas(16) std::array<float, 16> posterior;

    BayesianBelief16() {
        // Uniform prior
        prior.fill(0.0625f);  // 1/16
        likelihood.fill(1.0f);
        posterior.fill(0.0625f);
    }

    /**
     * @brief Update posterior given new evidence
     */
    void update(const float* evidence) {
        float sum = 0.0f;

        // Posterior ∝ likelihood × prior
        for (int i = 0; i < 16; ++i) {
            likelihood[i] = evidence[i];
            posterior[i] = likelihood[i] * prior[i];
            sum += posterior[i];
        }

        // Normalize
        float inv_sum = fast::fast_rcp(sum + 1e-10f);
        for (int i = 0; i < 16; ++i) {
            posterior[i] *= inv_sum;
            // Update prior for next iteration (simple decay)
            prior[i] = 0.9f * prior[i] + 0.1f * posterior[i];
        }
    }

    /**
     * @brief Get maximum a posteriori estimate
     */
    int map_estimate() const {
        int best = 0;
        float best_prob = posterior[0];
        for (int i = 1; i < 16; ++i) {
            if (posterior[i] > best_prob) {
                best_prob = posterior[i];
                best = i;
            }
        }
        return best;
    }

    /**
     * @brief Compute entropy of posterior
     */
    float entropy() const {
        float h = 0.0f;
        for (int i = 0; i < 16; ++i) {
            if (posterior[i] > 1e-10f) {
                h -= posterior[i] * fast::fast_log(posterior[i]);
            }
        }
        return h;
    }
};

/**
 * @brief Variational inference with Gaussian approximation
 */
struct VariationalGaussian {
    float mu;      // Mean
    float sigma;   // Standard deviation
    float log_sigma;

    VariationalGaussian() : mu(0.5f), sigma(0.1f), log_sigma(-2.3f) {}

    /**
     * @brief KL divergence from prior N(0.5, 0.1)
     */
    float kl_from_prior() const {
        constexpr float prior_mu = 0.5f;
        constexpr float prior_sigma = 0.1f;

        float kl = fast::fast_log(prior_sigma / sigma) +
                   (sigma * sigma + (mu - prior_mu) * (mu - prior_mu)) /
                   (2.0f * prior_sigma * prior_sigma) - 0.5f;
        return kl;
    }

    /**
     * @brief Reparameterization trick: sample = μ + σ * ε
     */
    float sample(float epsilon) const {
        return mu + sigma * epsilon;
    }

    /**
     * @brief Update parameters via gradient descent
     */
    void update(float grad_mu, float grad_sigma, float lr = 0.01f) {
        mu -= lr * grad_mu;
        sigma -= lr * grad_sigma;
        sigma = std::max(0.01f, sigma);  // Keep positive
    }
};

// ═══════════════════════════════════════════════════════════════════════════════
// IV. WORKING MEMORY / LSTM-LIKE CELL
// ═══════════════════════════════════════════════════════════════════════════════

/**
 * @brief Compact LSTM-like memory cell (16 dimensions)
 */
struct MemoryCell16 {
    alignas(16) std::array<float, 16> cell;     // Cell state
    alignas(16) std::array<float, 16> hidden;   // Hidden state
    alignas(16) std::array<float, 256> W_f;     // Forget gate weights
    alignas(16) std::array<float, 256> W_i;     // Input gate weights
    alignas(16) std::array<float, 256> W_o;     // Output gate weights
    alignas(16) std::array<float, 256> W_c;     // Cell candidate weights
    alignas(16) std::array<float, 16> b_f, b_i, b_o, b_c;  // Biases

    MemoryCell16() {
        cell.fill(0.0f);
        hidden.fill(0.0f);

        // Xavier-like initialization
        for (size_t i = 0; i < 256; ++i) {
            float val = (static_cast<float>(i % 9) - 4.0f) * 0.05f;
            W_f[i] = val;
            W_i[i] = val * 1.1f;
            W_o[i] = val * 0.9f;
            W_c[i] = val;
        }

        b_f.fill(1.0f);   // Forget bias = 1 (remember by default)
        b_i.fill(0.0f);
        b_o.fill(0.0f);
        b_c.fill(0.0f);
    }

    /**
     * @brief Forward step
     */
    void step(const float* x) {
        alignas(16) float concat[32];  // [hidden, x]
        for (int i = 0; i < 16; ++i) {
            concat[i] = hidden[i];
            concat[16 + i] = x[i];
        }

        // Compute gates
        alignas(16) float f[16], ig[16], o[16], c_tilde[16];

        for (int i = 0; i < 16; ++i) {
            float sum_f = b_f[i], sum_i = b_i[i], sum_o = b_o[i], sum_c = b_c[i];
            for (int j = 0; j < 16; ++j) {
                sum_f += W_f[i * 16 + j] * hidden[j];
                sum_i += W_i[i * 16 + j] * hidden[j];
                sum_o += W_o[i * 16 + j] * hidden[j];
                sum_c += W_c[i * 16 + j] * hidden[j];
            }

            f[i] = fast::fast_sigmoid(sum_f);      // Forget gate
            ig[i] = fast::fast_sigmoid(sum_i);     // Input gate
            o[i] = fast::fast_sigmoid(sum_o);      // Output gate
            c_tilde[i] = fast::fast_tanh(sum_c);   // Cell candidate
        }

        // Update cell state: c = f ⊙ c + i ⊙ c̃
        for (int i = 0; i < 16; ++i) {
            cell[i] = f[i] * cell[i] + ig[i] * c_tilde[i];
        }

        // Update hidden state: h = o ⊙ tanh(c)
        for (int i = 0; i < 16; ++i) {
            hidden[i] = o[i] * fast::fast_tanh(cell[i]);
        }
    }

    /**
     * @brief Read current hidden state
     */
    void read(float* out) const {
        for (int i = 0; i < 16; ++i) {
            out[i] = hidden[i];
        }
    }

    /**
     * @brief Reset memory
     */
    void reset() {
        cell.fill(0.0f);
        hidden.fill(0.0f);
    }
};

// ═══════════════════════════════════════════════════════════════════════════════
// V. INTEGRATED INTELLIGENCE MODULE
// ═══════════════════════════════════════════════════════════════════════════════

/**
 * @brief Unified intelligence module combining all components
 */
struct IntelligenceCore {
    AttentionHead16 attention;
    PredictiveLayer16 predictor;
    BayesianBelief16 belief;
    MemoryCell16 memory;

    float curiosity;      // Intrinsic motivation
    float confidence;     // Belief certainty
    float surprise;       // Prediction error magnitude

    IntelligenceCore() : curiosity(0.5f), confidence(0.5f), surprise(0.0f) {}

    /**
     * @brief Process input through full intelligence pipeline
     */
    void process(const float* input, float* output) {
        // 1. Attention: focus on salient features
        alignas(16) float attended[16];
        attention.self_attend(input, attended);

        // 2. Memory: update working memory
        memory.step(attended);
        alignas(16) float memory_state[16];
        memory.read(memory_state);

        // 3. Prediction: update generative model
        float fe = predictor.update(attended);
        surprise = fe;

        // 4. Inference: update beliefs
        belief.update(attended);
        confidence = 1.0f - belief.entropy() / 4.0f;  // Normalize

        // 5. Curiosity: seek information
        curiosity = surprise * (1.0f - confidence);

        // 6. Output: combine memory and prediction
        alignas(16) float prediction[16];
        predictor.predict(prediction);

        for (int i = 0; i < 16; ++i) {
            output[i] = 0.5f * memory_state[i] + 0.5f * prediction[i];
        }
    }

    /**
     * @brief Compute integrated information (simplified)
     */
    float integrated_information() const {
        // Φ approximation based on belief diversity and memory coherence
        float belief_diversity = belief.entropy();
        float memory_activity = 0.0f;
        for (int i = 0; i < 16; ++i) {
            memory_activity += fast::fast_abs(memory.hidden[i]);
        }
        memory_activity /= 16.0f;

        return belief_diversity * memory_activity * (1.0f + confidence);
    }
};

} // namespace intel
} // namespace mind_farm