#include "NaiveIRGenerator.h"
#include <omp.h>


using mnd::NaiveIRGenerator;


namespace mnd
{
    template class NaiveIRGenerator<float>;
    template class NaiveIRGenerator<double>;
    template class NaiveIRGenerator<mnd::DoubleDouble>;
    template class NaiveIRGenerator<mnd::QuadDouble>;
}

namespace mnd::eval
{


    struct Load { size_t index; };
    struct Store
    {
        size_t index;
        std::unique_ptr<EvalNode> v;
    };

    struct BinaryOperation
    {
        std::unique_ptr<EvalNode> a;
        std::unique_ptr<EvalNode> b;
    };

    struct UnaryOperation
    {
        std::unique_ptr<EvalNode> a;
    };

    struct Add : BinaryOperation {};
    struct Sub : BinaryOperation {};
    struct Mul : BinaryOperation {};
    struct Div : BinaryOperation {};

    struct Neg : UnaryOperation {};


    struct Atan2 : BinaryOperation {};
    struct Pow : BinaryOperation {};
    struct Cos : UnaryOperation {};
    struct Sin : UnaryOperation {};
    struct Exp : UnaryOperation {};
    struct Ln : UnaryOperation {};


    using namespace mnd;
    using namespace mnd::ir;
    template<typename T>
    struct ToEvalVisitor
    {
        EvalStruct<T>& es;

        std::unique_ptr<EvalNode> visit(ir::Node* f)
        {
            std::any& nodeData = getNodeData(f);
            if (EvalNode** en = std::any_cast<EvalNode*>(&nodeData)) {
                
                size_t tmpStore;
                if (Store* s = std::get_if<Store>(*en)) {
                    tmpStore = s->index;
                }
                else {
                    tmpStore = createTemp();
                    EvalNode store = Store{ tmpStore, std::make_unique<EvalNode>(std::move(**en)) };
                    **en = std::move(store);
                }
                auto l = std::make_unique<EvalNode>(Load{ tmpStore });
                setNodeData(f, l.get());
                return l;
            }
            EvalNode n = std::visit(*this, *f);
            auto r = std::make_unique<EvalNode>(std::move(n));
            setNodeData(f, r.get());
            return r;
        }

        std::any& getNodeData(ir::Node* n)
        {
            return std::visit([](auto& x) -> std::any& { return x.nodeData; }, *n);
        }

        void setNodeData(ir::Node* n, EvalNode* en)
        {
            std::visit([en](auto& x) { x.nodeData = en; }, *n);
        }

        size_t createTemp(void)
        {
            es.variables.push_back(0);
            return es.variables.size() - 1;
        }

        size_t createConstant(mnd::Real& value)
        {
            es.variables.push_back(mnd::convert<T>(value));
            return es.variables.size() - 1;
        }

        size_t createVariable(std::string& value)
        {
            es.variables.push_back(0);
            es.variableNames.emplace(value, es.variables.size() - 1);
            return es.variables.size() - 1;
        }

        EvalNode operator()(ir::Constant& x) {
            return Load{ createConstant(x.value) };
        }
        EvalNode operator()(ir::Variable& x) {
            return Load{ createVariable(x.name) };
        }
        EvalNode operator()(ir::Addition& x) {
            return Add{ visit(x.left), visit(x.right) };
        }
        EvalNode operator()(ir::Subtraction& x) {
            return Sub{ visit(x.left), visit(x.right) };
        }
        EvalNode operator()(ir::Multiplication& x) {
            return Mul{ visit(x.left), visit(x.right) };
        }
        EvalNode operator()(ir::Division& x) {
            return Div{ visit(x.left), visit(x.right) };
        }
        EvalNode operator()(ir::Negation& x) {
            return Neg{ visit(x.value) };
        }
        EvalNode operator()(ir::Atan2& x) {
            return Atan2{ visit(x.left), visit(x.right) };
        }
        EvalNode operator()(ir::Pow& x) {
            return Pow{ visit(x.left), visit(x.right) };
        }
        EvalNode operator()(ir::Cos& x) {
            return Cos{ visit(x.value) };
        }
        EvalNode operator()(ir::Sin& x) {
            return Sin{ visit(x.value) };
        }
        EvalNode operator()(ir::Exp& x) {
            return Exp{ visit(x.value) };
        }
        EvalNode operator()(ir::Ln& x) {
            return Ln{ visit(x.value) };
        }
    };



    template<typename T>
    struct EvalVisitor
    {
        mnd::eval::EvalStruct<T>& es;

        T visit(const EvalNode& en) {
            return std::visit(*this, en);
        }

        T operator()(const Load& x) {
            return es.variables[x.index];
        }

        T operator()(const Store& x) {
            T r = visit(*x.v);
            es.variables[x.index] = r;
            return r;
        }

        T operator()(const Add& x) {
            return visit(*x.a) + visit(*x.b);
        }

        T operator()(const Sub& x) {
            return visit(*x.a) - visit(*x.b);
        }

        T operator()(const Mul& x) {
            return visit(*x.a) * visit(*x.b);
        }

        T operator()(const Div& x) {
            return visit(*x.a) / visit(*x.b);
        }

        T operator()(const Neg& x) {
            return -visit(*x.a);
        }

        T operator()(const Atan2& x) {
            return mnd::atan2(visit(*x.a), visit(*x.b));
        }

        T operator()(const Pow& x) {
            return mnd::pow(visit(*x.a), visit(*x.b));
        }

        T operator()(const Cos& x) {
            return mnd::cos(visit(*x.a));
        }

        T operator()(const Sin& x) {
            return mnd::sin(visit(*x.a));
        }

        T operator()(const Exp& x) {
            return mnd::exp(visit(*x.a));
        }

        T operator()(const Ln& x) {
            return mnd::log(visit(*x.a));
        }
    };
}


template<typename T>
NaiveIRGenerator<T>::NaiveIRGenerator(const mnd::ir::Formula& irf,
    mnd::Precision prec) :
    mnd::MandelGenerator{ prec },
    form{ irf }
{
    eval::ToEvalVisitor<T> tev{ es };
    newz_re = tev.visit(irf.newA);
    newz_im = tev.visit(irf.newB);
    start_re = tev.visit(irf.startA);
    start_im = tev.visit(irf.startB);
}


template<typename T>
void NaiveIRGenerator<T>::generate(const mnd::MandelInfo& info, float* data)
{
    const MandelViewport& view = info.view;

    const bool parallel = true;
    T viewx = mnd::convert<T>(view.x);
    T viewy = mnd::convert<T>(view.y);
    T wpp = mnd::convert<T>(view.width / info.bWidth);
    T hpp = mnd::convert<T>(view.height / info.bHeight);

#if defined(_OPENMP)
    if constexpr (parallel)
        omp_set_num_threads(omp_get_num_procs());
#   pragma omp parallel for schedule(static, 1) if (parallel)
#endif
    for (long j = 0; j < info.bHeight; j++) {
        T y = viewy + T(double(j)) * hpp;
        long i = 0;
        for (i; i < info.bWidth; i++) {
            T x = viewx + T(double(i)) * wpp;

            es.prepare(0, 0, x, y);

            eval::EvalVisitor<T> visitor{ es };

            T a = visitor.visit(*start_re);
            T b = visitor.visit(*start_im);

            es.prepare(a, b, x, y);

            int k = 0;
            for (k = 0; k < info.maxIter; k++) {
                T newA = visitor.visit(*newz_re);
                T newB = visitor.visit(*newz_im);
                a = newA;
                b = newB;
                if (a * a + b * b >= 16.0)
                    break;
            }
            data[i + j * info.bWidth] = float(k);
        }
    }
}