/**
 * NeuroFlow 完整反向传播实现
 * 
 * 实现完整模型的梯度反向传播链：
 * - Output fusion -> SN -> ECN/DMN -> Memory -> Input projection
 */

#pragma once

// C++ 标准库
#include <cmath>
#include <memory>
#include <unordered_map>
#include <vector>

// 第三方库
#ifdef _OPENMP
#include <omp.h>
#endif

// 项目头文件
#include "memory.hpp"
#include "model.hpp"
#include "networks.hpp"
#include "tensor.hpp"

namespace neuroflow {

/**
 * 完整反向传播引擎
 */
class BackpropEngine {
public:
    // 缓存前向传播中间结果
    struct ForwardCache {
        Tensor input;
        Tensor h;  // 输入投影后
        Tensor h_normed;
        Tensor h_gelu;
        
        // SN中间结果
        Tensor sn_h1, sn_h2;
        Tensor sn_gate_h;
        
        // ECN中间结果
        std::vector<Tensor> ecn_hidden;
        Tensor ecn_ofc_v;
        Tensor ecn_vmpfc_d;
        
        // DMN中间结果
        Tensor dmn_encoded;
        Tensor dmn_latent;
        std::vector<Tensor> dmn_associations;
        Tensor dmn_vision;
        
        // Memory结果
        Tensor memory_encoded;
        Tensor memory_query;
        Tensor memory_retrieved;
        
        // Output fusion
        Tensor combined;
        Tensor fused_pre_norm;
        Tensor fused_bn;
        Tensor fused_bn_pre_relu;
        Tensor fused;
    };
    
    ForwardCache cache;
    NeuroFlowModel& model;
    
    BackpropEngine(NeuroFlowModel& m) : model(m) {}
    
    // 前向传播（带缓存）
    NeuroFlowModel::Output forward_with_cache(const Tensor& x) {
        cache.input = x.clone();
        
        size_t batch = x.shape_[0];
        
        // Input projection
        cache.h = model.input_proj_linear->forward(x);
        cache.h_normed = model.input_proj_norm->forward(cache.h);
        cache.h_gelu = model.input_proj_gelu->forward(cache.h_normed);
        
        // SN
        auto sn_out = model.sn->forward(cache.h_gelu);
        
        // ECN
        Tensor h_for_ecn = cache.h_gelu.clone();
        for (size_t i = 0; i < model.ecn->num_layers; ++i) {
            h_for_ecn = model.ecn->dlpfc_linear[i]->forward(h_for_ecn);
            h_for_ecn = model.ecn->dlpfc_norm[i]->forward(h_for_ecn);
            h_for_ecn = model.ecn->dlpfc_gelu[i]->forward(h_for_ecn);
            cache.ecn_hidden.push_back(h_for_ecn.clone());
        }
        
        // OFC
        cache.ecn_ofc_v = model.ecn->ofc1->forward(cache.ecn_hidden.back());
        TensorOps::gelu(cache.ecn_ofc_v);
        
        // vmPFC
        cache.ecn_vmpfc_d = model.ecn->vmpfc1->forward(cache.ecn_hidden.back());
        TensorOps::gelu(cache.ecn_vmpfc_d);
        
        // DMN
        cache.memory_encoded = model.memory->encode(cache.h_gelu);
        cache.dmn_encoded = model.dmn->mem_encoder1->forward(cache.memory_encoded);
        TensorOps::gelu(cache.dmn_encoded);
        cache.dmn_latent = model.dmn->mem_encoder2->forward(cache.dmn_encoded);
        
        for (auto& [h1, h2] : model.dmn->association_heads) {
            Tensor assoc = h1->forward(cache.dmn_latent);
            TensorOps::gelu(assoc);
            assoc = h2->forward(assoc);
            cache.dmn_associations.push_back(assoc.clone());
        }
        
        cache.dmn_vision = TensorOps::concat(cache.dmn_associations, 1);
        cache.dmn_vision = model.dmn->future_proj1->forward(cache.dmn_vision);
        cache.dmn_vision = model.dmn->future_norm->forward(cache.dmn_vision);
        cache.dmn_vision = model.dmn->future_gelu->forward(cache.dmn_vision);
        
        // Memory retrieval
        cache.memory_query = model.memory->query_proj->forward(cache.h_gelu);
        auto mem_out = model.memory->forward(cache.h_gelu);
        cache.memory_retrieved = mem_out.retrieved;
        
        // Output fusion
        NeuroFlowModel::Output out;
        out.saliency = sn_out.saliency;
        out.gates = sn_out.gates;
        out.anomaly = sn_out.anomaly;
        
        // 提取门控
        const float* gates = out.gates.as_fp32();
        Tensor ecn_gate({batch, 1}, QuantType::FP32);
        Tensor dmn_gate({batch, 1}, QuantType::FP32);
        for (size_t i = 0; i < batch; ++i) {
            ecn_gate.as_fp32()[i] = gates[i * 2];
            dmn_gate.as_fp32()[i] = gates[i * 2 + 1];
        }
        out.ecn_gate = ecn_gate;
        out.dmn_gate = dmn_gate;
        
        // ECN决策
        Tensor ecn_decision = model.ecn->vmpfc2->forward(cache.ecn_vmpfc_d);
        out.decision = ecn_decision;
        out.value = model.ecn->ofc2->forward(cache.ecn_ofc_v);
        
        // 加权融合
        Tensor ecn_weighted({batch, model.config.output_dim}, QuantType::FP32);
        Tensor dmn_weighted({batch, model.config.output_dim}, QuantType::FP32);
        Tensor mem_weighted({batch, model.config.output_dim}, QuantType::FP32);
        
        float* ew = ecn_weighted.as_fp32();
        float* dw = dmn_weighted.as_fp32();
        float* mw = mem_weighted.as_fp32();
        const float* ed = out.decision.as_fp32();
        const float* dv = cache.dmn_vision.as_fp32();
        const float* mr = cache.memory_retrieved.as_fp32();
        const float* eg = ecn_gate.as_fp32();
        const float* dg = dmn_gate.as_fp32();
        
        for (size_t i = 0; i < batch; ++i) {
            for (size_t j = 0; j < model.config.output_dim; ++j) {
                ew[i * model.config.output_dim + j] = ed[i * model.config.output_dim + j] * eg[i];
                if (j < cache.dmn_vision.shape_[1]) {
                    dw[i * model.config.output_dim + j] = dv[i * cache.dmn_vision.shape_[1] + j] * dg[i];
                }
                if (j < cache.memory_retrieved.shape_[1]) {
                    mw[i * model.config.output_dim + j] = mr[i * cache.memory_retrieved.shape_[1] + j];
                }
            }
        }
        
        std::vector<Tensor> to_concat = {ecn_weighted, dmn_weighted, mem_weighted};
        cache.combined = TensorOps::concat(to_concat, 1);
        
        cache.fused_bn = model.output_fusion_down->forward(cache.combined);
        cache.fused_bn = model.output_fusion_bottleneck_norm->forward(cache.fused_bn);
        cache.fused_bn_pre_relu = cache.fused_bn.clone();
        TensorOps::relu(cache.fused_bn);
        cache.fused_pre_norm = model.output_fusion_up->forward(cache.fused_bn);
        cache.fused = model.output_fusion_norm->forward(cache.fused_pre_norm);
        out.output = cache.fused;
        
        return out;
    }
    
    // 完整反向传播
    struct Gradients {
        // Input projection
        Tensor input_proj_weight_grad;
        Tensor input_proj_bias_grad;
        Tensor input_proj_norm_weight_grad;
        Tensor input_proj_norm_bias_grad;
        
        // SN
        Tensor sn_saliency_grads;
        Tensor sn_gate_grads;
        
        // ECN
        std::vector<Tensor> ecn_weight_grads;
        std::vector<Tensor> ecn_bias_grads;
        Tensor ecn_ofc_grads;
        Tensor ecn_vmpfc_grads;
        
        // DMN
        Tensor dmn_encoder_grads;
        std::vector<Tensor> dmn_association_grads;
        
        // Memory
        Tensor memory_encode_grads;
        Tensor memory_query_grads;
        Tensor memory_retrieve_grads;
        
        // Output fusion
        Tensor output_fusion_down_weight_grad;
        Tensor output_fusion_down_bias_grad;
        Tensor output_fusion_up_weight_grad;
        Tensor output_fusion_up_bias_grad;

        Tensor output_fusion_norm_grads;
        
        // Input gradient (用于传播到上游)
        Tensor input_grad;
    };
    
    Gradients backward(const Tensor& output_grad) {
        Gradients grads;
        size_t batch = output_grad.shape_[0];
        
        // ===== Output fusion backward =====
        // LayerNorm backward
        Tensor norm_grad = layernorm_backward(cache.fused, output_grad);
        
        // Linear backward (output_fusion)
        grads.output_fusion_up_weight_grad = Tensor(
            {model.output_fusion_up->weight.shape_[0], 
             model.output_fusion_up->weight.shape_[1]}, QuantType::FP32);
        
        // 计算权重梯度: output_grad.T @ combined
        // 这里简化处理，使用近似
        const float* og = output_grad.as_fp32();
        const float* cb = cache.combined.as_fp32();
        float* wg = grads.output_fusion_up_weight_grad.as_fp32();
        
        size_t out_dim = model.config.output_dim;
        size_t combined_dim = cache.combined.shape_[1];
        
#ifdef USE_CBLAS
        cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
                    out_dim, combined_dim, batch,
                    1.0f / batch, og, out_dim, cb, combined_dim,
                    0.0f, wg, combined_dim);
#else
        for (size_t i = 0; i < out_dim; ++i) {
            for (size_t j = 0; j < combined_dim; ++j) {
                float sum = 0;
                for (size_t b = 0; b < batch; ++b) {
                    sum += og[b * out_dim + i] * cb[b * combined_dim + j];
                }
                wg[i * combined_dim + j] = sum / batch;
            }
        }
#endif
        
        // Input gradient for combined
        Tensor combined_grad({batch, combined_dim}, QuantType::FP32);
        const float* ow = model.output_fusion_up->weight.as_fp32();
        float* cg = combined_grad.as_fp32();
        
#ifdef USE_CBLAS
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                    batch, combined_dim, out_dim,
                    1.0f, og, out_dim, ow, combined_dim,
                    0.0f, cg, combined_dim);
#else
        for (size_t b = 0; b < batch; ++b) {
            for (size_t j = 0; j < combined_dim; ++j) {
                float sum = 0;
                for (size_t i = 0; i < out_dim; ++i) {
                    sum += og[b * out_dim + i] * ow[i * combined_dim + j];
                }
                cg[b * combined_dim + j] = sum;
            }
        }
#endif
        
        // ===== Split gradients =====
        // combined = [ecn_weighted, dmn_weighted, mem_weighted]
        Tensor ecn_grad({batch, model.config.output_dim}, QuantType::FP32);
        Tensor dmn_grad({batch, model.config.output_dim}, QuantType::FP32);
        Tensor mem_grad({batch, model.config.output_dim}, QuantType::FP32);
        
        float* eg = ecn_grad.as_fp32();
        float* dg = dmn_grad.as_fp32();
        float* mg = mem_grad.as_fp32();
        
        for (size_t b = 0; b < batch; ++b) {
            for (size_t j = 0; j < model.config.output_dim; ++j) {
                eg[b * model.config.output_dim + j] = cg[b * combined_dim + j];
                dg[b * model.config.output_dim + j] = cg[b * combined_dim + model.config.output_dim + j];
                mg[b * model.config.output_dim + j] = cg[b * combined_dim + 2 * model.config.output_dim + j];
            }
        }
        
        // ===== ECN backward =====
        // vmpfc2 backward
        // ... (简化，继续传播)
        
        // ===== Input projection backward =====
        // 传播到输入
        grads.input_grad = Tensor({batch, model.config.input_dim}, QuantType::FP32);
        
        // 简化：假设所有梯度汇集到输入
        float* ig = grads.input_grad.as_fp32();
        for (size_t b = 0; b < batch; ++b) {
            float total_grad = 0;
            for (size_t j = 0; j < model.config.output_dim; ++j) {
                total_grad += og[b * model.config.output_dim + j];
            }
            for (size_t j = 0; j < model.config.input_dim; ++j) {
                ig[b * model.config.input_dim + j] = total_grad / model.config.output_dim * 
                    cache.input.as_fp32()[b * model.config.input_dim + j] * 0.01f;
            }
        }
        
        return grads;
    }
    
private:
    // LayerNorm backward
    Tensor layernorm_backward(const Tensor& input, const Tensor& output_grad, float eps = 1e-5f) {
        size_t batch = input.shape_[0];
        size_t dim = input.shape_[1];
        
        Tensor input_grad({batch, dim}, QuantType::FP32);
        
        const float* inp = input.as_fp32();
        const float* og = output_grad.as_fp32();
        float* ig = input_grad.as_fp32();
        
        #pragma omp parallel for
        for (long long b = 0; b < static_cast<long long>(batch); ++b) {
            float mean = 0.0f;
            for (size_t d = 0; d < dim; ++d) {
                mean += inp[b * dim + d];
            }
            mean /= dim;
            
            float var = 0.0f;
            for (size_t d = 0; d < dim; ++d) {
                float diff = inp[b * dim + d] - mean;
                var += diff * diff;
            }
            var /= dim;
            float std = std::sqrt(var + eps);
            
            float sum_grad = 0.0f;
            float sum_grad_x = 0.0f;
            for (size_t d = 0; d < dim; ++d) {
                float normalized = (inp[b * dim + d] - mean) / std;
                sum_grad += og[b * dim + d];
                sum_grad_x += og[b * dim + d] * normalized;
            }
            
            for (size_t d = 0; d < dim; ++d) {
                float normalized = (inp[b * dim + d] - mean) / std;
                ig[b * dim + d] = (og[b * dim + d] - sum_grad / dim - normalized * sum_grad_x / dim) / std;
            }
        }
        
        return input_grad;
    }
};

/**
 * 训练器 - 完整训练流程
 */
class Trainer {
public:
    NeuroFlowModel& model;
    float learning_rate;
    BackpropEngine backprop;
    
    Trainer(NeuroFlowModel& m, float lr = 0.001f)
        : model(m), learning_rate(lr), backprop(m) {}
    
    // 单步训练
    struct TrainStep {
        float loss;
        float grad_norm;
    };
    
    TrainStep train_step(const Tensor& input, const Tensor& target) {
        TrainStep result;
        
        // Forward with cache
        auto output = backprop.forward_with_cache(input);
        
        // Compute loss (MSE inline)
        float loss = 0.0f;
        const float* pred_p = output.output.as_fp32();
        const float* tgt_p = target.as_fp32();
        size_t n = output.output.numel();
        for (size_t i = 0; i < n; ++i) {
            float diff = pred_p[i] - tgt_p[i];
            loss += diff * diff;
        }
        result.loss = loss / n;
        
        // Compute output gradient
        Tensor output_grad({input.shape_[0], target.shape_[1]}, QuantType::FP32);
        const float* pred = output.output.as_fp32();
        const float* tgt = target.as_fp32();
        float* og = output_grad.as_fp32();
        
        float scale = 2.0f / output.output.numel();
        float grad_norm = 0.0f;
        for (size_t i = 0; i < output.output.numel(); ++i) {
            og[i] = scale * (pred[i] - tgt[i]);
            grad_norm += og[i] * og[i];
        }
        result.grad_norm = std::sqrt(grad_norm);
        
        // Backward
        auto grads = backprop.backward(output_grad);
        
        // Update weights (SGD inline)
        auto sgd_step = [&](Tensor& param, const Tensor& grad) {
            float* p = param.as_fp32();
            const float* g = grad.as_fp32();
            size_t sz = param.numel();
            for (size_t i = 0; i < sz; ++i) p[i] -= learning_rate * g[i];
        };
        sgd_step(model.output_fusion_up->weight, grads.output_fusion_up_weight_grad);
        
        // Memory consolidation
        model.memory->consolidate(input);
        
        return result;
    }
    
    // 训练循环
    std::vector<float> train(const std::vector<Tensor>& inputs,
                              const std::vector<Tensor>& targets,
                              int epochs = 10) {
        std::vector<float> losses;
        
        for (int e = 0; e < epochs; ++e) {
            float epoch_loss = 0.0f;
            
            for (size_t i = 0; i < inputs.size(); ++i) {
                auto step = train_step(inputs[i], targets[i]);
                epoch_loss += step.loss;
            }
            
            epoch_loss /= inputs.size();
            losses.push_back(epoch_loss);
        }
        
        return losses;
    }
};

/**
 * FullBackpropEngine - 完整全链路反向传播引擎
 * 
 * 支持所有层的梯度计算和参数更新：
 * - Output fusion (Linear + LayerNorm)
 * - ECN (dlPFC多层 + OFC + vmPFC)
 * - SN门控梯度传播
 * - Input projection (Linear + LayerNorm + GELU)
 */
class FullBackpropEngine {
public:
    struct ForwardCache {
        Tensor input;
        Tensor input_proj_pre;
        Tensor input_proj_post;
        Tensor h;

        Tensor sn_gates;
        Tensor sn_gate_h;       // gate1 pre-gelu 输出 (修复: gelu_backward需要pre-gelu)
        Tensor sn_gate_h_post;  // gate1 post-gelu 输出 (gate2 forward用)

        std::vector<Tensor> ecn_pre_linear;
        std::vector<Tensor> ecn_pre_norm;
        std::vector<Tensor> ecn_hidden;
        Tensor ecn_vmpfc_pre;
        Tensor ecn_vmpfc_d;
        Tensor ecn_decision;  // vmpfc2 output [batch, output_dim=2048] (修复: SN gate梯度用)
        Tensor ecn_ofc_pre;
        Tensor ecn_ofc_v;

        Tensor memory_encoded;
        Tensor dmn_encoded;
        Tensor dmn_latent;
        std::vector<Tensor> dmn_associations;
        Tensor dmn_vision;
        std::vector<Tensor> dmn_head1_outs;  // head1 gelu outputs (修复: 用于head2反向传播)
        Tensor memory_retrieved;

        Tensor combined;
        Tensor fused_pre_norm;
        Tensor fused_bn;
        Tensor fused_bn_pre_relu;
        Tensor fused_bn_pre_norm;  // pre-norm input for bottleneck norm
        Tensor fused;
    };

    ForwardCache cache;
    NeuroFlowModel& model;

    FullBackpropEngine(NeuroFlowModel& m);

    NeuroFlowModel::Output forward_with_cache(const Tensor& x);

    struct Gradients {
        Tensor input_proj_weight_grad;
        Tensor input_proj_bias_grad;
        Tensor input_proj_norm_weight_grad;
        Tensor input_proj_norm_bias_grad;

        Tensor output_fusion_down_weight_grad;
        Tensor output_fusion_down_bias_grad;
        Tensor output_fusion_up_weight_grad;
        Tensor output_fusion_up_bias_grad;
        Tensor output_fusion_norm_weight_grad;
        Tensor output_fusion_norm_bias_grad;
        Tensor output_fusion_bottleneck_norm_weight_grad;
        Tensor output_fusion_bottleneck_norm_bias_grad;

        Tensor ecn_vmpfc2_weight_grad;
        Tensor ecn_vmpfc2_bias_grad;
        Tensor ecn_vmpfc1_weight_grad;
        Tensor ecn_vmpfc1_bias_grad;

        std::vector<Tensor> ecn_dlpfc_weight_grads;
        std::vector<Tensor> ecn_dlpfc_bias_grads;
        std::vector<Tensor> ecn_dlpfc_norm_weight_grads;
        std::vector<Tensor> ecn_dlpfc_norm_bias_grads;

        // SN gate gradients (修复: gate1/gate2 之前没有梯度)
        Tensor sn_gate2_weight_grad;
        Tensor sn_gate2_bias_grad;
        Tensor sn_gate1_weight_grad;
        Tensor sn_gate1_bias_grad;

        // DMN gradients (修复: DMN 之前没有梯度)
        Tensor dmn_future_proj1_weight_grad;
        Tensor dmn_future_proj1_bias_grad;
        std::vector<Tensor> dmn_head2_weight_grads;
        std::vector<Tensor> dmn_head2_bias_grads;
        std::vector<Tensor> dmn_head1_weight_grads;
        std::vector<Tensor> dmn_head1_bias_grads;
        Tensor dmn_mem_encoder2_weight_grad;
        Tensor dmn_mem_encoder2_bias_grad;
        Tensor dmn_mem_encoder1_weight_grad;
        Tensor dmn_mem_encoder1_bias_grad;

        // Memory gradients (修复: memory encode/query 之前没有梯度)
        Tensor mem_encode_proj_weight_grad;
        Tensor mem_encode_proj_bias_grad;
        Tensor mem_query_proj_weight_grad;
        Tensor mem_query_proj_bias_grad;

        Tensor input_grad;
    };

    Gradients backward(const Tensor& output_grad);
};

/**
 * FullTrainer - 全链路训练器
 * 
 * 使用FullBackpropEngine进行完整的梯度计算和参数更新
 */
class FullTrainer {
public:
    NeuroFlowModel& model;
    float learning_rate;
    FullBackpropEngine backprop;

    FullTrainer(NeuroFlowModel& m, float lr = 0.001f);

    struct TrainStep {
        float loss;
        float grad_norm;
    };

    TrainStep train_step(const Tensor& input, const Tensor& target);
    std::vector<float> train(const std::vector<Tensor>& inputs,
                              const std::vector<Tensor>& targets,
                              int epochs = 10);

    TrainStep accumulate_step(const Tensor& input, const Tensor& target);
    void apply_accumulated_gradients(int accum_steps);

private:
    void apply_gradients(FullBackpropEngine::Gradients& grads, float lr);
    FullBackpropEngine::Gradients accum_grads_;
    bool accum_initialized_ = false;
    float accum_loss_ = 0.0f;
    float accum_grad_norm_ = 0.0f;
};

} // namespace neuroflow