// Copyright (c) 2024-2026, NVIDIA Corporation. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
//////////////////////////////////////////////////////////////////////////////
/**
 * @file types.cpp
 * @brief Implementation of type conversion and data structure management functions
 * 
 * Provides string-to-enum conversions, data structure allocation/deallocation,
 * and serialization utilities for gradient boosting types.
 */

#include <stdexcept>
#include <string>
#include <cstring>
#include <iostream>
#include <fstream>
#include <stdexcept>

#include "types.h"
#include "utils.h"
#include "optimizer.h"
#ifdef USE_CUDA
#include "cuda_types.h"
#endif

scoreFunc stringToScoreFunc(std::string str) {
    if (str == "L2" || str == "l2") return scoreFunc::L2;
    if (str == "Cosine" || str == "cosine") return scoreFunc::Cosine;
    throw std::runtime_error("Invalid score function! Options are: Cosine/L2");
    return scoreFunc::L2;
}

generatorType stringTogeneratorType(std::string str) {
    if (str == "uniform" || str == "Uniform") return generatorType::Uniform;
    if (str == "quantile" || str == "Quantile") return generatorType::Quantile;
    throw std::runtime_error("Invalid generator function! Options are: Uniform/Quantile");
    return generatorType::Uniform;
}

growPolicy stringTogrowPolicy(std::string str) {
    if (str == "oblivious" || str == "Oblivious") return growPolicy::OBLIVIOUS;
    if (str == "greedy" || str == "Greedy") return growPolicy::GREEDY;
    throw std::runtime_error("Invalid generator function! Options are: Greedy/Oblivious");
    return growPolicy::GREEDY;
}

lossType stringTolossType(std::string str) {
    if (str == "MultiRMSE") return lossType::MultiRMSE;
    throw std::runtime_error("Invalid loss function! Options are: MultiRMSE");
    return lossType::MultiRMSE;
}

deviceType stringTodeviceType(std::string str) {
    if (str == "cpu") return deviceType::cpu;
    if (str == "cuda" || str == "gpu") return deviceType::gpu;
    throw std::runtime_error("Invalid device! Options are: cpu/cuda");
}

exportFormat stringToexportFormat(std::string str) {
    if (str == "fxp8") return exportFormat::EXP_FXP8;
    if (str == "fxp16") return exportFormat::EXP_FXP16;
    if (str == "float") return exportFormat::EXP_FLOAT;
    throw std::runtime_error("Invalid exportFormat! Options are: float/fxp8/fxp16");
}
exportType stringToexportType(std::string str) {
    if (str == "compact") return exportType::COMPACT;
    if (str == "full") return exportType::FULL;
    throw std::runtime_error("Invalid exportType Options are: full/compact");
}

optimizerAlgo stringToAlgoType(std::string str) {
    if (str == "Adam" || str == "adam") return optimizerAlgo::Adam;
    if (str == "SGD" || str == "sgd") return optimizerAlgo::SGD;
    throw std::runtime_error("Invalid Optimizer Algorithm! Options are: SGD/Adam");
    return optimizerAlgo::SGD;
}

schedulerFunc stringToSchedulerType(std::string str) {
    if (str == "Const" || str == "const") return schedulerFunc::Const;
    if (str == "Linear" || str == "linear") return schedulerFunc::Linear;
    throw std::runtime_error("Invalid Scheduler type! Options are: Linear/Const");
    return schedulerFunc::Const;
}

std::string scoreFuncToString(scoreFunc func) {
    switch (func) {
        case scoreFunc::L2:
            return "L2";
        case scoreFunc::Cosine:
            return "Cosine";
        default:
            throw std::runtime_error("Invalid score function.");
    }
}

std::string generatorTypeToString(generatorType type) {
    switch (type) {
        case generatorType::Uniform:
            return "Uniform";
        case generatorType::Quantile:
            return "Quantile";
        default:
            throw std::runtime_error("Invalid generator type.");
    }
}

std::string growPolicyToString(growPolicy type) {
    switch (type) {
        case growPolicy::OBLIVIOUS:
            return "Oblivous";
        case growPolicy::GREEDY:
            return "Greedy";
        default:
            throw std::runtime_error("Invalid generator type.");
    }
}


std::string lossTypeToString(lossType type) {
    switch (type) {
        case lossType::MultiRMSE:
            return "MultiRMSE";
        default:
            throw std::runtime_error("Invalid loss type.");
    }
}

std::string deviceTypeToString(deviceType type) {
    switch (type) {
        case deviceType::cpu:
            return "cpu";
        case deviceType::gpu:
            return "cuda"; // Assuming 'cuda' is the preferred string for GPU.
        default:
            throw std::runtime_error("Invalid device type.");
    }
}

std::string algoTypeToString(optimizerAlgo algo) {
    switch (algo) {
        case optimizerAlgo::Adam:
            return "Adam";
        case optimizerAlgo::SGD:
            return "SGD";
        default:
            throw std::runtime_error("Invalid optimizer algorithm.");
    }
}

std::string schedulerTypeToString(schedulerFunc func) {
    switch (func) {
        case schedulerFunc::Const:
            return "Const";
        case schedulerFunc::Linear:
            return "Linear";
        default:
            throw std::runtime_error("Invalid scheduler type.");
    }
}

ensembleMetaData* ensemble_metadata_alloc(
    int max_trees,
    int max_leaves,
    int max_trees_batch,
    int max_leaves_batch,
    int input_dim,
    int output_dim,
    int policy_dim,
    int max_depth,
    int min_data_in_leaf,
    int n_bins,
    int par_th,
    float cv_beta,
    int verbose,
    int batch_size,
    bool use_cv,
    scoreFunc split_score_func,
    generatorType generator_type,
    growPolicy grow_policy,
    int n_mono_constraints){

    ensembleMetaData *metadata = new ensembleMetaData;
    metadata->input_dim = input_dim; 
    metadata->output_dim = output_dim; 
    metadata->policy_dim = policy_dim;
    metadata->max_depth = max_depth; 
    metadata->min_data_in_leaf = min_data_in_leaf; 
    metadata->par_th = par_th; 
    metadata->n_bins = n_bins; 
    metadata->n_leaves = 0;
    metadata->n_trees = 0; 
    metadata->verbose = verbose; 
    metadata->batch_size = batch_size; 
    metadata->grow_policy = grow_policy;
    metadata->generator_type = generator_type;
    metadata->cv_beta = cv_beta;
    metadata->split_score_func = split_score_func;
    metadata->use_cv = use_cv;
    metadata->max_trees = max_trees;
    metadata->max_leaves = max_leaves; 
    metadata->max_trees_batch = max_trees_batch; 
    metadata->max_leaves_batch = max_leaves_batch;
    metadata->n_num_features = 0;
    metadata->n_cat_features = 0;
    metadata->iteration = 0;
    metadata->n_mono_constraints = n_mono_constraints;
    return metadata;
}

ensembleData* ensemble_data_alloc(ensembleMetaData *metadata){
    if (metadata == nullptr){
        std::cerr << "Error metadata is nullptr cannot allocate ensembleData." << std::endl;
        throw std::runtime_error("Error invalid pointer");
    }
    ensembleData *edata = new ensembleData;
    size_t data_size = 0;
    
    // Allocate sub-structs on CPU
    edata->ensemble_info = new ensembleInfo;
    edata->leaf_data = new leafData;
    edata->feature_data = new featureData;
    edata->mono_constraints = new monotonicConstraints;
    edata->mono_constraints->n_constraints = 0;  // Initialize to 0
    edata->feature_mappings = new featureMapping;
    
    edata->bias = new float[metadata->output_dim];
    data_size += sizeof(float) * metadata->output_dim;
    memset(edata->bias, 0, metadata->output_dim * sizeof(float));
    int split_sizes = (metadata->grow_policy == OBLIVIOUS) ? metadata->max_trees : metadata->max_leaves;
#ifdef DEBUG
    edata->n_samples = new int[metadata->max_leaves]; // debugging
    memset(edata->n_samples, 0, metadata->max_leaves * sizeof(int));
    data_size += sizeof(int) * metadata->max_leaves;
#endif
    // Ensemble info
    edata->ensemble_info->tree_indices = new int[metadata->max_trees];
    data_size += sizeof(int) * metadata->max_trees;
    memset(edata->ensemble_info->tree_indices, 0, metadata->max_trees * sizeof(int));
    edata->ensemble_info->depths = new int[split_sizes];
    data_size += sizeof(int) * split_sizes;
    memset(edata->ensemble_info->depths, 0, split_sizes * sizeof(int));
    
    // Leaf data
    edata->leaf_data->values = new float[metadata->max_leaves * metadata->output_dim];
    data_size += sizeof(float) * metadata->max_leaves * metadata->output_dim;
    memset(edata->leaf_data->values, 0, metadata->max_leaves * metadata->output_dim * sizeof(float));
    edata->leaf_data->edge_weights = new float[metadata->max_leaves * metadata->max_depth];
    data_size += sizeof(float) * metadata->max_leaves * metadata->max_depth;
    memset(edata->leaf_data->edge_weights, 0, metadata->max_leaves * metadata->max_depth * sizeof(float));
    
    // Feature data
    edata->feature_data->feature_indices = new int[split_sizes * metadata->max_depth];
    data_size += sizeof(int) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->feature_indices, 0, split_sizes * metadata->max_depth * sizeof(int));
    edata->feature_data->feature_values = new float[split_sizes * metadata->max_depth];
    data_size += sizeof(float) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->feature_values, 0, split_sizes * metadata->max_depth * sizeof(float));
    edata->feature_data->feature_weights = new float[metadata->input_dim];
    data_size += sizeof(float) * metadata->input_dim;
    memset(edata->feature_data->feature_weights, 0, metadata->input_dim * sizeof(float));
    edata->feature_data->is_numerics = new bool[split_sizes * metadata->max_depth];
    data_size += sizeof(bool) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->is_numerics, 0, split_sizes * metadata->max_depth * sizeof(bool));
    edata->feature_data->inequality_directions = new bool[metadata->max_leaves * metadata->max_depth]; 
    data_size += sizeof(bool) * metadata->max_leaves * metadata->max_depth;
    memset(edata->feature_data->inequality_directions, 0, metadata->max_leaves * metadata->max_depth * sizeof(bool));
    edata->feature_data->categorical_values = new char[split_sizes * metadata->max_depth * MAX_CHAR_SIZE];
    data_size += sizeof(char) * split_sizes * metadata->max_depth * MAX_CHAR_SIZE;
    memset(edata->feature_data->categorical_values, 0, split_sizes * metadata->max_depth * sizeof(char) * MAX_CHAR_SIZE);
    
    // Monotonic constraints
    edata->mono_constraints->feature_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memset(edata->mono_constraints->feature_idx, 0, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->output_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memset(edata->mono_constraints->output_idx, 0, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->constraint = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memset(edata->mono_constraints->constraint, 0, metadata->n_mono_constraints * sizeof(int));

    // Feature mappings
    edata->feature_mappings->reverse_num_feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memset(edata->feature_mappings->reverse_num_feature_mapping, 0, metadata->input_dim * sizeof(int));
    edata->feature_mappings->reverse_cat_feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memset(edata->feature_mappings->reverse_cat_feature_mapping, 0, metadata->input_dim * sizeof(int));
    edata->feature_mappings->feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memset(edata->feature_mappings->feature_mapping, 0, metadata->input_dim * sizeof(int));
    edata->feature_mappings->mapping_numerics = new bool[metadata->input_dim];
    data_size += sizeof(bool) * metadata->input_dim;
    memset(edata->feature_mappings->mapping_numerics, 0, metadata->input_dim * sizeof(bool));
    
    edata->alloc_data_size = data_size;
    return edata;
}

ensembleData* ensemble_copy_data_alloc(ensembleMetaData *metadata){
    // same as normal alloc but only allocate memory for existing size 
    if (metadata == nullptr){
        std::cerr << "Error metadata is nullptr cannot allocate ensembleData." << std::endl;
        throw std::runtime_error("Error invalid pointer");
    }
    ensembleData *edata = new ensembleData;
    size_t data_size = 0;
    
    // Allocate sub-structs on CPU
    edata->ensemble_info = new ensembleInfo;
    edata->leaf_data = new leafData;
    edata->feature_data = new featureData;
    edata->mono_constraints = new monotonicConstraints;
    edata->mono_constraints->n_constraints = 0;  // Initialize to 0
    edata->feature_mappings = new featureMapping;
    
    edata->bias = new float[metadata->output_dim];
    data_size += sizeof(float) * metadata->output_dim;
    memset(edata->bias, 0, metadata->output_dim * sizeof(float));
    int split_sizes = (metadata->grow_policy == OBLIVIOUS) ? metadata->n_trees : metadata->n_leaves;
#ifdef DEBUG
    edata->n_samples = new int[metadata->n_leaves]; // debugging
    memset(edata->n_samples, 0, metadata->n_leaves * sizeof(int));
    data_size += sizeof(int) * metadata->n_leaves;
#endif
    // Ensemble info
    edata->ensemble_info->tree_indices = new int[metadata->n_trees];
    data_size += sizeof(int) * metadata->n_trees;
    memset(edata->ensemble_info->tree_indices, 0, metadata->n_trees * sizeof(int));
    edata->ensemble_info->depths = new int[split_sizes];
    data_size += sizeof(int) * split_sizes;
    memset(edata->ensemble_info->depths, 0, split_sizes * sizeof(int));
    
    // Leaf data
    edata->leaf_data->values = new float[metadata->n_leaves * metadata->output_dim];
    data_size += sizeof(float) * metadata->n_leaves * metadata->output_dim;
    memset(edata->leaf_data->values, 0, metadata->n_leaves * metadata->output_dim * sizeof(float));
    edata->leaf_data->edge_weights = new float[metadata->n_leaves * metadata->max_depth];
    data_size += sizeof(float) * metadata->n_leaves * metadata->max_depth;
    memset(edata->leaf_data->edge_weights, 0, metadata->n_leaves * metadata->max_depth * sizeof(float));
    
    // Feature data
    edata->feature_data->feature_indices = new int[split_sizes * metadata->max_depth];
    data_size += sizeof(int) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->feature_indices, 0, split_sizes * metadata->max_depth * sizeof(int));
    edata->feature_data->feature_values = new float[split_sizes * metadata->max_depth];
    data_size += sizeof(float) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->feature_values, 0, split_sizes * metadata->max_depth * sizeof(float));
    edata->feature_data->feature_weights = new float[metadata->input_dim];
    data_size += sizeof(float) * metadata->input_dim;
    memset(edata->feature_data->feature_weights, 0, metadata->input_dim * sizeof(float));
    edata->feature_data->is_numerics = new bool[split_sizes * metadata->max_depth];
    data_size += sizeof(bool) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->is_numerics, 0, split_sizes * metadata->max_depth * sizeof(bool));
    edata->feature_data->inequality_directions = new bool[metadata->n_leaves * metadata->max_depth]; 
    data_size += sizeof(bool) * metadata->n_leaves * metadata->max_depth;
    memset(edata->feature_data->inequality_directions, 0, metadata->n_leaves * metadata->max_depth * sizeof(bool));
    edata->feature_data->categorical_values = new char[split_sizes * metadata->max_depth * MAX_CHAR_SIZE];
    data_size += sizeof(char) * split_sizes * metadata->max_depth * MAX_CHAR_SIZE;
    memset(edata->feature_data->categorical_values, 0, split_sizes * metadata->max_depth * sizeof(char) * MAX_CHAR_SIZE);
    
    // Monotonic constraints
    edata->mono_constraints->feature_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memset(edata->mono_constraints->feature_idx, 0, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->output_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memset(edata->mono_constraints->output_idx, 0, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->constraint = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memset(edata->mono_constraints->constraint, 0, metadata->n_mono_constraints * sizeof(int));

    // Feature mappings
    edata->feature_mappings->reverse_num_feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memset(edata->feature_mappings->reverse_num_feature_mapping, 0, metadata->input_dim * sizeof(int));
    edata->feature_mappings->reverse_cat_feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memset(edata->feature_mappings->reverse_cat_feature_mapping, 0, metadata->input_dim * sizeof(int));
    edata->feature_mappings->feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memset(edata->feature_mappings->feature_mapping, 0, metadata->input_dim * sizeof(int));
    edata->feature_mappings->mapping_numerics = new bool[metadata->input_dim];
    data_size += sizeof(bool) * metadata->input_dim;
    memset(edata->feature_mappings->mapping_numerics, 0, metadata->input_dim * sizeof(bool));
    
    edata->alloc_data_size = data_size;
    return edata;
}

ensembleData* copy_ensemble_data(ensembleData *other_edata, ensembleMetaData *metadata){
    if (metadata == nullptr || other_edata == nullptr){
        std::cerr << "Error metadata is nullptr cannot allocate ensembleData." << std::endl;
        throw std::runtime_error("Error invalid pointer");
    }
    ensembleData *edata = new ensembleData;
    size_t data_size = 0;
    
    // Allocate sub-structs on CPU
    edata->ensemble_info = new ensembleInfo;
    edata->leaf_data = new leafData;
    edata->feature_data = new featureData;
    edata->mono_constraints = new monotonicConstraints;
    edata->mono_constraints->n_constraints = 0;  // Initialize to 0
    edata->feature_mappings = new featureMapping;
    
    edata->bias = new float[metadata->output_dim];
    data_size += sizeof(float) * metadata->output_dim;
    memcpy(edata->bias, other_edata->bias, metadata->output_dim * sizeof(float));
    int split_sizes = (metadata->grow_policy == OBLIVIOUS) ? metadata->n_trees : metadata->n_leaves;
#ifdef DEBUG
    edata->n_samples = new int[metadata->n_leaves]; // debugging
    memcpy(edata->n_samples, other_edata->n_samples, metadata->n_leaves * sizeof(int));
    data_size += sizeof(float) * metadata->n_leaves;
#endif
    // Ensemble info
    edata->ensemble_info->tree_indices = new int[metadata->n_trees];
    data_size += sizeof(int) * metadata->n_trees;
    memcpy(edata->ensemble_info->tree_indices, other_edata->ensemble_info->tree_indices, metadata->n_trees * sizeof(int));
    edata->ensemble_info->depths = new int[split_sizes];
    data_size += sizeof(int) * split_sizes;
    memcpy(edata->ensemble_info->depths, other_edata->ensemble_info->depths, split_sizes * sizeof(int));
    
    // Leaf data
    edata->leaf_data->values = new float[metadata->n_leaves * metadata->output_dim];
    data_size += sizeof(float) * metadata->n_leaves * metadata->output_dim;
    memcpy(edata->leaf_data->values, other_edata->leaf_data->values, metadata->n_leaves * metadata->output_dim * sizeof(float));
    edata->leaf_data->edge_weights = new float[metadata->n_leaves * metadata->max_depth];
    data_size += sizeof(float) * metadata->n_leaves * metadata->max_depth;
    memcpy(edata->leaf_data->edge_weights, other_edata->leaf_data->edge_weights, metadata->n_leaves * metadata->max_depth * sizeof(float));
    
    // Feature data
    edata->feature_data->feature_indices = new int[split_sizes * metadata->max_depth];
    data_size += sizeof(int) * split_sizes * metadata->max_depth;
    memcpy(edata->feature_data->feature_indices, other_edata->feature_data->feature_indices, split_sizes * metadata->max_depth * sizeof(int));
    edata->feature_data->feature_values = new float[split_sizes * metadata->max_depth];
    data_size += sizeof(float) * split_sizes * metadata->max_depth;
    memcpy(edata->feature_data->feature_values, other_edata->feature_data->feature_values, split_sizes * metadata->max_depth * sizeof(float));
    edata->feature_data->feature_weights = new float[metadata->input_dim];
    data_size += sizeof(float) * metadata->input_dim;
    memcpy(edata->feature_data->feature_weights, other_edata->feature_data->feature_weights, metadata->input_dim * sizeof(float));
    edata->feature_data->is_numerics = new bool[split_sizes * metadata->max_depth];
    data_size += sizeof(bool) * split_sizes * metadata->max_depth;
    memcpy(edata->feature_data->is_numerics, other_edata->feature_data->is_numerics, split_sizes * metadata->max_depth * sizeof(bool));
    edata->feature_data->categorical_values = new char[split_sizes * metadata->max_depth * MAX_CHAR_SIZE];
    data_size += sizeof(char) * split_sizes * metadata->max_depth * MAX_CHAR_SIZE;
    memcpy(edata->feature_data->categorical_values, other_edata->feature_data->categorical_values, split_sizes * metadata->max_depth * sizeof(char) * MAX_CHAR_SIZE);
    edata->feature_data->inequality_directions = new bool[metadata->n_leaves * metadata->max_depth]; 
    data_size += sizeof(bool) * metadata->n_leaves * metadata->max_depth;
    memcpy(edata->feature_data->inequality_directions, other_edata->feature_data->inequality_directions, metadata->n_leaves * metadata->max_depth * sizeof(bool));
    
    // Monotonic constraints
    edata->mono_constraints->feature_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memcpy(edata->mono_constraints->feature_idx, other_edata->mono_constraints->feature_idx, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->output_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memcpy(edata->mono_constraints->output_idx, other_edata->mono_constraints->output_idx, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->constraint = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memcpy(edata->mono_constraints->constraint, other_edata->mono_constraints->constraint, metadata->n_mono_constraints * sizeof(int));

    // Feature mappings
    edata->feature_mappings->reverse_num_feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memcpy(edata->feature_mappings->reverse_num_feature_mapping, other_edata->feature_mappings->reverse_num_feature_mapping, metadata->input_dim * sizeof(int));
    edata->feature_mappings->reverse_cat_feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memcpy(edata->feature_mappings->reverse_cat_feature_mapping, other_edata->feature_mappings->reverse_cat_feature_mapping, metadata->input_dim * sizeof(int));
    edata->feature_mappings->feature_mapping = new int[metadata->input_dim];
    data_size += sizeof(int) * metadata->input_dim;
    memcpy(edata->feature_mappings->feature_mapping, other_edata->feature_mappings->feature_mapping, metadata->input_dim * sizeof(int));
    edata->feature_mappings->mapping_numerics = new bool[metadata->input_dim];
    data_size += sizeof(bool) * metadata->input_dim;
    memcpy(edata->feature_mappings->mapping_numerics, other_edata->feature_mappings->mapping_numerics, metadata->input_dim * sizeof(bool));
    
    metadata->max_trees = metadata->n_trees;
    metadata->max_leaves = metadata->n_leaves;
    edata->alloc_data_size = data_size;
    return edata;
}

ensembleData* copy_compressed_ensemble_data(ensembleData *other_edata, ensembleMetaData *metadata, const int *leaf_indices, const int *tree_indices, const int n_compressed_leaves, const int n_compressed_trees, const int *new_tree_indices){
    
    if (metadata == nullptr || other_edata == nullptr){
        std::cerr << "Error metadata is nullptr cannot allocate ensembleData." << std::endl;
        throw std::runtime_error("Error invalid pointer");
    }
    ensembleData *edata = new ensembleData;
    size_t data_size = 0;
    
    // Allocate sub-structs on CPU
    edata->ensemble_info = new ensembleInfo;
    edata->leaf_data = new leafData;
    edata->feature_data = new featureData;
    edata->mono_constraints = new monotonicConstraints;
    edata->mono_constraints->n_constraints = 0;  // Initialize to 0
    edata->feature_mappings = new featureMapping;
    
    edata->bias = new float[metadata->output_dim];
    data_size += sizeof(float) * metadata->output_dim;
    memcpy(edata->bias, other_edata->bias, metadata->output_dim * sizeof(float));
    
    // Feature mappings
    edata->feature_mappings->feature_mapping = new int[metadata->input_dim];
    memcpy(edata->feature_mappings->feature_mapping, other_edata->feature_mappings->feature_mapping, metadata->input_dim * sizeof(int));
    data_size += sizeof(int) * metadata->input_dim;
    edata->feature_mappings->reverse_num_feature_mapping = new int[metadata->input_dim];
    memcpy(edata->feature_mappings->reverse_num_feature_mapping, other_edata->feature_mappings->reverse_num_feature_mapping, metadata->input_dim * sizeof(int));
    data_size += sizeof(int) * metadata->input_dim;
    edata->feature_mappings->reverse_cat_feature_mapping = new int[metadata->input_dim];
    memcpy(edata->feature_mappings->reverse_cat_feature_mapping, other_edata->feature_mappings->reverse_cat_feature_mapping, metadata->input_dim * sizeof(int));
    data_size += sizeof(int) * metadata->input_dim;
    edata->feature_mappings->mapping_numerics = new bool[metadata->input_dim];
    memcpy(edata->feature_mappings->mapping_numerics, other_edata->feature_mappings->mapping_numerics, metadata->input_dim * sizeof(bool));
    data_size += sizeof(bool) * metadata->input_dim;
    
    // Feature data - feature_weights
    edata->feature_data->feature_weights = new float[metadata->input_dim];
    data_size += sizeof(float) * metadata->input_dim;
    memcpy(edata->feature_data->feature_weights, other_edata->feature_data->feature_weights, metadata->input_dim * sizeof(float));
    
    int split_sizes = (metadata->grow_policy == OBLIVIOUS) ? n_compressed_trees : n_compressed_leaves;
    const int *split_indices = (metadata->grow_policy == OBLIVIOUS) ? tree_indices : leaf_indices;
#ifdef DEBUG
    edata->n_samples = new int[n_compressed_leaves]; // debugging
    data_size += sizeof(int) * n_compressed_leaves;
    memset(edata->n_samples, 0, n_compressed_leaves * sizeof(int));
    selective_copy(n_compressed_leaves, leaf_indices, edata->n_samples, other_edata->n_samples, 1);
#endif
    // Ensemble info
    edata->ensemble_info->tree_indices = new int[n_compressed_trees];
    data_size += sizeof(int) * n_compressed_trees;
    memcpy(edata->ensemble_info->tree_indices, new_tree_indices, n_compressed_trees * sizeof(int));
    edata->ensemble_info->depths = new int[split_sizes];
    data_size += sizeof(int) * split_sizes;
    memset(edata->ensemble_info->depths, 0, split_sizes * sizeof(int));
    selective_copy(split_sizes, split_indices, edata->ensemble_info->depths, other_edata->ensemble_info->depths, 1);
    
    // Leaf data
    edata->leaf_data->values = new float[n_compressed_leaves * metadata->output_dim];
    data_size += sizeof(float) * n_compressed_leaves * metadata->output_dim;
    memset(edata->leaf_data->values, 0, n_compressed_leaves*metadata->output_dim * sizeof(float));
    selective_copy(n_compressed_leaves, leaf_indices, edata->leaf_data->values, other_edata->leaf_data->values, metadata->output_dim);
    edata->leaf_data->edge_weights = new float[n_compressed_leaves * metadata->max_depth];
    data_size += sizeof(float) * n_compressed_leaves * metadata->max_depth;
    memset(edata->leaf_data->edge_weights, 0, n_compressed_leaves * metadata->max_depth * sizeof(float));
    selective_copy(n_compressed_leaves, leaf_indices, edata->leaf_data->edge_weights, other_edata->leaf_data->edge_weights, metadata->max_depth);
    
    // Feature data
    edata->feature_data->feature_indices = new int[split_sizes * metadata->max_depth];
    data_size += sizeof(int) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->feature_indices, 0, split_sizes*metadata->max_depth * sizeof(int));
    selective_copy(split_sizes, split_indices, edata->feature_data->feature_indices, other_edata->feature_data->feature_indices, metadata->max_depth);
    edata->feature_data->feature_values = new float[split_sizes * metadata->max_depth];
    data_size += sizeof(float) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->feature_values, 0, split_sizes * metadata->max_depth * sizeof(float));
    selective_copy(split_sizes, split_indices, edata->feature_data->feature_values, other_edata->feature_data->feature_values, metadata->max_depth);
    edata->feature_data->is_numerics = new bool[split_sizes * metadata->max_depth];
    data_size += sizeof(bool) * split_sizes * metadata->max_depth;
    memset(edata->feature_data->is_numerics, 0, split_sizes * metadata->max_depth * sizeof(bool));
    selective_copy(split_sizes, split_indices, edata->feature_data->is_numerics, other_edata->feature_data->is_numerics, metadata->max_depth);
    edata->feature_data->categorical_values = new char[split_sizes * metadata->max_depth * MAX_CHAR_SIZE];
    data_size += sizeof(char) * split_sizes * metadata->max_depth * MAX_CHAR_SIZE;
    memset(edata->feature_data->categorical_values, 0, split_sizes * metadata->max_depth * sizeof(char) * MAX_CHAR_SIZE);
    selective_copy_char(split_sizes, split_indices, edata->feature_data->categorical_values, other_edata->feature_data->categorical_values, metadata->max_depth);
    edata->feature_data->inequality_directions = new bool[n_compressed_leaves * metadata->max_depth]; 
    data_size += sizeof(bool) * n_compressed_leaves * metadata->max_depth;
    memset(edata->feature_data->inequality_directions, 0, n_compressed_leaves * metadata->max_depth * sizeof(bool));
    selective_copy(n_compressed_leaves, leaf_indices, edata->feature_data->inequality_directions, other_edata->feature_data->inequality_directions, metadata->max_depth);
    
    // Monotonic constraints - allocate but keep empty for compressed
    edata->mono_constraints->feature_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memcpy(edata->mono_constraints->feature_idx, other_edata->mono_constraints->feature_idx, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->output_idx = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memcpy(edata->mono_constraints->output_idx, other_edata->mono_constraints->output_idx, metadata->n_mono_constraints * sizeof(int));
    edata->mono_constraints->constraint = new int[metadata->n_mono_constraints];
    data_size += sizeof(int) * metadata->n_mono_constraints;
    memcpy(edata->mono_constraints->constraint, other_edata->mono_constraints->constraint, metadata->n_mono_constraints * sizeof(int));
    
    metadata->max_trees = n_compressed_trees;
    metadata->max_leaves = n_compressed_leaves;
    metadata->n_leaves = n_compressed_leaves;
    metadata->n_trees = n_compressed_trees;
    edata->alloc_data_size = data_size;
    return edata;
}

void ensemble_data_dealloc(ensembleData *edata){
    delete[] edata->bias;
#ifdef DEBUG
    delete[] edata->n_samples;
#endif
    // Ensemble info
    delete[] edata->ensemble_info->tree_indices;
    delete[] edata->ensemble_info->depths;
    delete edata->ensemble_info;
    
    // Leaf data
    delete[] edata->leaf_data->values;
    delete[] edata->leaf_data->edge_weights;
    delete edata->leaf_data;
    
    // Feature data
    delete[] edata->feature_data->feature_indices;
    delete[] edata->feature_data->feature_values;
    delete[] edata->feature_data->feature_weights;
    delete[] edata->feature_data->is_numerics;
    delete[] edata->feature_data->categorical_values;
    delete[] edata->feature_data->inequality_directions;
    delete edata->feature_data;
    
    // Feature mappings
    delete[] edata->feature_mappings->reverse_num_feature_mapping;
    delete[] edata->feature_mappings->reverse_cat_feature_mapping;
    delete[] edata->feature_mappings->feature_mapping;
    delete[] edata->feature_mappings->mapping_numerics;
    delete edata->feature_mappings;
    
    // Monotonic constraints
    delete[] edata->mono_constraints->feature_idx;
    delete[] edata->mono_constraints->output_idx;
    delete[] edata->mono_constraints->constraint;
    delete edata->mono_constraints;
    
    delete edata;
}

void export_ensemble_data(std::ofstream& header_file, const std::string& model_name, ensembleData *edata, ensembleMetaData *metadata, deviceType device, std::vector<Optimizer*> opts, exportFormat export_format, exportType export_type, const std::string &prefix)
{
    std::string type_name;
    switch (export_format){
       case EXP_FLOAT:
            type_name = "float";
            break;
        case EXP_FXP8:
            type_name = "int16";
            break;
        case EXP_FXP16:
            type_name = "int32";
            break;
        default:
            std::cerr << "Invalid exportFormat!" << std::endl;
            return; // Exit the function if the format is invalid
    }

    if (export_type == exportType::COMPACT){
        if (metadata->max_depth > 6 || metadata->grow_policy != growPolicy::OBLIVIOUS){
            std::cerr << "Cannot only compact export with max depth <= 6 and oblivious trees" << std::endl;
            return; // Exit the function if the format is invalid
        }
    }
    ensembleData *edata_cpu = nullptr;
#ifdef USE_CUDA
    if (device == gpu){
        edata_cpu = ensemble_data_copy_gpu_cpu(metadata, edata, nullptr);
    }
#endif 
    if (device == cpu)
        edata_cpu = edata;
    
    int binary_splits = 0;
    for (int i  = 0; i < metadata->n_trees; ++i){
        binary_splits += edata_cpu->ensemble_info->depths[i];
    }

    for (size_t opt_idx = 0; opt_idx < opts.size(); ++opt_idx){
        optimizerAlgo algo = opts[opt_idx]->getAlgo();
        if (algo != SGD){
            std::cerr << "Error. Can only export SGD optimizers" << std::endl;
            header_file.close();
            throw std::runtime_error("Error. Can only export SGD optimizers");
            return;
        }
    }

    header_file << "#ifndef GBRL_MODEL_H\n";
    header_file << "#define GBRL_MODEL_H\n\n";

    
    header_file << "/*\n";

    if (!model_name.empty()) {
           header_file << "###########################\n";
        header_file << "model_name: " << model_name << "\n";
    }
    header_file << "###########################\n";
    header_file << "n_leaves: " << metadata->n_leaves << ", ";
    header_file << "n_trees: " << metadata->n_trees << ", ";
    header_file << "max_trees: " << metadata->max_trees << ", ";
    header_file << "max_leaves: " << metadata->max_leaves << ", ";
    header_file << "max_trees_batch: " << metadata->max_trees_batch << ", ";
    header_file << "max_leaves_batch: " << metadata->max_leaves_batch << ", ";
    header_file << "input_dim: " << metadata->input_dim << ", ";
    header_file << "output_dim: " << metadata->output_dim << ", ";
    header_file << "policy_dim: " << metadata->policy_dim << ", ";
    header_file << "\nmax_depth: " << metadata->max_depth << ", ";
    header_file << "min_data_in_leaf: " << metadata->min_data_in_leaf << ", ";
    header_file << "n_bins: " << metadata->n_bins << ", ";
    header_file << "par_th: " << metadata->par_th << ", ";
    header_file << "cv_beta: " << metadata->cv_beta << ", ";
    header_file << "verbose: " << metadata->verbose << ", ";
    header_file << "batch_size: " << metadata->batch_size << ", ";
    header_file << "use_cv: " << metadata->use_cv;
    header_file << "\nsplit_score_func: " << scoreFuncToString(metadata->split_score_func) << ", ";
    header_file << "generator_type: " << generatorTypeToString(metadata->generator_type) << ", ";
    header_file << "grow_policy: " << growPolicyToString(metadata->grow_policy) << ", ";
    header_file << "n_num_features: " << metadata->n_num_features << ", ";
    header_file << "n_cat_features: " << metadata->n_cat_features << ", ";
    header_file << "iteration: " << metadata->iteration;
    header_file << "alloc_data_size: " << edata->alloc_data_size;
    header_file << "\n*/\n";

    header_file << "#define " << prefix << "N_TREES " << metadata->n_trees << "\n";
    header_file << "#define " << prefix << "N_LEAVES " << metadata->n_leaves << "\n";
    header_file << "#define " << prefix << "BINARY_FEATURES " << binary_splits << "\n";
    header_file << "#define N_INPUTS " << metadata->input_dim << "\n";
    header_file << "#define " << prefix << "N_OUTPUTS " << metadata->output_dim << "\n";
    header_file << "#define " << prefix << "N_FEATURES " << metadata->n_num_features  << "\n\n";
    if (metadata->output_dim > 1){
        header_file << "static inline void gbrl_predict(" << type_name << " *results, const " << type_name << " *features){\n\n";
    } else {
        header_file << "static inline " << type_name << " gbrl_predict(const " << type_name << " *features){\n\n";
        header_file << "\t" << type_name << " result = ";
        switch (export_format){
            case EXP_FLOAT:
                header_file << "0.0f;\n";
                break;
            case EXP_FXP8:
                header_file << "0;\n";
                break;
            case EXP_FXP16:
                header_file << "0;\n";
                break;
        }
    }
    
    header_file << "\tunsigned int tree_idx, idx, leaf_ptr, cond_ptr";
    if (metadata->output_dim > 1){
        header_file << ", j";
    }
    if (export_type == exportType::FULL)
        header_file << ", depth, current_depth";
    header_file << ";\n";
    header_file << "\t/* Model data */\n";
    if (export_type == exportType::FULL){
        header_file << "\tconst unsigned int depths[" << prefix << "N_TREES] = {";
        for (int i  = 0; i < metadata->n_trees; ++i){
            header_file << edata_cpu->ensemble_info->depths[i];
            if (i < metadata->n_trees - 1)
                header_file << ", ";
        }
        header_file << "};\n";
    }
    if (metadata->output_dim > 1){
        header_file << "\tconst " << type_name << " bias[" << prefix << "N_OUTPUTS] = {";
        for (int i  = 0; i < metadata->output_dim; ++i){
            switch (export_format){
                case EXP_FLOAT:
                    header_file << edata_cpu->bias[i];
                    break;
                case EXP_FXP8:
                    header_file << float_to_int16(edata_cpu->bias[i]);
                    break;
                case EXP_FXP16:
                    header_file << float_to_int32(edata_cpu->bias[i]);
                    break;
            }
            if (i < metadata->output_dim - 1)
                header_file << ", ";
        }

    } else {
        header_file << "\tconst " << type_name << " bias = ";
        switch (export_format){
            case EXP_FLOAT:
                header_file << edata_cpu->bias[0];
                break;
            case EXP_FXP8:
                header_file << float_to_int16(edata_cpu->bias[0]);
                break;
            case EXP_FXP16:
                header_file << float_to_int32(edata_cpu->bias[0]);
                break;
        }
    }
    header_file << ";\n";

    std::string max_size = (metadata->input_dim < 255) ? "uint8" : "uint16";
    header_file << "\tconst " << max_size <<  " feature_indices[" << prefix << "BINARY_FEATURES] = {";
    for (int i  = 0; i < binary_splits; ++i){
        header_file << edata_cpu->feature_data->feature_indices[i];
        if (i < binary_splits - 1)
            header_file << ", ";
    }
    header_file << "};\n";

    header_file << "\tconst " << type_name << " feature_values[" << prefix << "BINARY_FEATURES] = {";
    for (int i  = 0; i < binary_splits; ++i){
        switch (export_format){
            case EXP_FLOAT:
                header_file << edata_cpu->feature_data->feature_values[i];
                break;
            case EXP_FXP8:
                   header_file << float_to_int16(edata_cpu->feature_data->feature_values[i]);
                break;
            case EXP_FXP16:
                   header_file << float_to_int32(edata_cpu->feature_data->feature_values[i]);
                break;
        }
        if (i < binary_splits - 1)
            header_file << ", ";
    }
    header_file << "};\n";

    header_file << "\tconst " << type_name << " leaf_values[" << prefix << "N_LEAVES*" << prefix << "N_OUTPUTS]  = {";
    int tree_idx = 0;
    int limit_leaf_idx = edata_cpu->ensemble_info->tree_indices[tree_idx];
    float value;
    for (int i  = 0; i < metadata->n_leaves; ++i){
        if (i > limit_leaf_idx){
            tree_idx += 1;
            limit_leaf_idx = edata_cpu->ensemble_info->tree_indices[tree_idx];
        }
        int value_idx = i*metadata->output_dim;
        for (size_t opt_idx = 0; opt_idx < opts.size(); ++opt_idx){
            for (int j=opts[opt_idx]->start_idx; j < opts[opt_idx]->stop_idx; ++j){
                value = -edata_cpu->leaf_data->values[value_idx + j] * opts[opt_idx]->scheduler->get_lr(tree_idx);
                switch (export_format){
                    case EXP_FLOAT:
                        header_file << value;
                        break;
                    case EXP_FXP8:
                        header_file << float_to_int16(value);
                        break;
                    case EXP_FXP16:
                        header_file << float_to_int32(value);
                        break;
                }
                if ((i < metadata->n_leaves - 1) || (j < metadata->output_dim - 1  && i == metadata->n_leaves - 1))
                    header_file << ", ";
            }
        }
    }
    header_file << "};\n";

    header_file << "\tleaf_ptr = 0;\n";
    header_file << "\tcond_ptr = 0;\n";
    header_file << "\tunsigned char pass;\n";
    header_file << "\tfor (tree_idx = 0; tree_idx < " << prefix << "N_TREES; ++tree_idx)\n";
    header_file << "\t{\n";
    if (export_type == exportType::COMPACT){
        header_file << "\t\tidx = 0;\n";
        for (int depth = 0; depth < metadata->max_depth; ++depth){
            
            header_file << "\t\tpass = (unsigned char)(features[feature_indices[cond_ptr + " << depth << "]] > feature_values[cond_ptr + " << depth << "]);\n";
            header_file << "\t\tidx |= (pass <<  (" << metadata->max_depth << " - 1 - " << depth << "));\n";
        }
    } else {
        header_file << "\t\tcurrent_depth = depths[tree_idx];\n";
        header_file << "\t\tidx = 0;\n";
        header_file << "\t\tfor (depth = 0; depth < current_depth; ++depth){\n";
        header_file << "\t\t\tpass = (unsigned char)(features[feature_indices[cond_ptr + depth]] > feature_values[cond_ptr + depth]);\n";
        header_file << "\t\t\tidx |= (pass <<  (current_depth - 1 - depth));\n";
        header_file << "\t\t}\n";
    }
    
    if (metadata->output_dim > 1){
        header_file << "\t\tfor (j = 0 ; j < " << prefix << "N_OUTPUTS; j++)\n";
        header_file << "\t\t\tresults[j] += leaf_values[(leaf_ptr + idx)*" << prefix << "N_OUTPUTS + j];\n";
    } else {
        header_file << "\t\tresult += leaf_values[leaf_ptr + idx];\n";
    }

    if (export_type == exportType::COMPACT){
        int tmp = 1 << metadata->max_depth;
        header_file << "\t\tleaf_ptr += " << tmp << ";\n";
        header_file << "\t\tcond_ptr += " << metadata->max_depth << ";\n";
    } else {
        header_file << "\t\tleaf_ptr += (1 << current_depth);\n";
        header_file << "\t\tcond_ptr += current_depth;\n";
    }
    header_file << "\t}\n";
    if (metadata->output_dim > 1){
        header_file << "\tfor (j = 0 ; j < " << prefix << "N_OUTPUTS; j++)\n";
        header_file << "\t\tresults[j] += bias[j];\n";
    } else {
        header_file << "\tresult += bias;\n";
        header_file << "\treturn result;\n";
    }
    header_file << "}\n";
    header_file << "#endif\n";

#ifdef USE_CUDA
    if (device == gpu){
        ensemble_data_dealloc(edata_cpu);
    }
#endif 
}

void save_ensemble_data(std::ofstream& file, ensembleData *edata, ensembleMetaData *metadata, deviceType device){
    if (!file.is_open() || file.fail()) {
        std::cerr << "Error file is not open for writing: " << std::endl;
        throw std::runtime_error("Error opening file");
    }
    ensembleData *edata_cpu = nullptr;
#ifdef USE_CUDA
    if (device == gpu){
        edata_cpu = ensemble_data_copy_gpu_cpu(metadata, edata, nullptr);
    }
#endif 
    if (device == cpu)
        edata_cpu = edata;
    NULL_CHECK check = edata_cpu->bias != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->bias != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->bias), metadata->output_dim * sizeof(float));
    check = edata_cpu->feature_data->feature_weights != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_data->feature_weights != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_data->feature_weights), metadata->input_dim * sizeof(float));
#ifdef DEBUG
    check = edata_cpu->n_samples != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->n_samples != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->n_samples), metadata->n_leaves * sizeof(int));
#endif 
    check = edata_cpu->ensemble_info->tree_indices != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->ensemble_info->tree_indices != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->ensemble_info->tree_indices), metadata->n_trees * sizeof(int));
    size_t sizes = (metadata->grow_policy == OBLIVIOUS) ? metadata->n_trees: metadata->n_leaves;
    check = edata_cpu->ensemble_info->depths != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->ensemble_info->depths != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->ensemble_info->depths), sizes * sizeof(int));
    check = edata_cpu->leaf_data->values != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->leaf_data->values != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->leaf_data->values), metadata->n_leaves * metadata->output_dim * sizeof(float));
    check = edata_cpu->feature_data->feature_indices != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_data->feature_indices != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_data->feature_indices), metadata->max_depth * sizes * sizeof(int));
    check = edata_cpu->feature_data->feature_values != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_data->feature_values != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_data->feature_values), metadata->max_depth * sizes * sizeof(float));
    check = edata_cpu->leaf_data->edge_weights != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->leaf_data->edge_weights != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->leaf_data->edge_weights), metadata->max_depth * metadata->n_leaves * sizeof(float));
    // monotonic constraints
    check = edata_cpu->mono_constraints->feature_idx != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->mono_constraints->feature_idx != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->mono_constraints->feature_idx), metadata->n_mono_constraints * sizeof(int));
    check = edata_cpu->mono_constraints->output_idx != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->mono_constraints->output_idx != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->mono_constraints->output_idx), metadata->n_mono_constraints * sizeof(int));
    check = edata_cpu->mono_constraints->constraint != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->mono_constraints->constraint != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->mono_constraints->constraint), metadata->n_mono_constraints * sizeof(int));    

    check = edata_cpu->feature_mappings->reverse_num_feature_mapping != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_mappings->reverse_num_feature_mapping != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_mappings->reverse_num_feature_mapping), metadata->input_dim * sizeof(int));
    check = edata_cpu->feature_mappings->reverse_cat_feature_mapping != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_mappings->reverse_cat_feature_mapping != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_mappings->reverse_cat_feature_mapping), metadata->input_dim * sizeof(int));
    check = edata_cpu->feature_mappings->feature_mapping != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_mappings->feature_mapping != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_mappings->feature_mapping), metadata->input_dim * sizeof(int));
    check = edata_cpu->feature_mappings->mapping_numerics != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_mappings->mapping_numerics != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_mappings->mapping_numerics), metadata->input_dim * sizeof(bool));
    check = edata_cpu->feature_data->is_numerics != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_data->is_numerics != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_data->is_numerics), metadata->max_depth * sizes * sizeof(bool));
    check = edata_cpu->feature_data->inequality_directions != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_data->inequality_directions != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_data->inequality_directions), metadata->max_depth * metadata->n_leaves * sizeof(bool));
    check = edata_cpu->feature_data->categorical_values != nullptr ? VALID : NULL_OPT;
    file.write(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (edata_cpu->feature_data->categorical_values != nullptr)
        file.write(reinterpret_cast<char*>(edata_cpu->feature_data->categorical_values), metadata->max_depth * sizes * sizeof(char) * MAX_CHAR_SIZE);

#ifdef USE_CUDA
    if (device == gpu){
        ensemble_data_dealloc(edata_cpu);
    }
#endif 
}

ensembleData* load_ensemble_data(std::ifstream& file, ensembleMetaData *metadata){
    if (!file.is_open() || file.fail()) {
        std::cerr << "Error file is not open for writing: " << std::endl;
        throw std::runtime_error("Error opening file");
    }
    ensembleData *edata_cpu = ensemble_data_alloc(metadata);
    NULL_CHECK check;
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->bias), metadata->output_dim * sizeof(float));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_data->feature_weights), metadata->input_dim * sizeof(float));
    } 
#ifdef DEBUG
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->n_samples), metadata->n_leaves * sizeof(int));
    } 
#endif 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->ensemble_info->tree_indices), metadata->n_trees * sizeof(int));
    } 
    size_t sizes = (metadata->grow_policy == OBLIVIOUS) ? metadata->n_trees: metadata->n_leaves;
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->ensemble_info->depths), sizes * sizeof(int));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->leaf_data->values), metadata->output_dim * metadata->n_leaves * sizeof(float));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_data->feature_indices), metadata->max_depth * sizes * sizeof(int));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_data->feature_values), metadata->max_depth * sizes * sizeof(float));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->leaf_data->edge_weights), metadata->max_depth * metadata->n_leaves * sizeof(float));
    } 
    // monotonic constraints
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->mono_constraints->feature_idx), metadata->n_mono_constraints * sizeof(int));
    }
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->mono_constraints->output_idx), metadata->n_mono_constraints * sizeof(int));
    }
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->mono_constraints->constraint), metadata->n_mono_constraints * sizeof(int));
    }

    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_mappings->reverse_num_feature_mapping), metadata->input_dim * sizeof(int));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_mappings->reverse_cat_feature_mapping), metadata->input_dim * sizeof(int));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_mappings->feature_mapping), metadata->input_dim * sizeof(int));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
       if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_mappings->mapping_numerics), metadata->input_dim * sizeof(bool));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_data->is_numerics), metadata->max_depth * sizes * sizeof(bool));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_data->inequality_directions), metadata->max_depth * metadata->n_leaves * sizeof(bool));
    } 
    file.read(reinterpret_cast<char*>(&check), sizeof(NULL_CHECK));
    if (check == VALID) {
        file.read(reinterpret_cast<char*>(edata_cpu->feature_data->categorical_values), metadata->max_depth * sizes * sizeof(char) * MAX_CHAR_SIZE);
    } 
    return edata_cpu;
}


void allocate_ensemble_memory(ensembleMetaData *metadata, ensembleData *edata){
    int leaf_idx = metadata->n_leaves, tree_idx = metadata->n_trees; 
    if ((leaf_idx >= metadata->max_leaves) || (tree_idx >= metadata->max_trees)){
        int new_size_leaves = metadata->n_leaves + metadata->max_leaves_batch;
        int new_tree_size = metadata->n_trees + metadata->max_trees_batch;
        metadata->max_leaves = new_size_leaves;
        metadata->max_trees = new_tree_size;
        ensembleData *new_data = ensemble_data_alloc(metadata);
        memcpy(new_data->bias, edata->bias, metadata->output_dim * sizeof(float));
        memcpy(new_data->feature_data->feature_weights, edata->feature_data->feature_weights, metadata->input_dim * sizeof(float));
#ifdef DEBUG
        memcpy(new_data->n_samples, edata->n_samples, leaf_idx * sizeof(int));
#endif 
        memcpy(new_data->leaf_data->values, edata->leaf_data->values, leaf_idx * metadata->output_dim * sizeof(float));
        memcpy(new_data->ensemble_info->tree_indices, edata->ensemble_info->tree_indices, tree_idx * sizeof(int));
        memcpy(new_data->feature_data->inequality_directions, edata->feature_data->inequality_directions, leaf_idx * metadata->max_depth * sizeof(bool));
        memcpy(new_data->leaf_data->edge_weights, edata->leaf_data->edge_weights, leaf_idx * metadata->max_depth * sizeof(float));
        memcpy(new_data->feature_mappings->reverse_cat_feature_mapping, edata->feature_mappings->reverse_cat_feature_mapping, metadata->input_dim * sizeof(int));
        memcpy(new_data->feature_mappings->reverse_num_feature_mapping, edata->feature_mappings->reverse_num_feature_mapping, metadata->input_dim * sizeof(int));
        memcpy(new_data->feature_mappings->feature_mapping, edata->feature_mappings->feature_mapping, metadata->input_dim * sizeof(int));
        // monotonic constraints
        memcpy(new_data->mono_constraints->feature_idx, edata->mono_constraints->feature_idx, metadata->n_mono_constraints * sizeof(int));
        memcpy(new_data->mono_constraints->output_idx, edata->mono_constraints->output_idx, metadata->n_mono_constraints * sizeof(int));
        memcpy(new_data->mono_constraints->constraint, edata->mono_constraints->constraint, metadata->n_mono_constraints * sizeof(int));
        memcpy(new_data->feature_mappings->mapping_numerics, edata->feature_mappings->mapping_numerics, metadata->input_dim * sizeof(bool));
        if (metadata->grow_policy == GREEDY){
            memcpy(new_data->ensemble_info->depths, edata->ensemble_info->depths, leaf_idx * sizeof(int));
            memcpy(new_data->feature_data->feature_indices, edata->feature_data->feature_indices, leaf_idx * metadata->max_depth * sizeof(int));
            memcpy(new_data->feature_data->feature_values, edata->feature_data->feature_values, leaf_idx * metadata->max_depth * sizeof(float));
            memcpy(new_data->feature_data->is_numerics, edata->feature_data->is_numerics, leaf_idx * metadata->max_depth * sizeof(bool));
            memcpy(new_data->feature_data->categorical_values, edata->feature_data->categorical_values, leaf_idx * metadata->max_depth * sizeof(char) * MAX_CHAR_SIZE);
        } else {
            memcpy(new_data->ensemble_info->depths, edata->ensemble_info->depths, tree_idx * sizeof(int));
            memcpy(new_data->feature_data->feature_indices, edata->feature_data->feature_indices, tree_idx * metadata->max_depth * sizeof(int));
            memcpy(new_data->feature_data->feature_values, edata->feature_data->feature_values, tree_idx * metadata->max_depth * sizeof(float));
            memcpy(new_data->feature_data->is_numerics, edata->feature_data->is_numerics, tree_idx * metadata->max_depth * sizeof(bool));
            memcpy(new_data->feature_data->categorical_values, edata->feature_data->categorical_values, tree_idx * metadata->max_depth * sizeof(char) * MAX_CHAR_SIZE);
        }
        delete[] edata->bias;
        delete[] edata->feature_data->feature_weights;
#ifdef DEBUG
        delete [] edata->n_samples;
#endif
        delete[] edata->ensemble_info->depths;
        delete[] edata->leaf_data->values;
        // monotonic constraints
        delete[] edata->mono_constraints->feature_idx;
        delete[] edata->mono_constraints->output_idx;
        delete[] edata->mono_constraints->constraint;
        // leaf data
        delete[] edata->feature_data->feature_indices;
        delete[] edata->ensemble_info->tree_indices;
        delete[] edata->feature_data->feature_values;
        delete[] edata->leaf_data->edge_weights;
        delete[] edata->feature_data->is_numerics;
        delete[] edata->feature_data->categorical_values;
        delete[] edata->feature_data->inequality_directions; 
        delete[] edata->feature_mappings->reverse_cat_feature_mapping;
        delete[] edata->feature_mappings->feature_mapping;
        delete[] edata->feature_mappings->reverse_num_feature_mapping;
        delete[] edata->feature_mappings->mapping_numerics;

        edata->bias = new_data->bias;
        edata->feature_data->feature_weights = new_data->feature_data->feature_weights;
#ifdef DEBUG
        edata->n_samples = new_data->n_samples;
#endif
        edata->ensemble_info->depths = new_data->ensemble_info->depths;
        edata->ensemble_info->tree_indices = new_data->ensemble_info->tree_indices;
        edata->leaf_data->values = new_data->leaf_data->values;
        edata->feature_data->inequality_directions = new_data->feature_data->inequality_directions;
        edata->feature_data->feature_indices = new_data->feature_data->feature_indices;
        edata->feature_data->feature_values = new_data->feature_data->feature_values;
        edata->leaf_data->edge_weights = new_data->leaf_data->edge_weights;
        edata->mono_constraints->feature_idx = new_data->mono_constraints->feature_idx;
        edata->mono_constraints->output_idx = new_data->mono_constraints->output_idx;
        edata->mono_constraints->constraint = new_data->mono_constraints->constraint;
        edata->feature_mappings->reverse_cat_feature_mapping = new_data->feature_mappings->reverse_cat_feature_mapping;
        edata->feature_mappings->reverse_num_feature_mapping = new_data->feature_mappings->reverse_num_feature_mapping;
        edata->feature_mappings->feature_mapping = new_data->feature_mappings->feature_mapping;
        edata->feature_mappings->mapping_numerics = new_data->feature_mappings->mapping_numerics;
        edata->feature_data->is_numerics = new_data->feature_data->is_numerics;
        edata->feature_data->categorical_values = new_data->feature_data->categorical_values;
        
        // Delete only the sub-struct containers and ensembleData, not the underlying arrays (they're now owned by edata)
        delete new_data->ensemble_info;
        delete new_data->leaf_data;
        delete new_data->feature_data;
        delete new_data->mono_constraints;
        delete new_data->feature_mappings;
        delete new_data;
    }
}