#include <iostream>
#include <iomanip>

#include <unistd.h>

#include "test.h"
#include "gmsh_io.h"
#include "sgr.hpp"
#include "entity_data.h"
#include "kernels_cpu.h"

using namespace sgr;

static void
make_geometry(int order, double mesh_h)
{
    gm::add("difftest");
    gmo::addBox(0.0, 0.0, 0.0, 1.0, 1.0, 1.0);
    gmo::synchronize();

    gvp_t vp;
    gm::getEntities(vp);
    gmm::setSize(vp, mesh_h);
}

int test_differentiation_convergence(int geometric_order, int approximation_order)
{
    std::vector<double> sizes({ 0.32, 0.16, 0.08, 0.04 });
    std::vector<double> errors;

    std::cout << cyanfg << "Testing geometric order " << geometric_order;
    std::cout << ", approximation order = " << approximation_order << nofg;
    std::cout << std::endl;

    auto f = [](const point_3d& pt) -> double {
        return std::sin(M_PI*pt.x())*std::sin(M_PI*pt.y())*std::sin(M_PI*pt.z()); 
    };
    auto df_dx = [](const point_3d& pt) -> double {
        return M_PI*std::cos(M_PI*pt.x())*std::sin(M_PI*pt.y())*std::sin(M_PI*pt.z()); 
    };
    auto df_dy = [](const point_3d& pt) -> double {
        return M_PI*std::sin(M_PI*pt.x())*std::cos(M_PI*pt.y())*std::sin(M_PI*pt.z()); 
    };
    auto df_dz = [](const point_3d& pt) -> double {
        return M_PI*std::sin(M_PI*pt.x())*std::sin(M_PI*pt.y())*std::cos(M_PI*pt.z()); 
    };

    std::vector<double> errs_x;
    std::vector<double> errs_y;
    std::vector<double> errs_z;

    for (size_t i = 0; i < sizes.size(); i++)
    {
        double h = sizes[i];
        make_geometry(0,h);
        model mod(geometric_order, approximation_order);

        auto& e0 = mod.entity_at(0);
        e0.sort_by_orientation();
        vecxd Pf_e0 = e0.project(f);
        vecxd df_dx_e0 = vecxd::Zero( Pf_e0.rows() );
        vecxd df_dy_e0 = vecxd::Zero( Pf_e0.rows() );
        vecxd df_dz_e0 = vecxd::Zero( Pf_e0.rows() );

        entity_data_cpu<1> ed;
        e0.populate_entity_data(ed);

        compute_field_derivatives(ed, Pf_e0, df_dx_e0, df_dy_e0, df_dz_e0);

        vecxd Pdf_dx_e0 = e0.project(df_dx);
        vecxd Pdf_dy_e0 = e0.project(df_dy);
        vecxd Pdf_dz_e0 = e0.project(df_dz);

        double err_x = 0.0;
        double err_y = 0.0;
        double err_z = 0.0;

        for (size_t iT = 0; iT < e0.num_elements(); iT++)
        {
            auto& pe = e0.elem(iT);
            auto& re = e0.refelem(pe);
            matxd mass = e0.mass_matrix(iT);
            auto num_bf = re.num_basis_functions();
            vecxd diff_x = Pdf_dx_e0.segment(iT*num_bf, num_bf) - df_dx_e0.segment(iT*num_bf, num_bf); 
            vecxd diff_y = Pdf_dy_e0.segment(iT*num_bf, num_bf) - df_dy_e0.segment(iT*num_bf, num_bf); 
            vecxd diff_z = Pdf_dz_e0.segment(iT*num_bf, num_bf) - df_dz_e0.segment(iT*num_bf, num_bf); 
        
            err_x += diff_x.dot(mass*diff_x);
            err_y += diff_y.dot(mass*diff_y);
            err_z += diff_z.dot(mass*diff_z);
        }

        std::cout << std::sqrt(err_x) << " " << std::sqrt(err_y) << " " << std::sqrt(err_z) << std::endl;
    
        errs_x.push_back( std::sqrt(err_x) );
        errs_y.push_back( std::sqrt(err_y) );
        errs_z.push_back( std::sqrt(err_z) );
    }

    double rate_x = 0.0;
    double rate_y = 0.0;
    double rate_z = 0.0;

    std::cout << Byellowfg << "rate df/dx  rate df/dy  rate df/dz" << reset << std::endl;
    for (size_t i = 1; i < sizes.size(); i++)
    {
        std::cout << (rate_x = std::log2(errs_x[i-1]/errs_x[i]) ) << " ";
        std::cout << (rate_y = std::log2(errs_y[i-1]/errs_y[i]) ) << " ";
        std::cout << (rate_z = std::log2(errs_z[i-1]/errs_z[i]) ) << std::endl;
    }

    COMPARE_VALUES_ABSOLUTE("df/dx", rate_x, double(approximation_order), 0.2);
    COMPARE_VALUES_ABSOLUTE("df/dy", rate_y, double(approximation_order), 0.2);
    COMPARE_VALUES_ABSOLUTE("df/dz", rate_z, double(approximation_order), 0.2);

    return 0;
}