/**
 * JIT template — auto-generated by cugenopt.solve_custom()
 *
 * Placeholders (replaced by Python at runtime):
 *   {{D1}}, {{D2}}             — Solution dimensions
 *   {{ENCODING}}               — EncodingType enum value
 *   {{ROW_MODE}}               — RowMode enum value
 *   {{DIM1}}, {{DIM2_DEFAULT}}, {{TOTAL_ELEMENTS}}
 *   {{VALUE_LOWER}}, {{VALUE_UPPER}}
 *   {{CROSS_ROW_PROB}}, {{PERM_REPEAT_COUNT}}
 *   {{OBJ_DEFS}}               — ObjDef array initializer
 *   {{DATA_FIELDS}}            — struct field declarations
 *   {{COMPUTE_OBJ}}            — user's compute_obj body
 *   {{COMPUTE_PENALTY}}        — user's compute_penalty body
 *   {{SHARED_MEM_EXPR}}        — shared_mem_bytes() return expression
 *   {{LOAD_SHARED_BODY}}       — load_shared() body
 *   {{DESTROY_BODY}}           — destroy() body
 *   {{DATA_LOAD_BODY}}         — main() data loading code
 *   {{SOLVER_CONFIG}}          — SolverConfig field assignments
 *   {{NUM_OBJ}}                — number of objectives
 *   {{CUSTOM_OP_DEFINES}}      — #define CUGENOPT_HAS_CUSTOM_OPS (or empty)
 *   {{CUSTOM_OP_SWITCH}}       — switch cases for execute_custom_op (or empty)
 *   {{CUSTOM_OP_REGISTRY}}     — add() calls for register_custom_operators (or empty)
 */

{{CUSTOM_OP_DEFINES}}

#include "types.cuh"
#include "cuda_utils.cuh"

struct CustomProblem;

#ifdef CUGENOPT_HAS_CUSTOM_OPS
template<typename Sol>
__device__ inline bool execute_custom_op(int seq_id, Sol& sol, int dim1,
                                          EncodingType encoding, curandState* rng,
                                          int val_lb, int val_ub,
                                          const void* prob_data);
#endif

#include "solver.cuh"
#include <cstdio>
#include <cstdlib>
#include <cstring>

struct CustomProblem : ProblemBase<CustomProblem, {{D1}}, {{D2}}> {
    {{DATA_FIELDS}}
    int _n;

    static constexpr ObjDef OBJ_DEFS[] = {
        {{OBJ_DEFS}}
    };

    __device__ float compute_obj(int idx, const Sol& sol) const {
        {{COMPUTE_OBJ}}
    }

    __device__ float compute_penalty(const Sol& sol) const {
        {{COMPUTE_PENALTY}}
    }

    ProblemConfig config() const {
        ProblemConfig cfg;
        cfg.encoding = {{ENCODING}};
        cfg.dim1 = {{DIM1}};
        cfg.dim2_default = {{DIM2_DEFAULT}};
        cfg.row_mode = {{ROW_MODE}};
        cfg.total_elements = {{TOTAL_ELEMENTS}};
        cfg.cross_row_prob = {{CROSS_ROW_PROB}};
        cfg.perm_repeat_count = {{PERM_REPEAT_COUNT}};
        cfg.value_lower_bound = {{VALUE_LOWER}};
        cfg.value_upper_bound = {{VALUE_UPPER}};
        fill_obj_config(cfg);
        return cfg;
    }

    size_t shared_mem_bytes() const {
        {{SHARED_MEM_EXPR}}
    }

    __device__ void load_shared(char* smem, int tid, int bsz) {
        {{LOAD_SHARED_BODY}}
    }

    void destroy() {
        {{DESTROY_BODY}}
    }
};

#ifdef CUGENOPT_HAS_CUSTOM_OPS
template<typename Sol>
__device__ inline bool execute_custom_op(int seq_id, Sol& sol, int dim1,
                                          EncodingType encoding, curandState* rng,
                                          int val_lb, int val_ub,
                                          const void* prob_data) {
    const CustomProblem* prob = static_cast<const CustomProblem*>(prob_data);
    switch (seq_id) {
{{CUSTOM_OP_SWITCH}}
    default:
        return false;
    }
}
#endif

#ifdef CUGENOPT_HAS_CUSTOM_OPS
inline void register_custom_operators(SeqRegistry& reg) {
    auto add = [&](int id, float w, float cap) {
        if (reg.count >= MAX_SEQ) return;
        reg.ids[reg.count] = id;
        reg.weights[reg.count] = w;
        reg.max_w[reg.count] = cap;
        reg.count++;
    };
{{CUSTOM_OP_REGISTRY}}
    float sum = 0.0f;
    for (int i = 0; i < reg.count; i++) sum += reg.weights[i];
    if (sum > 0.0f) {
        for (int i = 0; i < reg.count; i++) reg.weights[i] /= sum;
    }
}
#endif

static float* read_binary_floats(const char* path, int count) {
    FILE* f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "Cannot open %s\n", path); exit(1); }
    float* buf = new float[count];
    size_t r = fread(buf, sizeof(float), count, f);
    fclose(f);
    if ((int)r != count) { fprintf(stderr, "Short read %s: %d/%d\n", path, (int)r, count); exit(1); }
    return buf;
}

static int* read_binary_ints(const char* path, int count) {
    FILE* f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "Cannot open %s\n", path); exit(1); }
    int* buf = new int[count];
    size_t r = fread(buf, sizeof(int), count, f);
    fclose(f);
    if ((int)r != count) { fprintf(stderr, "Short read %s: %d/%d\n", path, (int)r, count); exit(1); }
    return buf;
}

int main(int argc, char** argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <data_dir>\n", argv[0]);
        return 1;
    }
    const char* data_dir = argv[1];
    char path[512];

    // Read problem size
    snprintf(path, sizeof(path), "%s/n.bin", data_dir);
    int n;
    { FILE* f = fopen(path, "rb"); fread(&n, sizeof(int), 1, f); fclose(f); }

    // Read data arrays and upload to GPU
    CustomProblem prob;
    prob._n = n;
    {{DATA_LOAD_BODY}}

    // Configure solver
    SolverConfig cfg;
    {{SOLVER_CONFIG}}

    // Solve
#ifdef CUGENOPT_HAS_CUSTOM_OPS
    auto result = solve(prob, cfg, nullptr, 0, register_custom_operators);
#else
    auto result = solve(prob, cfg);
#endif

    // Output JSON result
    auto& best = result.best_solution;
    printf("{\"objective\":%.10g", best.objectives[0]);
    printf(",\"penalty\":%.10g", best.penalty);
    printf(",\"elapsed_ms\":%.4f", result.elapsed_ms);
    printf(",\"generations\":%d", result.generations);
    printf(",\"stop_reason\":\"%s\"",
           result.stop_reason == StopReason::TimeLimit ? "time_limit" :
           result.stop_reason == StopReason::Stagnation ? "stagnation" : "max_gen");

    printf(",\"objectives\":[");
    for (int i = 0; i < {{NUM_OBJ}}; i++) {
        if (i > 0) printf(",");
        printf("%.10g", best.objectives[i]);
    }
    printf("]");

    printf(",\"solution\":[");
    for (int r = 0; r < {{DIM1}}; r++) {
        if (r > 0) printf(",");
        printf("[");
        int len = best.dim2_sizes[r];
        for (int c = 0; c < len; c++) {
            if (c > 0) printf(",");
            printf("%d", best.data[r][c]);
        }
        printf("]");
    }
    printf("]}\n");

    prob.destroy();
    return 0;
}