// SPDX-License-Identifier: MIT
// GERO probe invoking the actual nntrainer DivideLayer. Oracle lives in Python.
#include <array>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <vector>
#include <divide_layer.h>
#include <layer_context.h>
#include <var_grad.h>

static float fromBits(uint32_t u) { float v; std::memcpy(&v, &u, sizeof(v)); return v; }
static uint32_t bits(float v) { uint32_t u; std::memcpy(&u, &v, sizeof(u)); return u; }
int main() {
  unsigned n;
  if (!(std::cin >> n) || n == 0 || n % 64) return 2;
  std::vector<std::array<uint32_t,3>> inputs(n);
  for (auto &x: inputs) if (!(std::cin >> x[0] >> x[1] >> x[2])) return 3;
  const std::vector<nntrainer::TensorDim> shapes={
    nntrainer::TensorDim({1,1,1,n}), nntrainer::TensorDim({2,1,1,n/2}),
    nntrainer::TensorDim({4,2,8,n/64})};
  unsigned shape=0;
  for (const auto &dim: shapes) {
    nntrainer::DivideLayer layer;
    nntrainer::InitLayerContext init({dim,dim},{true},false,"divide_audit");
    layer.finalize(init);
    nntrainer::Var_Grad left(dim,nntrainer::Initializer::NONE,true,true,"left");
    nntrainer::Var_Grad right(dim,nntrainer::Initializer::NONE,true,true,"right");
    nntrainer::Var_Grad out(dim,nntrainer::Initializer::NONE,true,true,"out");
    nntrainer::RunLayerContext context("divide_audit",true,0.0f,false,1.0f,nullptr,false,{}, {&left,&right},{&out},{});
    for (unsigned i=0;i<n;++i) {
      left.getVariableRef().getData<float>()[i]=fromBits(inputs[i][0]);
      right.getVariableRef().getData<float>()[i]=fromBits(inputs[i][1]);
      out.getGradientRef().getData<float>()[i]=fromBits(inputs[i][2]);
    }
    layer.forwarding(context,true);
    for (unsigned repeat=0;repeat<2;++repeat) {
      layer.calcDerivative(context);
      for (unsigned i=0;i<n;++i) {
        bool unchanged=bits(left.getVariableRef().getData<float>()[i])==inputs[i][0] &&
          bits(right.getVariableRef().getData<float>()[i])==inputs[i][1] &&
          bits(out.getGradientRef().getData<float>()[i])==inputs[i][2];
        std::cout << shape << ' ' << repeat << ' ' << i << ' '
          << bits(out.getVariableRef().getData<float>()[i]) << ' '
          << bits(left.getGradientRef().getData<float>()[i]) << ' '
          << bits(right.getGradientRef().getData<float>()[i]) << ' ' << unchanged << '\n';
      }
    }
    ++shape;
  }
}