//Released under BSD-3 Cluse licsence
//How to run: root -b stackedsp.cpp
//or: c++ stackedsp.cpp -o stackedsp -O3 -lEtaler `root-config --cflags --ldflags --glibs` && ./stackedsp
//Assuming you have Etaler installed in /use/local/lib. Change this if you have it in other places
#pragma cling load("/usr/local/lib/libEtaler.so")
#include <Etaler/Etaler.hpp>
#include <Etaler/Algorithms/SpatialPooler.hpp>
#include <Etaler/Encoders/GridCell1d.hpp>
#include <Etaler/Encoders/Scalar.hpp>
#include <Etaler/Backends/OpenCLBackend.hpp>
using namespace et;

#include <random>

#include <xtensor/xarray.hpp>

constexpr intmax_t INPUT_SDR_SIZE = 2048;

#include <TGraph.h>
#include <TCanvas.h>
#include <TAxis.h>
#include <TH2.h>
#include <TStyle.h>

void run_experiment(int num_sps)
{
	//Make a stack of N SpatialPoolers
	std::vector<SpatialPooler> sps;
	for(int i=0;i<num_sps;i++) {
		SpatialPooler sp(Shape({INPUT_SDR_SIZE}), Shape({INPUT_SDR_SIZE}));
		sp.setBoostingFactor(10);
		sp.setGlobalDensity(0.04);
		sps.push_back(sp);
	}

	//Encode a scalar using a 1D GridCell encoder. A ScalarEncoder does the same job.
	auto encode = [](float v){return encoder::gridCell1d(v,INPUT_SDR_SIZE/16);};
	//Like an RBM, inference on each layer, train it and use it's outout as the
	//next layer's input
	auto train = [&](auto x){for(auto& sp : sps) {auto y = sp.compute(x); sp.learn(x, y); x = y;}};

	//Train the SP with random inputs
	std::mt19937 rng;
	std::uniform_real_distribution<float> dist;
	for(size_t i=0;i<1000;i++) {
		Tensor x = encode(dist(rng));
		train(x);
	}

	auto inference = [&](auto sdr) {Tensor x = sdr; for(auto& sp : sps) x = sp.compute(x); return x;};

	//Generate a graph of the on bits in the SDR
        const int num_steps = 128;
        auto c1 = std::make_unique<TCanvas>("c1", "canvas", 3000, 700);
        gPad->SetLeftMargin(0.06);
        gPad->SetRightMargin(0.01);
	auto t = std::make_unique<TH2F>("h2", ("plot of SDR generated by SP"+std::to_string(num_sps)).c_str(), INPUT_SDR_SIZE, 0, INPUT_SDR_SIZE, num_steps, 0, 1);

        auto bins = xt::linspace<float>(0.f, 1.f, num_steps);
        for(auto v : bins) {
                auto sdr = inference(encode(v)).toHost<uint8_t>();
                for(size_t i=0;i<sdr.size();i++)
                        t->Fill(i, v, sdr[i]+0.0001);
        }
        gStyle->SetOptStat(0);
        t->GetXaxis()->SetTitle("SDR");
        t->GetYaxis()->SetTitle("value sent into SP stack");
        t->Draw("col");
        c1->SaveAs(("plot_" + std::to_string(num_sps) + ".png").c_str());
	c1->Close();
}

void stackedsp_plot_sdr()
{
	//C++17 type reduction is awesome
        //setDefaultBackend(std::make_shared<OpenCLBackend>());

	std::vector num_sps = {0, 1, 2, 4, 8, 16};
	for(auto n : num_sps)
		run_experiment(n);
}

int main()
{
	stackedsp();
}