#ifndef NEUROFLOW_DPO_HPP
#define NEUROFLOW_DPO_HPP

#include <cstddef>
#include <random>
#include <string>
#include <vector>

#include "adamw.hpp"
#include "alignment_common.hpp"
#include "causal_lm.hpp"
#include "scheduler.hpp"
#include "tokenizer.hpp"

namespace neuroflow {

struct DPOTrainConfig {
    std::string data_path;
    std::string sft_ckpt_path;
    std::string tokenizer_path;
    std::string output_dir;
    float learning_rate = 1e-6f;
    int epochs = 3;
    size_t max_seq_len = 512;
    float warmup_ratio = 0.05f;
    float weight_decay = 0.0f;
    float grad_clip = 1.0f;
    float beta = 0.1f;
    float adam_beta1 = 0.9f;
    float adam_beta2 = 0.999f;
    float adam_eps = 1e-8f;
    size_t save_interval = 1000;
    size_t log_interval = 10;
    unsigned seed = 42;
};

class DPODataLoader {
public:
    DPODataLoader(const std::string& jsonl_path, size_t max_samples = 0);

    bool has_next() const;
    DPOSample next();
    void reset();
    void shuffle(std::mt19937& rng);
    size_t total_samples() const { return samples_.size(); }
    size_t invalid_count() const { return invalid_count_; }

private:
    std::vector<DPOSample> samples_;
    size_t cursor_ = 0;
    size_t invalid_count_ = 0;
};

struct DPOLossOutput {
    float loss;
    float alpha;
    float reward_chosen;
    float reward_rejected;
};

float compute_log_prob(CausalLMHead& model, const std::vector<size_t>& token_ids,
                       size_t prompt_len, size_t vocab_size);

DPOLossOutput compute_dpo_loss(float log_prob_chosen_policy,
                                float log_prob_rejected_policy,
                                float log_prob_chosen_ref,
                                float log_prob_rejected_ref,
                                float beta);

class DPOTrainer {
public:
    DPOTrainConfig config;

    DPOTrainer(const DPOTrainConfig& cfg);

    void train();
    float train_on_sample(const DPOSample& sample);

private:
    std::unique_ptr<CausalLMHead> policy_;
    std::unique_ptr<CausalLMHead> reference_;
    std::unique_ptr<BPETokenizer> tokenizer_;
    std::unique_ptr<AdamW> optimizer_;
    std::unique_ptr<CosineScheduler> scheduler_;

    float compute_w_embed_checksum();
};

} // namespace neuroflow

#endif // NEUROFLOW_DPO_HPP