#include "neuroflow/sft.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <numeric>
#include <sstream>

namespace neuroflow {

SFTDataLoader::SFTDataLoader(const std::string& jsonl_path, size_t max_samples) {
    std::ifstream ifs(jsonl_path);
    if (!ifs) {
        std::cerr << "SFT数据文件无法打开: " << jsonl_path << std::endl;
        return;
    }
    std::string line;
    while (std::getline(ifs, line)) {
        if (line.empty() || line[0] == '#') continue;
        std::string instruction = extract_json_string(line, "instruction");
        std::string response = extract_json_string(line, "response");
        unescape_json(instruction);
        unescape_json(response);
        if (instruction.empty() || response.empty()) {
            invalid_count_++;
            continue;
        }
        samples_.push_back({instruction, response});
        if (max_samples > 0 && samples_.size() >= max_samples) break;
    }
    std::cerr << "SFT数据加载: " << samples_.size() << " 样本, "
              << invalid_count_ << " 无效" << std::endl;
}

bool SFTDataLoader::has_next() const {
    return cursor_ < samples_.size();
}

SFTSample SFTDataLoader::next() {
    return samples_[cursor_++];
}

void SFTDataLoader::reset() {
    cursor_ = 0;
}

void SFTDataLoader::shuffle(std::mt19937& rng) {
    std::shuffle(samples_.begin(), samples_.end(), rng);
}

SFTTrainingTensors build_sft_training_tensors(const SFTSample& sample,
                                                BPETokenizer& tokenizer,
                                                size_t max_seq_len) {
    SFTTrainingTensors result;
    std::string prompt = sample.instruction + "\n";
    std::string full_text = prompt + sample.response;

    std::vector<size_t> prompt_ids = tokenizer.encode(prompt, max_seq_len);
    std::vector<size_t> full_ids = tokenizer.encode(full_text, max_seq_len);

    if (full_ids.size() > max_seq_len) {
        full_ids.resize(max_seq_len);
    }
    if (full_ids.size() < 2) {
        result.instruction_len = 0;
        return result;
    }

    result.input_ids.assign(full_ids.begin(), full_ids.end() - 1);
    result.target_ids.assign(full_ids.begin() + 1, full_ids.end());
    result.instruction_len = std::min(prompt_ids.size(), result.input_ids.size());

    result.loss_mask.resize(result.target_ids.size(), 0.0f);
    for (size_t i = result.instruction_len; i < result.target_ids.size(); ++i) {
        result.loss_mask[i] = 1.0f;
    }

    return result;
}

MaskedCEOutput masked_cross_entropy(const Tensor& logits, size_t vocab_size,
                                     const std::vector<size_t>& target_ids,
                                     const std::vector<float>& loss_mask) {
    MaskedCEOutput output;
    output.loss = 0.0f;
    output.valid_token_count = 0;

    size_t seq_len = target_ids.size();
    if (seq_len == 0 || logits.numel() == 0) {
        output.logits_grad = Tensor({1, vocab_size}, QuantType::FP32);
        memset(output.logits_grad.as_fp32(), 0, output.logits_grad.data_size_);
        return output;
    }

    output.logits_grad = Tensor({seq_len, vocab_size}, QuantType::FP32);
    memset(output.logits_grad.as_fp32(), 0, output.logits_grad.data_size_);
    float* lg = output.logits_grad.as_fp32();

    for (size_t t = 0; t < seq_len; ++t) {
        if (loss_mask[t] < 0.5f) continue;

        size_t target = target_ids[t];
        if (target >= vocab_size) target = 1;

        const float* pred = logits.as_fp32() + t * vocab_size;

        float max_val = -1e30f;
        for (size_t j = 0; j < vocab_size; ++j) {
            if (pred[j] > max_val) max_val = pred[j];
        }
        float sum_exp = 0.0f;
        for (size_t j = 0; j < vocab_size; ++j) {
            sum_exp += std::exp(pred[j] - max_val);
        }
        float log_sum_exp = max_val + std::log(sum_exp);
        float token_loss = -(pred[target] - log_sum_exp);

        if (!std::isfinite(token_loss)) continue;

        output.loss += token_loss;
        output.valid_token_count++;

        float* grad_row = lg + t * vocab_size;
        for (size_t j = 0; j < vocab_size; ++j) {
            float softmax_val = std::exp(pred[j] - max_val) / sum_exp;
            grad_row[j] = softmax_val;
            if (j == target) grad_row[j] -= 1.0f;
        }
    }

    if (output.valid_token_count > 0) {
        output.loss /= static_cast<float>(output.valid_token_count);
        float inv_n = 1.0f / static_cast<float>(output.valid_token_count);
        for (size_t i = 0; i < seq_len * vocab_size; ++i) {
            lg[i] *= inv_n;
        }
    }

    return output;
}

SFTTrainer::SFTTrainer(const SFTTrainConfig& cfg) : config(cfg) {
    CausalLMConfig lm_config;
    lm_config.vocab_size = 128000;
    lm_config.d_model = 512;
    lm_config.max_seq_len = cfg.max_seq_len;
    lm_config.num_attn_layers = 4;
    lm_config.num_attn_heads = 8;
    lm_config.n_kv_heads = 2;
    lm_config.use_rope = true;
    lm_config.use_qk_norm = true;
    lm_config.use_swiglu = true;
    lm_config.use_bridge = true;
    lm_config.weight_tying = true;
    lm_config.pooling = "last";

    model_ = std::make_unique<CausalLMHead>(lm_config);
    if (!cfg.ckpt_path.empty()) {
        load_lm_checkpoint(*model_, cfg.ckpt_path);
    }

    tokenizer_ = std::make_unique<BPETokenizer>(cfg.tokenizer_path);

    size_t total_steps = 0;
    {
        SFTDataLoader tmp_loader(cfg.data_path);
        total_steps = tmp_loader.total_samples() * cfg.epochs;
    }

    optimizer_ = std::make_unique<AdamW>(cfg.learning_rate, cfg.adam_beta1,
                                          cfg.adam_beta2, cfg.adam_eps,
                                          cfg.weight_decay);

    model_->register_trainable_params(*optimizer_, cfg.learning_rate, cfg.weight_decay);

    scheduler_ = std::make_unique<CosineScheduler>(cfg.learning_rate, total_steps,
                                                     0.1f, cfg.warmup_ratio);
}

void SFTTrainer::train() {
    SFTDataLoader loader(config.data_path);
    if (loader.total_samples() == 0) {
        std::cerr << "SFT训练: 无有效样本" << std::endl;
        return;
    }

    std::cerr << "SFT训练开始: " << loader.total_samples() << " 样本, "
              << config.epochs << " epochs" << std::endl;

    model_->train();

    size_t global_step = 0;
    auto train_start = std::chrono::steady_clock::now();

    for (int epoch = 0; epoch < config.epochs; ++epoch) {
        auto epoch_start = std::chrono::steady_clock::now();
        loader.reset();
        std::mt19937 shuffle_rng(config.seed + epoch);
        loader.shuffle(shuffle_rng);

        float epoch_loss = 0.0f;
        size_t step_count = 0;

        while (loader.has_next()) {
            SFTSample sample = loader.next();
            float lr = scheduler_->get_lr(global_step);
            optimizer_->set_lr(lr);
            float sample_loss = train_on_sample(sample);
            global_step++;

            if (std::isfinite(sample_loss)) {
                epoch_loss += sample_loss;
                step_count++;
            }

            if (config.log_interval > 0 && global_step % config.log_interval == 0) {
                auto now = std::chrono::steady_clock::now();
                float elapsed = static_cast<float>(
                    std::chrono::duration<double>(now - train_start).count());
                std::cerr << "[SFT] step=" << global_step
                          << " epoch=" << (epoch + 1)
                          << " loss=" << sample_loss
                          << " lr=" << optimizer_->get_lr()
                          << " elapsed=" << elapsed << "s" << std::endl;
            }

            if (config.save_interval > 0 && global_step % config.save_interval == 0) {
                std::string cdir = config.output_dir + "/checkpoint_step" + std::to_string(global_step);
                std::filesystem::create_directories(cdir);
                save_lm_checkpoint(*model_, cdir + "/lm_head.nfv1");
                std::cerr << "[SFT] Checkpoint: step=" << global_step << std::endl;
            }
        }

        float avg_loss = (step_count > 0) ? epoch_loss / static_cast<float>(step_count) : 0.0f;
        auto epoch_end = std::chrono::steady_clock::now();
        float epoch_elapsed = static_cast<float>(
            std::chrono::duration<double>(epoch_end - epoch_start).count());
        std::cerr << "[SFT] Epoch " << (epoch + 1) << "/" << config.epochs
                  << " avg_loss=" << avg_loss
                  << " elapsed=" << epoch_elapsed << "s" << std::endl;

        std::string cdir = config.output_dir + "/checkpoint_epoch" + std::to_string(epoch + 1);
        std::filesystem::create_directories(cdir);
        save_lm_checkpoint(*model_, cdir + "/lm_head.nfv1");
    }

    std::filesystem::create_directories(config.output_dir);
    save_lm_checkpoint(*model_, config.output_dir + "/lm_head_sft_final.nfv1");
    std::cerr << "SFT训练完成, 模型已保存: " << config.output_dir << "/lm_head_sft_final.nfv1" << std::endl;
}

float SFTTrainer::train_on_sample(const SFTSample& sample) {
    SFTTrainingTensors tensors = build_sft_training_tensors(sample, *tokenizer_, config.max_seq_len);
    if (tensors.input_ids.empty() || tensors.instruction_len == 0) {
        return 0.0f;
    }

    size_t seq_len = tensors.input_ids.size();
    size_t vocab_size = model_->config_.vocab_size;

    float total_loss = 0.0f;
    float total_grad_norm = 0.0f;

    for (size_t t = 0; t < seq_len; ++t) {
        std::vector<size_t> input_prefix(tensors.input_ids.begin(),
                                          tensors.input_ids.begin() + t + 1);

        Tensor logits = model_->forward_for_training(input_prefix);

        size_t target_id = tensors.target_ids[t];
        if (target_id >= vocab_size) target_id = 1;

        const float* pred = logits.as_fp32();
        float max_val = -1e30f;
        for (size_t j = 0; j < vocab_size; ++j) {
            if (pred[j] > max_val) max_val = pred[j];
        }
        float sum_exp = 0.0f;
        for (size_t j = 0; j < vocab_size; ++j) {
            sum_exp += std::exp(pred[j] - max_val);
        }
        float log_sum_exp = max_val + std::log(sum_exp);
        float token_loss = -(pred[target_id] - log_sum_exp);

        if (tensors.loss_mask[t] < 0.5f) {
            continue;
        }

        if (!std::isfinite(token_loss)) {
            std::cerr << "[SFT WARN] NaN/Inf loss at token " << t << ", skipping" << std::endl;
            continue;
        }

        Tensor logits_grad({1, vocab_size}, QuantType::FP32);
        float* lg = logits_grad.as_fp32();
        float grad_norm = 0.0f;
        for (size_t j = 0; j < 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;
            grad_norm += lg[j] * lg[j];
        }

        float gn = std::sqrt(grad_norm);
        float clip_scale = 1.0f;
        if (!std::isfinite(gn) || (gn > config.grad_clip && config.grad_clip > 0.0f)) {
            clip_scale = config.grad_clip / gn;
        }
        if (clip_scale < 1.0f) {
            float* lg2 = logits_grad.as_fp32();
            for (size_t j = 0; j < vocab_size; ++j) lg2[j] *= clip_scale;
        }

        auto lm_grads = model_->backward_from_logits(logits_grad);

        model_->assign_grads_to_optimizer(*optimizer_, lm_grads);
        optimizer_->step();

        total_loss += token_loss;
        total_grad_norm += grad_norm;
    }

    size_t valid_count = 0;
    for (size_t i = 0; i < tensors.loss_mask.size(); ++i) {
        if (tensors.loss_mask[i] > 0.5f) valid_count++;
    }

    return (valid_count > 0) ? total_loss / static_cast<float>(valid_count) : 0.0f;
}

} // namespace neuroflow