// Optimization R interface
// Training and fine-tuning support through ggml-opt

#include <R.h>
#include <Rinternals.h>
#include "ggml.h"
#include "ggml-backend.h"
#include "ggml-cpu.h"
#include "ggml-opt.h"

extern int ggmlR_get_n_threads(void);

// Update all CPU backends in scheduler with current ggmlR thread setting
static void opt_sync_cpu_threads(ggml_backend_sched_t sched) {
    int n_threads = ggmlR_get_n_threads();
    int n = ggml_backend_sched_get_n_backends(sched);
    for (int i = 0; i < n; i++) {
        ggml_backend_t b = ggml_backend_sched_get_backend(sched, i);
        if (ggml_backend_is_cpu(b)) {
            ggml_backend_cpu_set_n_threads(b, n_threads);
        }
    }
}

// ============================================================================
// Loss Type Constants
// ============================================================================

SEXP R_ggml_opt_loss_type_mean(void) {
    return ScalarInteger(GGML_OPT_LOSS_TYPE_MEAN);
}

SEXP R_ggml_opt_loss_type_sum(void) {
    return ScalarInteger(GGML_OPT_LOSS_TYPE_SUM);
}

SEXP R_ggml_opt_loss_type_cross_entropy(void) {
    return ScalarInteger(GGML_OPT_LOSS_TYPE_CROSS_ENTROPY);
}

SEXP R_ggml_opt_loss_type_mse(void) {
    return ScalarInteger(GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR);
}

// ============================================================================
// Optimizer Type Constants
// ============================================================================

SEXP R_ggml_opt_optimizer_type_adamw(void) {
    return ScalarInteger(GGML_OPT_OPTIMIZER_TYPE_ADAMW);
}

SEXP R_ggml_opt_optimizer_type_sgd(void) {
    return ScalarInteger(GGML_OPT_OPTIMIZER_TYPE_SGD);
}

// ============================================================================
// Dataset Functions
// ============================================================================

// Create a new dataset
SEXP R_ggml_opt_dataset_init(SEXP type_data, SEXP type_label,
                              SEXP ne_datapoint, SEXP ne_label,
                              SEXP ndata, SEXP ndata_shard) {
    enum ggml_type t_data = (enum ggml_type)asInteger(type_data);
    enum ggml_type t_label = (enum ggml_type)asInteger(type_label);
    int64_t n_dp = (int64_t)asReal(ne_datapoint);
    int64_t n_lb = (int64_t)asReal(ne_label);
    int64_t n_data = (int64_t)asReal(ndata);
    int64_t n_shard = (int64_t)asReal(ndata_shard);

    ggml_opt_dataset_t dataset = ggml_opt_dataset_init(t_data, t_label, n_dp, n_lb, n_data, n_shard);

    if (dataset == NULL) {
        error("Failed to create dataset");
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(dataset, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Free dataset
SEXP R_ggml_opt_dataset_free(SEXP dataset_ptr) {
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    if (dataset != NULL) {
        ggml_opt_dataset_free(dataset);
        R_ClearExternalPtr(dataset_ptr);
    }

    return R_NilValue;
}

// Get number of data points
SEXP R_ggml_opt_dataset_ndata(SEXP dataset_ptr) {
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }

    int64_t n = ggml_opt_dataset_ndata(dataset);
    return ScalarReal((double)n);
}

// Get data tensor
SEXP R_ggml_opt_dataset_data(SEXP dataset_ptr) {
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }

    struct ggml_tensor * data = ggml_opt_dataset_data(dataset);

    SEXP ptr = PROTECT(R_MakeExternalPtr(data, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get labels tensor
SEXP R_ggml_opt_dataset_labels(SEXP dataset_ptr) {
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }

    struct ggml_tensor * labels = ggml_opt_dataset_labels(dataset);

    if (labels == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(labels, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Shuffle dataset
SEXP R_ggml_opt_dataset_shuffle(SEXP opt_ctx_ptr, SEXP dataset_ptr, SEXP idata) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }
    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }

    int64_t i = (int64_t)asReal(idata);
    ggml_opt_dataset_shuffle(opt_ctx, dataset, i);

    return R_NilValue;
}

// Get batch from dataset
SEXP R_ggml_opt_dataset_get_batch(SEXP dataset_ptr, SEXP data_batch_ptr,
                                   SEXP labels_batch_ptr, SEXP ibatch) {
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);
    struct ggml_tensor * data_batch = (struct ggml_tensor *)R_ExternalPtrAddr(data_batch_ptr);
    struct ggml_tensor * labels_batch = NULL;

    if (labels_batch_ptr != R_NilValue) {
        labels_batch = (struct ggml_tensor *)R_ExternalPtrAddr(labels_batch_ptr);
    }

    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }
    if (data_batch == NULL) {
        error("Invalid data_batch tensor pointer");
    }

    int64_t i = (int64_t)asReal(ibatch);
    ggml_opt_dataset_get_batch(dataset, data_batch, labels_batch, i);

    return R_NilValue;
}

// ============================================================================
// Optimizer Context Functions
// ============================================================================

// Get default optimizer params
SEXP R_ggml_opt_default_params(SEXP sched_ptr, SEXP loss_type) {
    ggml_backend_sched_t sched = (ggml_backend_sched_t)R_ExternalPtrAddr(sched_ptr);

    if (sched == NULL) {
        error("Invalid scheduler pointer");
    }

    enum ggml_opt_loss_type lt = (enum ggml_opt_loss_type)asInteger(loss_type);
    struct ggml_opt_params params = ggml_opt_default_params(sched, lt);

    // Return as a list
    SEXP result = PROTECT(allocVector(VECSXP, 4));
    SEXP names = PROTECT(allocVector(STRSXP, 4));

    SET_STRING_ELT(names, 0, mkChar("loss_type"));
    SET_STRING_ELT(names, 1, mkChar("build_type"));
    SET_STRING_ELT(names, 2, mkChar("opt_period"));
    SET_STRING_ELT(names, 3, mkChar("optimizer"));

    SET_VECTOR_ELT(result, 0, ScalarInteger(params.loss_type));
    SET_VECTOR_ELT(result, 1, ScalarInteger(params.build_type));
    SET_VECTOR_ELT(result, 2, ScalarInteger(params.opt_period));
    SET_VECTOR_ELT(result, 3, ScalarInteger(params.optimizer));

    setAttrib(result, R_NamesSymbol, names);
    UNPROTECT(2);
    return result;
}

// Initialize optimizer context
SEXP R_ggml_opt_init(SEXP sched_ptr, SEXP loss_type, SEXP optimizer_type, SEXP opt_period,
                     SEXP ctx_compute_ptr, SEXP inputs_ptr, SEXP outputs_ptr) {
    ggml_backend_sched_t sched = (ggml_backend_sched_t)R_ExternalPtrAddr(sched_ptr);

    if (sched == NULL) {
        error("Invalid scheduler pointer");
    }

    enum ggml_opt_loss_type lt = (enum ggml_opt_loss_type)asInteger(loss_type);
    struct ggml_opt_params params = ggml_opt_default_params(sched, lt);

    params.optimizer = (enum ggml_opt_optimizer_type)asInteger(optimizer_type);
    params.opt_period = asInteger(opt_period);

    // Optional: set ctx_compute, inputs, outputs for static graph mode
    if (ctx_compute_ptr != R_NilValue) {
        params.ctx_compute = (struct ggml_context *)R_ExternalPtrAddr(ctx_compute_ptr);
    }
    if (inputs_ptr != R_NilValue) {
        params.inputs = (struct ggml_tensor *)R_ExternalPtrAddr(inputs_ptr);
    }
    if (outputs_ptr != R_NilValue) {
        params.outputs = (struct ggml_tensor *)R_ExternalPtrAddr(outputs_ptr);
    }

    opt_sync_cpu_threads(sched);

    ggml_opt_context_t opt_ctx = ggml_opt_init(params);

    if (opt_ctx == NULL) {
        error("Failed to initialize optimizer context");
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(opt_ctx, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Free optimizer context
SEXP R_ggml_opt_free(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx != NULL) {
        ggml_opt_free(opt_ctx);
        R_ClearExternalPtr(opt_ctx_ptr);
    }

    return R_NilValue;
}

// Reset optimizer context
SEXP R_ggml_opt_reset(SEXP opt_ctx_ptr, SEXP optimizer) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    bool reset_optimizer = asLogical(optimizer);
    ggml_opt_reset(opt_ctx, reset_optimizer);

    return R_NilValue;
}

// Check if using static graphs
SEXP R_ggml_opt_static_graphs(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    bool is_static = ggml_opt_static_graphs(opt_ctx);
    return ScalarLogical(is_static);
}

// Get inputs tensor
SEXP R_ggml_opt_inputs(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    struct ggml_tensor * inputs = ggml_opt_inputs(opt_ctx);

    if (inputs == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(inputs, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get outputs tensor
SEXP R_ggml_opt_outputs(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    struct ggml_tensor * outputs = ggml_opt_outputs(opt_ctx);

    if (outputs == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(outputs, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get labels tensor
SEXP R_ggml_opt_labels(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    struct ggml_tensor * labels = ggml_opt_labels(opt_ctx);

    if (labels == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(labels, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get loss tensor
SEXP R_ggml_opt_loss(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    struct ggml_tensor * loss = ggml_opt_loss(opt_ctx);

    if (loss == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(loss, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get predictions tensor
SEXP R_ggml_opt_pred(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    struct ggml_tensor * pred = ggml_opt_pred(opt_ctx);

    if (pred == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(pred, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get number of correct predictions tensor
SEXP R_ggml_opt_ncorrect(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    struct ggml_tensor * ncorrect = ggml_opt_ncorrect(opt_ctx);

    if (ncorrect == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(ncorrect, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get optimizer type
SEXP R_ggml_opt_context_optimizer_type(SEXP opt_ctx_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    enum ggml_opt_optimizer_type opt_type = ggml_opt_context_optimizer_type(opt_ctx);
    return ScalarInteger((int)opt_type);
}

// Get optimizer name
SEXP R_ggml_opt_optimizer_name(SEXP optimizer_type) {
    enum ggml_opt_optimizer_type opt_type = (enum ggml_opt_optimizer_type)asInteger(optimizer_type);
    const char * name = ggml_opt_optimizer_name(opt_type);
    return mkString(name);
}

// ============================================================================
// Result Functions
// ============================================================================

// Initialize result
SEXP R_ggml_opt_result_init(void) {
    ggml_opt_result_t result = ggml_opt_result_init();

    if (result == NULL) {
        error("Failed to initialize result");
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(result, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Free result
SEXP R_ggml_opt_result_free(SEXP result_ptr) {
    ggml_opt_result_t result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);

    if (result != NULL) {
        ggml_opt_result_free(result);
        R_ClearExternalPtr(result_ptr);
    }

    return R_NilValue;
}

// Reset result
SEXP R_ggml_opt_result_reset(SEXP result_ptr) {
    ggml_opt_result_t result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);

    if (result == NULL) {
        error("Invalid result pointer");
    }

    ggml_opt_result_reset(result);
    return R_NilValue;
}

// Get number of data points from result
SEXP R_ggml_opt_result_ndata(SEXP result_ptr) {
    ggml_opt_result_t result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);

    if (result == NULL) {
        error("Invalid result pointer");
    }

    int64_t ndata;
    ggml_opt_result_ndata(result, &ndata);
    return ScalarReal((double)ndata);
}

// Get loss from result
SEXP R_ggml_opt_result_loss(SEXP result_ptr) {
    ggml_opt_result_t result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);

    if (result == NULL) {
        error("Invalid result pointer");
    }

    double loss, unc;
    ggml_opt_result_loss(result, &loss, &unc);

    SEXP r = PROTECT(allocVector(REALSXP, 2));
    REAL(r)[0] = loss;
    REAL(r)[1] = unc;

    SEXP names = PROTECT(allocVector(STRSXP, 2));
    SET_STRING_ELT(names, 0, mkChar("loss"));
    SET_STRING_ELT(names, 1, mkChar("uncertainty"));
    setAttrib(r, R_NamesSymbol, names);

    UNPROTECT(2);
    return r;
}

// Get accuracy from result
SEXP R_ggml_opt_result_accuracy(SEXP result_ptr) {
    ggml_opt_result_t result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);

    if (result == NULL) {
        error("Invalid result pointer");
    }

    double accuracy, unc;
    ggml_opt_result_accuracy(result, &accuracy, &unc);

    SEXP r = PROTECT(allocVector(REALSXP, 2));
    REAL(r)[0] = accuracy;
    REAL(r)[1] = unc;

    SEXP names = PROTECT(allocVector(STRSXP, 2));
    SET_STRING_ELT(names, 0, mkChar("accuracy"));
    SET_STRING_ELT(names, 1, mkChar("uncertainty"));
    setAttrib(r, R_NamesSymbol, names);

    UNPROTECT(2);
    return r;
}

// ============================================================================
// Computation Functions
// ============================================================================

// Allocate graph for evaluation
SEXP R_ggml_opt_alloc(SEXP opt_ctx_ptr, SEXP backward) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    bool do_backward = asLogical(backward);
    ggml_opt_alloc(opt_ctx, do_backward);

    return R_NilValue;
}

// Evaluate (forward pass, optionally backward pass)
SEXP R_ggml_opt_eval(SEXP opt_ctx_ptr, SEXP result_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);
    ggml_opt_result_t result = NULL;

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }

    if (result_ptr != R_NilValue) {
        result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);
    }

    ggml_opt_eval(opt_ctx, result);

    return R_NilValue;
}

// ============================================================================
// High-Level Training Function
// ============================================================================

// Helper: get batch size from last dimension of tensor
static int64_t r_ggml_opt_batch_size(const struct ggml_tensor * t) {
    return t->ne[ggml_n_dims(t) - 1];
}

// Fit model to dataset, returning history (loss/accuracy per epoch)
SEXP R_ggml_opt_fit(SEXP sched_ptr, SEXP ctx_compute_ptr,
                    SEXP inputs_ptr, SEXP outputs_ptr,
                    SEXP dataset_ptr, SEXP loss_type, SEXP optimizer_type,
                    SEXP nepoch, SEXP nbatch_logical, SEXP val_split, SEXP silent) {

    ggml_backend_sched_t sched = (ggml_backend_sched_t)R_ExternalPtrAddr(sched_ptr);
    struct ggml_context * ctx_compute = (struct ggml_context *)R_ExternalPtrAddr(ctx_compute_ptr);
    struct ggml_tensor * inputs = (struct ggml_tensor *)R_ExternalPtrAddr(inputs_ptr);
    struct ggml_tensor * outputs = (struct ggml_tensor *)R_ExternalPtrAddr(outputs_ptr);
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    if (sched == NULL) {
        error("Invalid scheduler pointer");
    }
    if (ctx_compute == NULL) {
        error("Invalid compute context pointer");
    }
    if (inputs == NULL) {
        error("Invalid inputs tensor pointer");
    }
    if (outputs == NULL) {
        error("Invalid outputs tensor pointer");
    }
    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }

    enum ggml_opt_loss_type lt = (enum ggml_opt_loss_type)asInteger(loss_type);
    enum ggml_opt_optimizer_type ot = (enum ggml_opt_optimizer_type)asInteger(optimizer_type);
    int64_t n_epoch = (int64_t)asReal(nepoch);
    int64_t n_batch = (int64_t)asReal(nbatch_logical);
    float v_split = (float)asReal(val_split);
    bool is_silent = asLogical(silent);

    // Compute parameters (mirroring ggml_opt_fit logic from ggml-opt.cpp)
    const int64_t ndata = ggml_opt_dataset_data(dataset)->ne[1];
    const int64_t nbatch_physical = r_ggml_opt_batch_size(inputs);
    const int64_t opt_period = n_batch / nbatch_physical;
    const int64_t nbatches_logical = ndata / n_batch;
    const int64_t ibatch_split = (int64_t)(((1.0f - v_split) * nbatches_logical)) * opt_period;
    const int64_t idata_split = ibatch_split * nbatch_physical;

    int64_t epoch = 1;

    struct ggml_opt_params params = ggml_opt_default_params(sched, lt);
    params.ctx_compute     = ctx_compute;
    params.inputs          = inputs;
    params.outputs         = outputs;
    params.opt_period      = opt_period;
    params.get_opt_pars    = ggml_opt_get_default_optimizer_params;
    params.get_opt_pars_ud = &epoch;
    params.optimizer       = ot;
    ggml_opt_context_t opt_ctx = ggml_opt_init(params);

    if (n_batch < ndata) {
        ggml_opt_dataset_shuffle(opt_ctx, dataset, -1);
    }

    ggml_opt_result_t result_train = ggml_opt_result_init();
    ggml_opt_result_t result_val   = ggml_opt_result_init();

    ggml_opt_epoch_callback epoch_callback = is_silent ? NULL : ggml_opt_epoch_callback_progress_bar;

    opt_sync_cpu_threads(sched);

    // Allocate history arrays
    double * hist_train_loss = (double *)R_alloc(n_epoch, sizeof(double));
    double * hist_train_acc  = (double *)R_alloc(n_epoch, sizeof(double));
    double * hist_val_loss   = (double *)R_alloc(n_epoch, sizeof(double));
    double * hist_val_acc    = (double *)R_alloc(n_epoch, sizeof(double));

    for (; epoch <= n_epoch; ++epoch) {
        if (n_batch < idata_split) {
            ggml_opt_dataset_shuffle(opt_ctx, dataset, idata_split);
        }

        ggml_opt_result_reset(result_train);
        ggml_opt_result_reset(result_val);

        if (!is_silent) {
            Rprintf("Epoch %d/%d:\n", (int)epoch, (int)n_epoch);
        }

        ggml_opt_epoch(opt_ctx, dataset, result_train, result_val,
                       idata_split, epoch_callback, epoch_callback);

        if (!is_silent) {
            Rprintf("\n");
        }

        // Collect metrics
        int idx = (int)(epoch - 1);

        ggml_opt_result_loss(result_train, &hist_train_loss[idx], NULL);
        ggml_opt_result_accuracy(result_train, &hist_train_acc[idx], NULL);

        if (v_split > 0.0f) {
            ggml_opt_result_loss(result_val, &hist_val_loss[idx], NULL);
            ggml_opt_result_accuracy(result_val, &hist_val_acc[idx], NULL);
        } else {
            hist_val_loss[idx] = NA_REAL;
            hist_val_acc[idx]  = NA_REAL;
        }
    }

    ggml_opt_free(opt_ctx);
    ggml_opt_result_free(result_train);
    ggml_opt_result_free(result_val);

    // Build R list with history
    SEXP r_train_loss = PROTECT(allocVector(REALSXP, n_epoch));
    SEXP r_train_acc  = PROTECT(allocVector(REALSXP, n_epoch));
    SEXP r_val_loss   = PROTECT(allocVector(REALSXP, n_epoch));
    SEXP r_val_acc    = PROTECT(allocVector(REALSXP, n_epoch));

    memcpy(REAL(r_train_loss), hist_train_loss, n_epoch * sizeof(double));
    memcpy(REAL(r_train_acc),  hist_train_acc,  n_epoch * sizeof(double));
    memcpy(REAL(r_val_loss),   hist_val_loss,   n_epoch * sizeof(double));
    memcpy(REAL(r_val_acc),    hist_val_acc,    n_epoch * sizeof(double));

    SEXP result = PROTECT(allocVector(VECSXP, 4));
    SEXP names = PROTECT(allocVector(STRSXP, 4));

    SET_VECTOR_ELT(result, 0, r_train_loss);
    SET_VECTOR_ELT(result, 1, r_train_acc);
    SET_VECTOR_ELT(result, 2, r_val_loss);
    SET_VECTOR_ELT(result, 3, r_val_acc);

    SET_STRING_ELT(names, 0, mkChar("train_loss"));
    SET_STRING_ELT(names, 1, mkChar("train_accuracy"));
    SET_STRING_ELT(names, 2, mkChar("val_loss"));
    SET_STRING_ELT(names, 3, mkChar("val_accuracy"));

    setAttrib(result, R_NamesSymbol, names);
    UNPROTECT(6);
    return result;
}

// ============================================================================
// Additional Functions
// ============================================================================

// Get gradient accumulator for a node
SEXP R_ggml_opt_grad_acc(SEXP opt_ctx_ptr, SEXP node_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);
    struct ggml_tensor * node = (struct ggml_tensor *)R_ExternalPtrAddr(node_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }
    if (node == NULL) {
        error("Invalid node tensor pointer");
    }

    struct ggml_tensor * grad_acc = ggml_opt_grad_acc(opt_ctx, node);

    if (grad_acc == NULL) {
        return R_NilValue;
    }

    SEXP ptr = PROTECT(R_MakeExternalPtr(grad_acc, R_NilValue, R_NilValue));
    UNPROTECT(1);
    return ptr;
}

// Get predictions from result (returns integer vector of length ndata)
SEXP R_ggml_opt_result_pred(SEXP result_ptr) {
    ggml_opt_result_t result = (ggml_opt_result_t)R_ExternalPtrAddr(result_ptr);

    if (result == NULL) {
        error("Invalid result pointer");
    }

    // First get ndata
    int64_t ndata;
    ggml_opt_result_ndata(result, &ndata);

    if (ndata <= 0) {
        return allocVector(INTSXP, 0);
    }

    // Allocate R integer vector
    SEXP r_pred = PROTECT(allocVector(INTSXP, (R_xlen_t)ndata));

    // Get predictions
    ggml_opt_result_pred(result, INTEGER(r_pred));

    UNPROTECT(1);
    return r_pred;
}

// Prepare allocation for non-static graphs
SEXP R_ggml_opt_prepare_alloc(SEXP opt_ctx_ptr, SEXP ctx_compute_ptr,
                               SEXP graph_ptr, SEXP inputs_ptr, SEXP outputs_ptr) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);
    struct ggml_context * ctx_compute = (struct ggml_context *)R_ExternalPtrAddr(ctx_compute_ptr);
    struct ggml_cgraph * gf = (struct ggml_cgraph *)R_ExternalPtrAddr(graph_ptr);
    struct ggml_tensor * inputs = (struct ggml_tensor *)R_ExternalPtrAddr(inputs_ptr);
    struct ggml_tensor * outputs = (struct ggml_tensor *)R_ExternalPtrAddr(outputs_ptr);

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }
    if (ctx_compute == NULL) {
        error("Invalid compute context pointer");
    }
    if (gf == NULL) {
        error("Invalid graph pointer");
    }
    if (inputs == NULL) {
        error("Invalid inputs tensor pointer");
    }
    if (outputs == NULL) {
        error("Invalid outputs tensor pointer");
    }

    ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs);

    return R_NilValue;
}

// ============================================================================
// R Callback Support for ggml_opt_epoch
// ============================================================================

// Global storage for R callback functions (protected from GC)
// Initialized to NULL, will be set to R_NilValue when needed
static SEXP g_callback_train = NULL;
static SEXP g_callback_eval = NULL;

// C wrapper that calls R callback function
static void r_callback_wrapper(
        bool               train,
        ggml_opt_context_t opt_ctx,
        ggml_opt_dataset_t dataset,
        ggml_opt_result_t  result,
        int64_t            ibatch,
        int64_t            ibatch_max,
        int64_t            t_start_us) {

    // Select the appropriate R callback
    SEXP callback = train ? g_callback_train : g_callback_eval;

    if (callback == NULL || callback == R_NilValue || !Rf_isFunction(callback)) {
        return;
    }

    // Build the call: callback(train, ibatch, ibatch_max, t_start_us, result)
    // We pass simplified arguments that are easy to use in R
    SEXP call = PROTECT(Rf_allocVector(LANGSXP, 6));
    SETCAR(call, callback);

    SEXP args = CDR(call);
    SETCAR(args, Rf_ScalarLogical(train));

    args = CDR(args);
    SETCAR(args, Rf_ScalarReal((double)ibatch));

    args = CDR(args);
    SETCAR(args, Rf_ScalarReal((double)ibatch_max));

    args = CDR(args);
    SETCAR(args, Rf_ScalarReal((double)t_start_us));

    args = CDR(args);
    // Pass result as external pointer so user can query loss/accuracy
    SEXP result_ptr = PROTECT(R_MakeExternalPtr(result, R_NilValue, R_NilValue));
    SETCAR(args, result_ptr);

    // Evaluate the call in the global environment
    // Use R_tryEval to catch R errors without crashing
    int error_occurred = 0;
    R_tryEval(call, R_GlobalEnv, &error_occurred);

    if (error_occurred) {
        Rf_warning("Error in R callback function");
    }

    UNPROTECT(2);
}

// Run one epoch with R callback support
SEXP R_ggml_opt_epoch(SEXP opt_ctx_ptr, SEXP dataset_ptr,
                       SEXP result_train_ptr, SEXP result_eval_ptr,
                       SEXP idata_split, SEXP callback_train, SEXP callback_eval) {
    ggml_opt_context_t opt_ctx = (ggml_opt_context_t)R_ExternalPtrAddr(opt_ctx_ptr);
    ggml_opt_dataset_t dataset = (ggml_opt_dataset_t)R_ExternalPtrAddr(dataset_ptr);

    ggml_opt_result_t result_train = NULL;
    ggml_opt_result_t result_eval = NULL;

    if (opt_ctx == NULL) {
        error("Invalid optimizer context pointer");
    }
    if (dataset == NULL) {
        error("Invalid dataset pointer");
    }

    if (result_train_ptr != R_NilValue) {
        result_train = (ggml_opt_result_t)R_ExternalPtrAddr(result_train_ptr);
    }
    if (result_eval_ptr != R_NilValue) {
        result_eval = (ggml_opt_result_t)R_ExternalPtrAddr(result_eval_ptr);
    }

    int64_t split = (int64_t)asReal(idata_split);

    // Determine callbacks
    ggml_opt_epoch_callback cb_train_func = NULL;
    ggml_opt_epoch_callback cb_eval_func = NULL;

    // Store R callbacks in global variables (protected)
    // Check if callback_train is a function or special value
    if (Rf_isFunction(callback_train)) {
        R_PreserveObject(callback_train);
        g_callback_train = callback_train;
        cb_train_func = r_callback_wrapper;
    } else if (Rf_isLogical(callback_train) && asLogical(callback_train)) {
        // TRUE means use built-in progress bar
        cb_train_func = ggml_opt_epoch_callback_progress_bar;
    } else {
        g_callback_train = NULL;
    }

    if (Rf_isFunction(callback_eval)) {
        R_PreserveObject(callback_eval);
        g_callback_eval = callback_eval;
        cb_eval_func = r_callback_wrapper;
    } else if (Rf_isLogical(callback_eval) && asLogical(callback_eval)) {
        // TRUE means use built-in progress bar
        cb_eval_func = ggml_opt_epoch_callback_progress_bar;
    } else {
        g_callback_eval = NULL;
    }

    // Run the epoch
    ggml_opt_epoch(opt_ctx, dataset, result_train, result_eval, split, cb_train_func, cb_eval_func);

    // Release R callbacks
    if (g_callback_train != NULL && Rf_isFunction(g_callback_train)) {
        R_ReleaseObject(g_callback_train);
        g_callback_train = NULL;
    }
    if (g_callback_eval != NULL && Rf_isFunction(g_callback_eval)) {
        R_ReleaseObject(g_callback_eval);
        g_callback_eval = NULL;
    }

    return R_NilValue;
}

// ============================================================================
// LR-controllable optimizer context for R-side epoch loop
// ============================================================================

// Userdata for R-controlled LR: holds current optimizer params
// (updated by R between epochs via R_ggml_opt_set_lr)
typedef struct {
    struct ggml_opt_optimizer_params params;
} r_opt_lr_userdata;

// C callback: simply returns the stored params (R updates them between epochs)
static struct ggml_opt_optimizer_params r_opt_get_constant_lr(void * userdata) {
    r_opt_lr_userdata * ud = (r_opt_lr_userdata *)userdata;
    return ud->params;
}

// Finalizer: free userdata when external pointer is GC'd
static void r_opt_lr_userdata_finalizer(SEXP ptr) {
    void * ud = R_ExternalPtrAddr(ptr);
    if (ud != NULL) {
        free(ud);
        R_ClearExternalPtr(ptr);
    }
}

// Initialize optimizer context for R-side epoch loop.
// Returns a list: list(opt_ctx=<ptr>, lr_ud=<ptr>) so R can call R_ggml_opt_set_lr.
// Default LR is taken from ggml_opt_get_default_optimizer_params.
SEXP R_ggml_opt_init_for_fit(SEXP sched_ptr, SEXP loss_type, SEXP optimizer_type,
                              SEXP opt_period, SEXP ctx_compute_ptr,
                              SEXP inputs_ptr, SEXP outputs_ptr) {
    ggml_backend_sched_t sched = (ggml_backend_sched_t)R_ExternalPtrAddr(sched_ptr);
    if (sched == NULL) error("Invalid scheduler pointer");

    enum ggml_opt_loss_type lt = (enum ggml_opt_loss_type)asInteger(loss_type);
    struct ggml_opt_params params = ggml_opt_default_params(sched, lt);
    params.optimizer = (enum ggml_opt_optimizer_type)asInteger(optimizer_type);
    params.opt_period = asInteger(opt_period);

    if (ctx_compute_ptr != R_NilValue)
        params.ctx_compute = (struct ggml_context *)R_ExternalPtrAddr(ctx_compute_ptr);
    if (inputs_ptr != R_NilValue)
        params.inputs = (struct ggml_tensor *)R_ExternalPtrAddr(inputs_ptr);
    if (outputs_ptr != R_NilValue)
        params.outputs = (struct ggml_tensor *)R_ExternalPtrAddr(outputs_ptr);

    // Allocate userdata with default LR from ggml defaults
    r_opt_lr_userdata * ud = (r_opt_lr_userdata *)malloc(sizeof(r_opt_lr_userdata));
    if (ud == NULL) error("Failed to allocate LR userdata");
    // Get default params (epoch=1 just to get defaults)
    int64_t dummy_epoch = 1;
    ud->params = ggml_opt_get_default_optimizer_params(&dummy_epoch);

    params.get_opt_pars    = r_opt_get_constant_lr;
    params.get_opt_pars_ud = ud;

    ggml_opt_context_t opt_ctx = ggml_opt_init(params);
    if (opt_ctx == NULL) { free(ud); error("Failed to initialize optimizer context"); }

    SEXP opt_ptr = PROTECT(R_MakeExternalPtr(opt_ctx, R_NilValue, R_NilValue));
    SEXP ud_ptr  = PROTECT(R_MakeExternalPtr(ud, R_NilValue, R_NilValue));
    R_RegisterCFinalizerEx(ud_ptr, r_opt_lr_userdata_finalizer, TRUE);

    SEXP result = PROTECT(allocVector(VECSXP, 2));
    SEXP names  = PROTECT(allocVector(STRSXP, 2));
    SET_VECTOR_ELT(result, 0, opt_ptr);
    SET_VECTOR_ELT(result, 1, ud_ptr);
    SET_STRING_ELT(names, 0, mkChar("opt_ctx"));
    SET_STRING_ELT(names, 1, mkChar("lr_ud"));
    setAttrib(result, R_NamesSymbol, names);

    UNPROTECT(4);
    return result;
}

// Update learning rate in the userdata (called between epochs from R).
// adamw_lr: new AdamW LR (NA to keep current)
// sgd_lr:   new SGD LR (NA to keep current)
SEXP R_ggml_opt_set_lr(SEXP ud_ptr, SEXP adamw_lr, SEXP sgd_lr) {
    r_opt_lr_userdata * ud = (r_opt_lr_userdata *)R_ExternalPtrAddr(ud_ptr);
    if (ud == NULL) error("Invalid LR userdata pointer");

    if (!ISNA(asReal(adamw_lr)))
        ud->params.adamw.alpha = (float)asReal(adamw_lr);
    if (!ISNA(asReal(sgd_lr)))
        ud->params.sgd.alpha = (float)asReal(sgd_lr);

    return R_NilValue;
}

// Get current LR from userdata
SEXP R_ggml_opt_get_lr(SEXP ud_ptr) {
    r_opt_lr_userdata * ud = (r_opt_lr_userdata *)R_ExternalPtrAddr(ud_ptr);
    if (ud == NULL) error("Invalid LR userdata pointer");

    SEXP result = PROTECT(allocVector(REALSXP, 2));
    SEXP names  = PROTECT(allocVector(STRSXP, 2));
    REAL(result)[0] = (double)ud->params.adamw.alpha;
    REAL(result)[1] = (double)ud->params.sgd.alpha;
    SET_STRING_ELT(names, 0, mkChar("adamw"));
    SET_STRING_ELT(names, 1, mkChar("sgd"));
    setAttrib(result, R_NamesSymbol, names);
    UNPROTECT(2);
    return result;
}
