#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <iostream>
#include <vector>
#include "neuroflow/generative.hpp"
#include "neuroflow/model.hpp"
#include "neuroflow/tensor.hpp"

using namespace neuroflow;

int main() {
    std::cerr << "=== CausalLMHead 通路测试 ===" << std::endl;

    CausalLMConfig cfg;
    cfg.vocab_size = 100;
    cfg.d_model = 32;
    cfg.max_seq_len = 16;
    cfg.causal_window_size = 4;
    cfg.sae_k = 8;
    cfg.ntm_memory_slots = 4;
    cfg.use_mla = false;
    cfg.weight_tying = true;
    cfg.num_attn_layers = 1;
    cfg.num_attn_heads = 2;
    cfg.pooling = "mean";

    std::cerr << "1. 构造 CausalLMHead..." << std::endl;
    CausalLMHead lm_head(cfg);
    lm_head.tie_weights();
    std::cerr << "   OK" << std::endl;

    std::cerr << "2. forward (推理)..." << std::endl;
    std::vector<size_t> token_ids = {1, 5, 10, 20, 30};
    Tensor logits = lm_head.forward(token_ids);
    std::cerr << "   logits shape: [" << logits.shape_[0] << "," << logits.shape_[1] << "]" << std::endl;
    std::cerr << "   OK" << std::endl;

    std::cerr << "3. forward_for_training..." << std::endl;
    std::vector<size_t> train_ids = {1, 5, 10, 20};
    Tensor train_logits = lm_head.forward_for_training(train_ids);
    std::cerr << "   train_logits shape: [" << train_logits.shape_[0] << "," << train_logits.shape_[1] << "]" << std::endl;
    std::cerr << "   OK" << std::endl;

    std::cerr << "4. backward_from_logits..." << std::endl;
    size_t target_id = 30;
    const float* pred = train_logits.as_fp32();
    float max_val = -1e30f;
    for (size_t j = 0; j < cfg.vocab_size; ++j) {
        if (pred[j] > max_val) max_val = pred[j];
    }
    float sum_exp = 0.0f;
    for (size_t j = 0; j < cfg.vocab_size; ++j) {
        sum_exp += std::exp(pred[j] - max_val);
    }
    float loss = -(pred[target_id] - max_val - std::log(sum_exp));
    std::cerr << "   loss = " << loss << std::endl;

    Tensor logits_grad({1, cfg.vocab_size}, QuantType::FP32);
    float* lg = logits_grad.as_fp32();
    for (size_t j = 0; j < cfg.vocab_size; ++j) {
        float softmax_val = std::exp(pred[j] - max_val) / sum_exp;
        lg[j] = softmax_val;
        if (j == target_id) lg[j] -= 1.0f;
    }

    auto grads = lm_head.backward_from_logits(logits_grad);
    std::cerr << "   attn_grads.size() = " << grads.attn_grads.size() << std::endl;
    std::cerr << "   w_proj_weight_grad shape: [" << grads.w_proj_weight_grad.shape_[0] << "," << grads.w_proj_weight_grad.shape_[1] << "]" << std::endl;
    std::cerr << "   embed_grad shape: [" << grads.embed_grad.shape_[0] << "," << grads.embed_grad.shape_[1] << "]" << std::endl;
    std::cerr << "   dw_kernel_grad shape: [" << grads.dw_kernel_grad.shape_[0] << "," << grads.dw_kernel_grad.shape_[1] << "]" << std::endl;

    // NaN diagnostics
    auto check_nan = [](const std::string& name, const Tensor& t) {
        if (t.numel() == 0 || t.data_size_ == 0) return;
        const float* d = t.as_fp32();
        size_t nan_count = 0, inf_count = 0;
        float max_abs = 0.0f;
        for (size_t i = 0; i < t.numel(); ++i) {
            if (std::isnan(d[i])) nan_count++;
            else if (std::isinf(d[i])) inf_count++;
            else max_abs = std::max(max_abs, std::abs(d[i]));
        }
        std::cerr << "   [DIAG] " << name << ": nan=" << nan_count << " inf=" << inf_count << " max_abs=" << max_abs << std::endl;
    };
    check_nan("w_proj_weight_grad", grads.w_proj_weight_grad);
    std::cerr << "   [DEBUG] w_out_weight_grad numel=" << grads.w_out_weight_grad.numel() << std::endl;
    check_nan("w_out_weight_grad", grads.w_out_weight_grad);
    check_nan("ln_weight_grad", grads.ln_weight_grad);
    check_nan("sae_encode_weight_grad", grads.sae_encode_weight_grad);
    check_nan("sae_decode_weight_grad", grads.sae_decode_weight_grad);
    check_nan("ntm_read_weight_grad", grads.ntm_read_weight_grad);
    check_nan("dw_kernel_grad", grads.dw_kernel_grad);
    check_nan("pw_conv_weight_grad", grads.pw_conv_weight_grad);
    check_nan("embed_grad", grads.embed_grad);
    if (!grads.attn_grads.empty()) {
        check_nan("attn0.w_q_weight_grad", grads.attn_grads[0].w_q_weight_grad);
        check_nan("attn0.w_k_weight_grad", grads.attn_grads[0].w_k_weight_grad);
        check_nan("attn0.w_v_weight_grad", grads.attn_grads[0].w_v_weight_grad);
        check_nan("attn0.w_out_weight_grad", grads.attn_grads[0].w_out_weight_grad);
        check_nan("attn0.input_grad", grads.attn_grads[0].input_grad);
    }
    std::cerr << "   OK" << std::endl;

    std::cerr << "5. apply_lm_gradients..." << std::endl;
    lm_head.apply_lm_gradients(grads, 1e-5f);
    std::cerr << "   OK" << std::endl;

    std::cerr << "6. 第二次 forward_for_training (验证梯度更新有效)..." << std::endl;
    Tensor train_logits2 = lm_head.forward_for_training(train_ids);
    const float* pred2 = train_logits2.as_fp32();
    float max_val2 = -1e30f;
    for (size_t j = 0; j < cfg.vocab_size; ++j) {
        if (pred2[j] > max_val2) max_val2 = pred2[j];
    }
    float sum_exp2 = 0.0f;
    for (size_t j = 0; j < cfg.vocab_size; ++j) {
        sum_exp2 += std::exp(pred2[j] - max_val2);
    }
    float loss2 = -(pred2[target_id] - max_val2 - std::log(sum_exp2));
    std::cerr << "   loss2 = " << loss2 << std::endl;
    if (loss2 != loss) {
        std::cerr << "   梯度更新有效 (loss变化)" << std::endl;
    } else {
        std::cerr << "   警告: loss未变化" << std::endl;
    }

    std::cerr << "7. 保存 LM Head..." << std::endl;
    {
        auto sl = [](std::ofstream& o, const std::string& n, const Tensor& t) {
            uint32_t nl = n.size(); o.write((char*)&nl, 4); o.write(n.data(), nl);
            uint32_t nd = t.shape_.size(); o.write((char*)&nd, 4);
            for (auto d : t.shape_) { uint32_t dd = d; o.write((char*)&dd, 4); }
            uint32_t ds = t.data_size_; o.write((char*)&ds, 4);
            o.write((char*)t.data_.get(), ds);
        };
        std::ofstream o("D:/neuroflow-C++/test_run/lm_head_test.nfv1", std::ios::binary);
        o.write("LMH2", 4);
        sl(o, "w_embed", lm_head.w_embed_);
        sl(o, "w_proj.weight", lm_head.w_proj_->weight);
        sl(o, "w_out.weight", lm_head.w_out_->weight);
        sl(o, "dw_kernel", lm_head.dw_kernel_);
        sl(o, "sae_encode.weight", lm_head.sae_w_encode_->weight);
        sl(o, "sae_decode.weight", lm_head.sae_w_decode_->weight);
        sl(o, "ntm_read.weight", lm_head.ntm_w_read_->weight);
        sl(o, "ntm_write.weight", lm_head.ntm_w_write_->weight);
        sl(o, "ntm_erase.weight", lm_head.ntm_w_erase_->weight);
        sl(o, "ntm_memory", lm_head.ntm_memory_);
        sl(o, "ln.weight", lm_head.ln_->weight);
        sl(o, "ln.bias", lm_head.ln_->bias);
        for (size_t i = 0; i < lm_head.attn_layers_.size(); ++i) {
            std::string p = "attn" + std::to_string(i) + ".";
            sl(o, p + "w_q.weight", lm_head.attn_layers_[i]->w_q->weight);
            sl(o, p + "w_q.bias", lm_head.attn_layers_[i]->w_q->bias);
            sl(o, p + "w_k.weight", lm_head.attn_layers_[i]->w_k->weight);
            sl(o, p + "w_k.bias", lm_head.attn_layers_[i]->w_k->bias);
            sl(o, p + "w_v.weight", lm_head.attn_layers_[i]->w_v->weight);
            sl(o, p + "w_v.bias", lm_head.attn_layers_[i]->w_v->bias);
            sl(o, p + "w_out.weight", lm_head.attn_layers_[i]->w_out->weight);
            sl(o, p + "w_out.bias", lm_head.attn_layers_[i]->w_out->bias);
            sl(o, p + "norm.weight", lm_head.attn_layers_[i]->norm->weight);
            sl(o, p + "norm.bias", lm_head.attn_layers_[i]->norm->bias);
        }
        uint32_t z = 0; o.write((char*)&z, 4); o.close();
    }
    std::cerr << "   OK" << std::endl;

    std::cerr << "=== 所有通路测试通过! ===" << std::endl;
    return 0;
}