edge_wrapper.cpp

/**
 * weather_edge_wrapper.cpp
 * ======================
 * Edge runtime wrapper for the Weather RL model (.mnn).
 *
 * Design:
 *   - External GRU hidden-state management (hidden_in → hidden_out per step).
 *   - Action masking applied here, after logits are read from the model.
 *   - Vulkan preferred (Mali-G31), CPU fallback if Vulkan session fails.
 *   - Input tensors cached at init; not re-fetched on every step.
 *   - C interface exposed for processing_nodes integration.
 *
 * Tensor names (match StatelessInferenceWrapper in mnn_export.py):
 *   Inputs (in canonical sorted key order, then hidden_in):
 *     "basin_context"        [1, 4]                      flat: BASIN_FLAT (schema v3+)
 *     "forecast_precip"      [1, N_ZONES, HORIZON_DAYS]  flat: N_ZONES * HORIZON_DAYS
 *     "forecast_uncertainty" [1, N_ZONES]
 *     "prior_belief"         [1, 1]
 *     "zone_belief"          [1, N_ZONES]
 *     "hidden_in"            [1, 1, HIDDEN_SIZE]         flat: HIDDEN_SIZE
 *   Outputs:
 *     "action_logits"        [1, N_ACTIONS]
 *     "hidden_out"           [1, 1, HIDDEN_SIZE]         flat: HIDDEN_SIZE
 *
 * Notes:
 *   - action_mask is NOT in the MNN graph — applied externally here.
 *     Terminate action (index N_ZONES) is never masked.
 *   - export_keys in mnn_export.py are sorted(obs_keys) with action_mask
 *     excluded, so the canonical input order is alphabetical:
 *       basin_context, forecast_precip, forecast_uncertainty,
 *       prior_belief, zone_belief
 *   - basin_context (schema v3) is [enso_oni, iod_dmi, itcz_latitude_deg,
 *     mslp_regional_hpa], matching zone_observation.BasinContext field
 *     order. When no basin data is available at the edge, pass the neutral
 *     default {0.0f, 0.0f, 0.0f, 1013.25f} (same default
 *     weather_forecast_env.py uses for episodes without basin data).
 *     BREAKING vs the pre-v3 interface: runStep()/weather_run() gained a
 *     basin_context parameter, and the MNN graph gained the input — old
 *     callers and old .mnn models must be updated together.
 *   - forecast_precip is [N_ZONES, HORIZON_DAYS] — pass it row-major
 *     (all days for zone 0, then all days for zone 1, etc.).
 *   - forecast_uncertainty is [N_ZONES] — one scalar per zone, not per day.
 *   - HIDDEN_SIZE default is 64 (gru_weather_policy.py default).
 *     N_ZONES and HORIZON_DAYS vary by curriculum phase; they are set here
 *     to the final curriculum phase (heatwave/humidity: 4 zones, 30 days).
 *     Always confirm these match your deployed checkpoint.
 *   - Backend order: Vulkan → CPU.
 *
 * Bug 2.5 fix:
 *   weather_run() and weather_destroy() now guard against null handles.
 *   weather_create() already returns nullptr on failure, but callers that
 *   omit the null check and pass the result directly to weather_run() would
 *   previously trigger undefined behaviour (segfault on dereference).
 *   weather_run() now returns -1 immediately on a null handle.
 *   weather_destroy() now no-ops on a null handle (delete nullptr is safe
 *   in C++, but the explicit guard makes the contract visible to callers).
 */

#include <MNN/Interpreter.hpp>
#include <MNN/Tensor.hpp>
#include <algorithm>
#include <cassert>
#include <cstring>
#include <iostream>
#include <memory>
#include <string>
#include <vector>

// ---------------------------------------------------------------------------
// Shape constants — must match the exported checkpoint's ForecastConfig
//
// Default ForecastConfig (zone_observation.py):
//   horizon_days = 30
//   n_zones      = 1   (single-zone baseline)
//
// Final curriculum phase (train_curriculum.py heatwave/humidity):
//   horizon_days = 30   (ForecastConfig default, unchanged by curriculum)
//   n_zones      = 4
//
// GRU policy default (gru_weather_policy.py):
//   hidden_size  = 64
//
// ---------------------------------------------------------------------------
static constexpr int N_ZONES        = 4;    // zones in the deployed checkpoint
static constexpr int HORIZON_DAYS   = 30;   // ForecastConfig.horizon_days default
static constexpr int N_LAYERS       = 1;    // GRU num_layers (always 1)
static constexpr int HIDDEN_SIZE    = 64;   // GRU hidden_size default
static constexpr int N_ACTIONS      = N_ZONES + 1;  // zone actions + terminate

// Derived sizes (flat float counts)
static constexpr int PRECIP_FLAT    = N_ZONES * HORIZON_DAYS;  // forecast_precip flat size
static constexpr int HIDDEN_FLAT    = N_LAYERS * 1 * HIDDEN_SIZE;  // batch=1
static constexpr int BASIN_FLAT     = 4;    // basin_context: [enso_oni, iod_dmi, itcz_lat, mslp_hpa]

// Neutral basin default (matches weather_forecast_env._BASIN_CONTEXT_NEUTRAL).
static constexpr float BASIN_NEUTRAL[BASIN_FLAT] = {0.0f, 0.0f, 0.0f, 1013.25f};

// ---------------------------------------------------------------------------
// WeatherEdgeWrapper
// ---------------------------------------------------------------------------
class WeatherEdgeWrapper {
public:
    explicit WeatherEdgeWrapper(const std::string& model_path)
        : model_path_(model_path)
    {
        net_ = std::shared_ptr<MNN::Interpreter>(
            MNN::Interpreter::createFromFile(model_path.c_str()),
            MNN::Interpreter::destroy
        );

        if (!net_) {
            std::cerr << "[WeatherWrapper] FATAL: failed to load model from "
                      << model_path << std::endl;
            return;
        }

        // --- Try Vulkan first (Mali-G31), fall back to CPU ---
        session_ = tryCreateSession(MNN_FORWARD_VULKAN, "Vulkan");
        if (!session_) {
            std::cerr << "[WeatherWrapper] Vulkan unavailable, falling back to CPU" << std::endl;
            session_ = tryCreateSession(MNN_FORWARD_CPU, "CPU");
        }

        if (!session_) {
            std::cerr << "[WeatherWrapper] FATAL: could not create any MNN session" << std::endl;
            return;
        }

        // --- Resize all input tensors before resizeSession ---
        // mnn_export.py uses dynamic batch axes; providing explicit shapes
        // ensures MNN allocates memory correctly for batch=1 inference.
        // Input order matches sorted export_keys + hidden_in (alphabetical):
        //   basin_context, forecast_precip, forecast_uncertainty,
        //   prior_belief, zone_belief, hidden_in
        auto* t_basin = net_->getSessionInput(session_, "basin_context");
        if (t_basin) net_->resizeTensor(t_basin, {1, BASIN_FLAT});

        auto* t_precip = net_->getSessionInput(session_, "forecast_precip");
        if (t_precip) net_->resizeTensor(t_precip, {1, N_ZONES, HORIZON_DAYS});

        auto* t_unc = net_->getSessionInput(session_, "forecast_uncertainty");
        if (t_unc) net_->resizeTensor(t_unc, {1, N_ZONES});

        auto* t_prior = net_->getSessionInput(session_, "prior_belief");
        if (t_prior) net_->resizeTensor(t_prior, {1, 1});

        auto* t_belief = net_->getSessionInput(session_, "zone_belief");
        if (t_belief) net_->resizeTensor(t_belief, {1, N_ZONES});

        // hidden_in shape: [1, 1, HIDDEN_SIZE]  (n_layers=1, batch=1, hidden_size)
        auto* t_hidden = net_->getSessionInput(session_, "hidden_in");
        if (t_hidden) net_->resizeTensor(t_hidden, {1, 1, HIDDEN_SIZE});

        net_->resizeSession(session_);

        // --- Cache input tensor pointers (avoid per-step lookup) ---
        in_forecast_precip_ = checkedGetInput("forecast_precip");
        in_forecast_unc_    = checkedGetInput("forecast_uncertainty");
        in_prior_belief_    = checkedGetInput("prior_belief");
        in_zone_belief_     = checkedGetInput("zone_belief");
        in_hidden_          = checkedGetInput("hidden_in");

        // --- Cache output tensor pointers ---
        // Output names from StatelessInferenceWrapper._export_onnx:
        //   output_names = ["action_logits", "hidden_out"]
        out_action_logits_ = checkedGetOutput("action_logits");
        out_hidden_        = checkedGetOutput("hidden_out");

        ready_ = in_basin_context_ && in_forecast_precip_ && in_forecast_unc_
              && in_prior_belief_ && in_zone_belief_ && in_hidden_
              && out_action_logits_ && out_hidden_;

        if (ready_) {
            std::cout << "[WeatherWrapper] Ready. Backend: "
                      << (usingVulkan_ ? "Vulkan" : "CPU")
                      << "  Model: " << model_path << std::endl;
        } else {
            std::cerr << "[WeatherWrapper] WARNING: one or more tensor names not found. "
                      << "Check tensor names against mnn_export.py." << std::endl;
        }
    }

    bool isReady() const { return ready_; }

    /**
     * Run one inference step.
     *
     * Inputs  (flat float arrays, caller-owned):
     *   forecast_precip     [N_ZONES * HORIZON_DAYS]
     *                       Row-major: all HORIZON_DAYS for zone 0, then zone 1, etc.
     *                       Matches forecast_precip[n_zones, horizon_days] in Python.
     *   forecast_uncertainty[N_ZONES]  one uncertainty value per zone (NOT per day)
     *   prior_belief        [1]
     *   zone_belief         [N_ZONES]
     *   hidden_in           [N_LAYERS * 1 * HIDDEN_SIZE]  zeros on episode start
     *   action_mask         [N_ACTIONS]  1.0f = valid, 0.0f = invalid
     *                       Terminate action (index N_ZONES) is always valid.
     *
     * Outputs (flat float arrays, caller-allocated):
     *   logits_out          [N_ACTIONS]  masked: invalid actions set to -1e9
     *   hidden_out          [N_LAYERS * 1 * HIDDEN_SIZE]  store for next step
     *
     * Returns: chosen action index (argmax over masked logits), or -1 on error.
     */
    int runStep(
        const float* basin_context,         // length BASIN_FLAT (schema v3)
        const float* forecast_precip,       // length N_ZONES * HORIZON_DAYS
        const float* forecast_uncertainty,  // length N_ZONES
        const float* prior_belief,          // length 1
        const float* zone_belief,           // length N_ZONES
        const float* hidden_in,             // length HIDDEN_FLAT
        const float* action_mask,           // length N_ACTIONS; 1=valid, 0=invalid
        float*       logits_out,            // length N_ACTIONS  (output)
        float*       hidden_out             // length HIDDEN_FLAT (output)
    ) {
        if (!ready_) {
            std::cerr << "[WeatherWrapper] runStep called on unready wrapper" << std::endl;
            return -1;
        }

        // --- Fill input tensors via host-side wrappers ---
        // Alphabetical order matches mnn_export.py export_keys sort.
        copyIn(in_basin_context_,   basin_context,        BASIN_FLAT);
        copyIn(in_forecast_precip_, forecast_precip,      PRECIP_FLAT);
        copyIn(in_forecast_unc_,    forecast_uncertainty, N_ZONES);
        copyIn(in_prior_belief_,    prior_belief,         1);
        copyIn(in_zone_belief_,     zone_belief,          N_ZONES);
        copyIn(in_hidden_,          hidden_in,            HIDDEN_FLAT);

        // --- Run ---
        if (net_->runSession(session_) != MNN::NO_ERROR) {
            std::cerr << "[WeatherWrapper] runSession failed" << std::endl;
            return -1;
        }

        // --- Read outputs ---
        copyOut(out_action_logits_, logits_out,  N_ACTIONS);
        copyOut(out_hidden_,        hidden_out,  HIDDEN_FLAT);

        // --- Apply action mask ---
        // Terminate action (index N_ZONES) is ALWAYS valid regardless of mask.
        // All other actions: masked if action_mask[i] <= 0.5.
        // Protocol matches mnn_export.py EDGE_INFERENCE_NOTE:
        //   logits[action_mask == 0] = -1e9; action = argmax(logits)
        int   best_action = -1;
        float best_logit  = -1e38f;
        for (int i = 0; i < N_ACTIONS; ++i) {
            bool valid = (i == N_ZONES) ? true : (action_mask[i] > 0.5f);
            if (!valid) {
                logits_out[i] = -1e9f;
            } else if (logits_out[i] > best_logit) {
                best_logit  = logits_out[i];
                best_action = i;
            }
        }

        return best_action;
    }

private:
    // --- Helpers ---

    MNN::Session* tryCreateSession(MNNForwardType type, const char* label) {
        MNN::ScheduleConfig config;
        config.type      = type;
        config.numThread = 2;

        MNN::BackendConfig backendConfig;
        backendConfig.precision = MNN::BackendConfig::Precision_Low;  // fp16 on GPU
        backendConfig.memory    = MNN::BackendConfig::Memory_Low;
        config.backendConfig    = &backendConfig;

        auto* s = net_->createSession(config);
        if (s) {
            std::cout << "[WeatherWrapper] Session created on " << label << std::endl;
            if (type == MNN_FORWARD_VULKAN) usingVulkan_ = true;
        }
        return s;
    }

    MNN::Tensor* checkedGetInput(const char* name) {
        auto* t = net_->getSessionInput(session_, name);
        if (!t) std::cerr << "[WeatherWrapper] WARNING: input tensor not found: " << name << std::endl;
        return t;
    }

    MNN::Tensor* checkedGetOutput(const char* name) {
        auto* t = net_->getSessionOutput(session_, name);
        if (!t) std::cerr << "[WeatherWrapper] WARNING: output tensor not found: " << name << std::endl;
        return t;
    }

    // Copy host float array → MNN tensor via a temporary host-layout wrapper.
    static void copyIn(MNN::Tensor* dst, const float* src, int n) {
        MNN::Tensor host(dst, MNN::Tensor::TENSORFLOW);
        std::memcpy(host.host<float>(), src, n * sizeof(float));
        dst->copyFromHostTensor(&host);
    }

    // Copy MNN tensor → host float array via a temporary host-layout wrapper.
    static void copyOut(MNN::Tensor* src, float* dst, int n) {
        MNN::Tensor host(src, MNN::Tensor::TENSORFLOW);
        src->copyToHostTensor(&host);
        std::memcpy(dst, host.host<float>(), n * sizeof(float));
    }

    // --- Members ---
    std::string                       model_path_;
    std::shared_ptr<MNN::Interpreter> net_;
    MNN::Session*                     session_     = nullptr;
    bool                              ready_       = false;
    bool                              usingVulkan_ = false;

    // Cached input tensors (alphabetical — matches export_keys sort order)
    MNN::Tensor* in_basin_context_   = nullptr;  // schema v3
    MNN::Tensor* in_forecast_precip_ = nullptr;
    MNN::Tensor* in_forecast_unc_    = nullptr;
    MNN::Tensor* in_prior_belief_    = nullptr;
    MNN::Tensor* in_zone_belief_     = nullptr;
    MNN::Tensor* in_hidden_          = nullptr;

    // Cached output tensors
    MNN::Tensor* out_action_logits_ = nullptr;
    MNN::Tensor* out_hidden_        = nullptr;
};


// ---------------------------------------------------------------------------
// C interface — for processing_nodes integration
// ---------------------------------------------------------------------------
extern "C" {

    /**
     * Create a wrapper instance.
     * Returns opaque handle, or nullptr on failure.
     */
    void* weather_create(const char* model_path) {
        auto* w = new WeatherEdgeWrapper(model_path);
        if (!w->isReady()) {
            delete w;
            return nullptr;
        }
        return w;
    }

    /**
     * Run one inference step.
     *
     * Bug 2.5 fix: null handle guard added. weather_create() returns nullptr
     * on failure; callers that omit the null check would previously trigger
     * undefined behaviour (segfault) here. Now returns -1 immediately.
     *
     * basin_context:       float[BASIN_FLAT]  [enso_oni, iod_dmi,
     *                      itcz_latitude_deg, mslp_regional_hpa] (schema v3).
     *                      Pass BASIN_NEUTRAL {0,0,0,1013.25} when unknown.
     * forecast_precip:     float[N_ZONES * HORIZON_DAYS], row-major
     *                      (all days for zone 0, then zone 1, etc.)
     * forecast_uncertainty:float[N_ZONES]   one value per zone
     * prior_belief:        float[1]
     * zone_belief:         float[N_ZONES]
     * hidden_in:           float[HIDDEN_SIZE]   zeros at episode start
     * action_mask:         float[N_ACTIONS]  1.0=valid, 0.0=masked
     *                      Terminate (index N_ZONES) always valid.
     * logits_out:          float[N_ACTIONS]  masked logits (out)
     * hidden_out:          float[HIDDEN_SIZE]  new hidden state (out)
     *
     * Returns chosen action index (0 .. N_ACTIONS-1), or -1 on error.
     */
    int weather_run(
        void*        handle,
        const float* basin_context,         // float[BASIN_FLAT] (schema v3)
        const float* forecast_precip,       // float[N_ZONES * HORIZON_DAYS]
        const float* forecast_uncertainty,  // float[N_ZONES]
        const float* prior_belief,          // float[1]
        const float* zone_belief,           // float[N_ZONES]
        const float* hidden_in,             // float[HIDDEN_SIZE]
        const float* action_mask,           // float[N_ACTIONS]
        float*       logits_out,            // float[N_ACTIONS]        (out)
        float*       hidden_out             // float[HIDDEN_SIZE]      (out)
    ) {
        // Bug 2.5 fix: guard against null handle (weather_create failed)
        if (!handle) {
            std::cerr << "[WeatherWrapper] weather_run called with null handle" << std::endl;
            return -1;
        }
        auto* w = static_cast<WeatherEdgeWrapper*>(handle);
        return w->runStep(basin_context,
                          forecast_precip, forecast_uncertainty,
                          prior_belief, zone_belief,
                          hidden_in, action_mask,
                          logits_out, hidden_out);
    }

    /**
     * Destroy a wrapper instance.
     *
     * Bug 2.5 fix: null handle guard added for symmetry with weather_run.
     * delete nullptr is safe in C++ but the explicit check makes the
     * contract visible and prevents a confusing double-free if a caller
     * passes nullptr after a failed weather_create().
     */
    void weather_destroy(void* handle) {
        if (!handle) return;
        delete static_cast<WeatherEdgeWrapper*>(handle);
    }

    /** Query shape constants — lets calling code stay in sync without hardcoding. */
    int weather_n_zones()       { return N_ZONES; }
    int weather_n_actions()     { return N_ACTIONS; }
    int weather_horizon_days()  { return HORIZON_DAYS; }
    int weather_hidden_size()   { return HIDDEN_SIZE; }
    int weather_precip_flat()   { return PRECIP_FLAT; }  // N_ZONES * HORIZON_DAYS
    int weather_basin_flat()    { return BASIN_FLAT; }   // 4 (schema v3)

} // extern "C"