#pragma once
#include "algorithms/tebd.hpp"
#include "mps.hpp"
#include "models/transverse_field_ising.hpp"
#include "models/bose_hubbard.hpp"
#include "emscripten.hpp"
#include <cstdlib>
#include <iterator>
#include <limits>
#include <numeric>
#include <print>

using TfiModel = TransverseFieldIsingModel1D<dtype>;
using BhModel = BoseHubbardModel1D<dtype>;

template<typename Model>
class Results {
    public:
        virtual ~Results() {};
        virtual postUpdateFunction_t<dtype> getPostUpdateFunction(std::shared_ptr<Model> model) = 0;
        virtual bool postUpdateFunction(std::shared_ptr<Model> model, const MPS<dtype>& psi) = 0;
        virtual void finalize(std::shared_ptr<Model> model, const MPS<dtype>& psi) = 0;
        /// Norm of the MPS
        dtype finalNorm = 0;
};


class TfiResults : public Results<TfiModel> {
    public:
        using Model_t = TfiModel;
        TfiResults(ems::val callback, double convergence) : callback(callback), convergence(convergence), previousE(std::numeric_limits<double>::max()) {
            sigmaZ.setValues({{1, 0}, {0, -1}});
            sigmaX.setValues({{0, 1}, {1, 0}});
        };
        /// @todo write a second postUpdateFunction without the JS callback and return that one if callback is null
        postUpdateFunction_t<dtype> getPostUpdateFunction(std::shared_ptr<TfiModel> model) override {
            return [this, model](const MPS<dtype>& psi, unsigned i){ return this->postUpdateFunction(model, psi); };
        }
        bool postUpdateFunction(std::shared_ptr<TfiModel> model, const MPS<dtype>& psi) override {
            auto Mval = psi.getSiteExpectationValueSum(this->sigmaZ).real() / psi.getL();
            auto Eval = model->energy(psi).real();
            auto Sval = psi.getEntanglementEntropy(psi.getL()/2).real();
            auto chiVal = psi.getChiMax();
            callback(Mval, Eval, Sval, chiVal);
            // hand control back to browser, otherwise site hangs and callback will not be executed
            emscripten_sleep(0);
            if (std::abs(Eval - previousE) < convergence) return true;
            previousE = Eval;
            return false;
        }
        void finalize(std::shared_ptr<TfiModel> model, const MPS<dtype>& psi) override {
            this->finalNorm = psi.collapse(0);
            this->finalStatesZ = psi.getSiteExpectationValues(this->sigmaZ);
            this->finalStatesX = psi.getSiteExpectationValues(this->sigmaX);
        }
        std::vector<dtype> finalStatesZ;
        std::vector<dtype> finalStatesX;
    private:
        e::TensorFixedSize<dtype, e::Sizes<2, 2>> sigmaZ;
        e::TensorFixedSize<dtype, e::Sizes<2, 2>> sigmaX;
        ems::val callback;
        double convergence;
        double previousE;
};


/**
 * @brief Measures <N>, E and S 
 */
class BoseHubbardResults : public Results<BhModel> {
    public:
        BoseHubbardResults(ems::val callback, double convergence) : callback(callback), convergence(convergence), previousE(std::numeric_limits<double>::max()){};
        postUpdateFunction_t<dtype> getPostUpdateFunction(std::shared_ptr<BhModel> model) override {
            return [this, model](const MPS<dtype>& psi, unsigned) { return this->postUpdateFunction(model, psi); };
        }
        bool postUpdateFunction(std::shared_ptr<BhModel> model, const MPS<dtype>& psi) override {
            e::TensorMap<const e::Tensor<dtype, 2>> particleNumberOperator(model->getLadderOperators().getN().data(), model->localDim, model->localDim);
            auto Nval = psi.getSiteExpectationValueSum(particleNumberOperator).real() / model->L;
            auto Eval = model->energy(psi).real();
            auto Sval = psi.getEntanglementEntropy(model->L/2).real();
            auto chiVal = psi.getChiMax();
            callback(Nval, Eval, Sval, chiVal);
            // hand control back to browser, otherwise site hangs and callback will not be executed
            emscripten_sleep(0);
            if (std::abs(Eval - previousE) < convergence) return true;
            previousE = Eval;
            return false;
        }
        void finalize(std::shared_ptr<BhModel> model, const MPS<dtype>& psi) override {
            e::TensorMap<const e::Tensor<dtype, 2>> particleNumberOperator(model->getLadderOperators().getN().data(), model->localDim, model->localDim);
            this->finalNorm = psi.collapse(0);
            this->finalStatesN = psi.getSiteExpectationValues(particleNumberOperator);
        }
        std::vector<dtype> finalStatesN;
    private:
        ems::val callback;
        double convergence;
        double previousE;
};
