#include <libasr/asr.h>
#include <libasr/containers.h>
#include <libasr/exception.h>
#include <libasr/asr_utils.h>
#include <libasr/asr_verify.h>
#include <libasr/pass/loop_vectorise.h>
#include <libasr/pass/pass_utils.h>

#include <cmath>


namespace LCompilers {

using ASR::down_cast;
using ASR::is_a;

/*

This ASR pass converts loops over arrays with a call to internal
optimization routine for loop vectorisation. This allows backend
specific code generation for vector operations resulting in optimized
code.

Converts:

    i: i32
    for i in range(9216):
        a[i] = b[i]

to:

    for i in range(0, 9216, 8):
        vector_copy(8, a[i*8:(i+1)*8], b[i*8:(i+1)*8])

*/
class LoopVectoriseVisitor : public PassUtils::SkipOptimizationFunctionVisitor<LoopVectoriseVisitor>
{
private:
    ASR::TranslationUnit_t &unit;

    std::string rl_path;

    // To make sure that FMA is applied only for
    // the nodes implemented in this class
    bool from_loop_vectorise;

    const int64_t avx512 = 512;

public:
    LoopVectoriseVisitor(Allocator &al_, ASR::TranslationUnit_t &unit_,
                         const std::string& rl_path_) : SkipOptimizationFunctionVisitor(al_),
    unit(unit_), rl_path(rl_path_), from_loop_vectorise(false)
    {
        pass_result.reserve(al, 1);
    }

    bool is_loop(const ASR::stmt_t& x) {
        return (ASR::is_a<ASR::DoLoop_t>(x) ||
                ASR::is_a<ASR::WhileLoop_t>(x));
    }

    bool is_vector_copy(ASR::stmt_t* x, Vec<ASR::expr_t*>& arrays) {
        if( !ASR::is_a<ASR::Assignment_t>(*x) ) {
            return false;
        }
        ASR::Assignment_t* x_assignment = ASR::down_cast<ASR::Assignment_t>(x);
        ASR::expr_t* target = x_assignment->m_target;
        ASR::expr_t* value = x_assignment->m_value;
        if( !ASR::is_a<ASR::ArrayItem_t>(*target) ||
            !ASR::is_a<ASR::ArrayItem_t>(*value) ) {
            return false;
        }
        ASR::ArrayItem_t* target_array_ref = ASR::down_cast<ASR::ArrayItem_t>(target);
        ASR::ArrayItem_t* value_array_ref = ASR::down_cast<ASR::ArrayItem_t>(value);
        if( target_array_ref->m_args->m_left ||
            target_array_ref->m_args->m_step ||
            value_array_ref->m_args->m_left ||
            value_array_ref->m_args->m_step ) {
            return false;
        }
        arrays.push_back(al, target_array_ref->m_v);
        arrays.push_back(al, value_array_ref->m_v);
        return true;
    }

    int64_t select_vector_instruction_size() {
        return avx512;
    }

    int64_t get_vector_length(ASR::ttype_t* type) {
        int kind = ASRUtils::extract_kind_from_ttype_t(type);
        int64_t instruction_length = select_vector_instruction_size();
        // TODO: Add check for divisibility and raise error if needed.
        return instruction_length/(kind * 8);
    }

    void get_vector_intrinsic(ASR::stmt_t* loop_stmt, ASR::expr_t* index,
                              ASR::expr_t*& vector_length,
                              Vec<ASR::stmt_t*>& vectorised_loop_body) {
        LCOMPILERS_ASSERT(vectorised_loop_body.reserve_called);
        Vec<ASR::expr_t*> arrays;
        arrays.reserve(al, 2);
        if( is_vector_copy(loop_stmt, arrays) ) {
            ASR::expr_t *target_sym = arrays[0], *value_sym = arrays[1];
            ASR::ttype_t* target_type = ASRUtils::expr_type(target_sym);
            int64_t vector_length_int = get_vector_length(target_type);
            vector_length = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, loop_stmt->base.loc,
                                vector_length_int, ASRUtils::expr_type(index)));
            ASR::expr_t* start = ASRUtils::EXPR(ASR::make_IntegerBinOp_t(al, loop_stmt->base.loc, index, ASR::binopType::Mul,
                                        vector_length, ASRUtils::expr_type(index), nullptr));
            ASR::expr_t* one = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, loop_stmt->base.loc,
                                    1, ASRUtils::expr_type(index)));
            ASR::expr_t* next_index = ASRUtils::EXPR(ASR::make_IntegerBinOp_t(al, loop_stmt->base.loc, index, ASR::binopType::Add,
                                            one, ASRUtils::expr_type(index), nullptr));
            ASR::expr_t* end = ASRUtils::EXPR(ASR::make_IntegerBinOp_t(al, loop_stmt->base.loc, next_index, ASR::binopType::Mul,
                                        vector_length, ASRUtils::expr_type(index), nullptr));
            vectorised_loop_body.push_back(al, PassUtils::get_vector_copy(target_sym, value_sym, start,
                end, one, vector_length,
                al, unit,
                current_scope, loop_stmt->base.loc));
            return ;
        }
    }

    void vectorise_loop(const ASR::DoLoop_t& x) {
        // Do Vectorisation of Loop inside this function
        ASR::expr_t* loop_start = x.m_head.m_start;
        ASR::expr_t* loop_end = x.m_head.m_end;
        ASR::expr_t* loop_inc = x.m_head.m_increment;
        LCOMPILERS_ASSERT(loop_start);
        LCOMPILERS_ASSERT(loop_end);
        if (!loop_inc) {
            int a_kind = ASRUtils::extract_kind_from_ttype_t(ASRUtils::expr_type(x.m_head.m_v));
            ASR::ttype_t *type = ASRUtils::TYPE(ASR::make_Integer_t(al, x.base.base.loc, a_kind));
            loop_inc = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, x.base.base.loc, 1, type));
        }
        ASR::expr_t* loop_start_value = ASRUtils::expr_value(loop_start);
        ASR::expr_t* loop_end_value = ASRUtils::expr_value(loop_end);
        ASR::expr_t* loop_inc_value = ASRUtils::expr_value(loop_inc);
        if( !ASRUtils::is_value_constant(loop_start_value) ||
            !ASRUtils::is_value_constant(loop_end_value) ||
            !ASRUtils::is_value_constant(loop_inc_value) ) {
            // Skip vectorisation of variable sized loops
            return ;
        }
        ASR::stmt_t* loop_stmt = x.m_body[0];
        int64_t loop_start_int = -1, loop_end_int = -1, loop_inc_int = -1;
        ASRUtils::extract_value(loop_start_value, loop_start_int);
        ASRUtils::extract_value(loop_end_value, loop_end_int);
        ASRUtils::extract_value(loop_inc_value, loop_inc_int);
        int64_t loop_size = (loop_end_int - loop_start_int) / loop_inc_int + 1;
        ASR::expr_t* vector_length = nullptr;
        Vec<ASR::stmt_t*> vectorised_loop_body;
        vectorised_loop_body.reserve(al, x.n_body);
        get_vector_intrinsic(loop_stmt, x.m_head.m_v,
                             vector_length, vectorised_loop_body);
        if( !vector_length ) {
            return ;
        }
        ASR::do_loop_head_t vectorised_loop_head;
        vectorised_loop_head.m_start = loop_start;
        int64_t vector_length_int = -1;
        ASRUtils::extract_value(vector_length, vector_length_int);
        // TODO: Add tests for loop sizes not divisible by vector_length.
        vectorised_loop_head.m_end = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, x.m_head.m_v->base.loc,
                                        loop_size/vector_length_int - 1, ASRUtils::expr_type(x.m_head.m_v)));
        // TODO: Update this for increments other 1.
        ASR::expr_t* one = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, loop_stmt->base.loc,
                                    1, ASRUtils::expr_type(x.m_head.m_v)));
        vectorised_loop_head.m_increment = one;
        vectorised_loop_head.m_v = x.m_head.m_v;
        vectorised_loop_head.loc = x.m_head.loc;

        ASR::stmt_t* vectorised_loop = ASRUtils::STMT(ASR::make_DoLoop_t(al, x.base.base.loc,
                                            x.m_name, vectorised_loop_head, vectorised_loop_body.p,
                                            vectorised_loop_body.size(), nullptr, 0));
        pass_result.push_back(al, vectorised_loop);
    }

    void visit_DoLoop(const ASR::DoLoop_t& x) {
        from_loop_vectorise = true;
        if( x.n_body == 1 ) {
            // Vectorise
            if( is_loop(*x.m_body[0]) ) {
                from_loop_vectorise = false;
                visit_stmt(*x.m_body[0]);
                return ;
            }
            vectorise_loop(x);
        }
        from_loop_vectorise = false;
    }

};

void pass_loop_vectorise(Allocator &al, ASR::TranslationUnit_t &unit,
                         const LCompilers::PassOptions& pass_options) {
    std::string rl_path = pass_options.runtime_library_dir;
    LoopVectoriseVisitor v(al, unit, rl_path);
    v.visit_TranslationUnit(unit);
    PassUtils::UpdateDependenciesVisitor u(al);
    u.visit_TranslationUnit(unit);
}


} // namespace LCompilers
