#pragma once

#include <iostream>
#include <string>
#include <vector>

#include "KomputeModelML.hpp"

namespace godot {

KomputeModelML::KomputeModelML()
{
    std::cout << "CALLING CONSTRUCTOR" << std::endl;
    this->_init();
}

void
KomputeModelML::train(Array yArr, Array xIArr, Array xJArr)
{

    assert(yArr.size() == xIArr.size());
    assert(xIArr.size() == xJArr.size());

    std::vector<float> yData;
    std::vector<float> xIData;
    std::vector<float> xJData;
    std::vector<float> zerosData;

    for (size_t i = 0; i < yArr.size(); i++) {
        yData.push_back(yArr[i]);
        xIData.push_back(xIArr[i]);
        xJData.push_back(xJArr[i]);
        zerosData.push_back(0);
    }

    uint32_t ITERATIONS = 100;
    float learningRate = 0.1;

    {
        kp::Manager mgr;

        std::shared_ptr<kp::Tensor> xI = mgr.tensor(xIData);
        std::shared_ptr<kp::Tensor> xJ = mgr.tensor(xJData);

        std::shared_ptr<kp::Tensor> y = mgr.tensor(yData);

        std::shared_ptr<kp::Tensor> wIn = mgr.tensor({ 0.001, 0.001 });
        std::shared_ptr<kp::Tensor> wOutI = mgr.tensor(zerosData);
        std::shared_ptr<kp::Tensor> wOutJ = mgr.tensor(zerosData);

        std::shared_ptr<kp::Tensor> bIn = mgr.tensor({ 0 });
        std::shared_ptr<kp::Tensor> bOut = mgr.tensor(zerosData);

        std::shared_ptr<kp::Tensor> lOut = mgr.tensor(zerosData);

        std::vector<std::shared_ptr<kp::Tensor>> params = { xI,  xJ,    y,
                                                            wIn, wOutI, wOutJ,
                                                            bIn, bOut,  lOut };

        {
            std::vector<uint32_t> spirv(
              (uint32_t*)
                kp::shader_data::shaders_glsl_logisticregression_comp_spv,
              (uint32_t*)(kp::shader_data::
                            shaders_glsl_logisticregression_comp_spv +
                          kp::shader_data::
                            shaders_glsl_logisticregression_comp_spv_len));

            std::shared_ptr<kp::Algorithm> algo = mgr.algorithm(params, spirv);

            mgr.sequence()->eval<kp::OpTensorSyncDevice>(params);

            std::shared_ptr<kp::Sequence> sq =
              mgr.sequence()
                ->record<kp::OpTensorSyncDevice>({ wIn, bIn })
                ->record<kp::OpAlgoDispatch>(algo)
                ->record<kp::OpTensorSyncLocal>({ wOutI, wOutJ, bOut, lOut });

            // Iterate across all expected iterations
            for (size_t i = 0; i < ITERATIONS; i++) {

                sq->eval();

                for (size_t j = 0; j < bOut->size(); j++) {
                    wIn->data()[0] -= learningRate * wOutI->data()[j];
                    wIn->data()[1] -= learningRate * wOutJ->data()[j];
                    bIn->data()[0] -= learningRate * bOut->data()[j];
                }
            }
        }

        KP_LOG_INFO("RESULT: <<<<<<<<<<<<<<<<<<<");
        KP_LOG_INFO(wIn->data()[0]);
        KP_LOG_INFO(wIn->data()[1]);
        KP_LOG_INFO(bIn->data()[0]);

        this->mWeights = wIn;
        this->mBias = bIn;
    }
}

Array
KomputeModelML::predict(Array xI, Array xJ)
{
    assert(xI.size() == xJ.size());

    Array retArray;

    // We run the inference in the CPU for simplicity
    // BUt you can also implement the inference on GPU
    // GPU implementation would speed up minibatching
    for (size_t i = 0; i < xI.size(); i++) {
        float xIVal = xI[i];
        float xJVal = xJ[i];
        float result =
          (xIVal * this->mWeights->data()[0] +
           xJVal * this->mWeights->data()[1] + this->mBias->data()[0]);

        // Instead of using sigmoid we'll just return full numbers
        Variant var = result > 0 ? 1 : 0;
        retArray.push_back(var);
    }

    return retArray;
}

Array
KomputeModelML::get_params()
{
    Array retArray;

    KP_LOG_INFO(this->mWeights->size() + this->mBias->size());

    if (this->mWeights->size() + this->mBias->size() == 0) {
        return retArray;
    }

    retArray.push_back(this->mWeights->data()[0]);
    retArray.push_back(this->mWeights->data()[1]);
    retArray.push_back(this->mBias->data()[0]);
    retArray.push_back(99.0);

    return retArray;
}

void
KomputeModelML::_init()
{
    std::cout << "CALLING INIT" << std::endl;
}

void
KomputeModelML::_process(float delta)
{}

void
KomputeModelML::_register_methods()
{
    register_method((char*)"_process", &KomputeModelML::_process);
    register_method((char*)"_init", &KomputeModelML::_init);

    register_method((char*)"train", &KomputeModelML::train);
    register_method((char*)"predict", &KomputeModelML::predict);
    register_method((char*)"get_params", &KomputeModelML::get_params);
}

}
