#include "gmsh_io.h"
#include "reference_element.h"

reference_element::reference_element()
{}

size_t
reference_element::num_basis_functions(void) const
{
    return m_num_bf;
}

size_t
reference_element::num_integration_points(void) const
{
    return m_quadpoints.size();
}

vecxd
reference_element::basis_functions(size_t iQp) const
{
    assert(iQp < num_integration_points());
    return m_basis_functions.segment(m_num_bf*iQp, m_num_bf);
}

matxd
reference_element::basis_gradients(size_t iQp) const
{
    assert(iQp < num_integration_points());
    return m_basis_gradients.block(m_num_bf * iQp, 0, m_num_bf, 3);
}

vecxd
reference_element::basis_functions(const point_3d& p) const
{
    std::vector<double> ip;
    ip.push_back(p.x());
    ip.push_back(p.y());
    ip.push_back(p.z());

    int bf_nc, bf_no;
    std::vector<double> basisFunctions;
    gmm::getBasisFunctions(m_gmsh_elem_type, ip, basis_func_name(m_approx_order), bf_nc,
        basisFunctions, bf_no);

    assert( basisFunctions.size() == m_num_bf*bf_no );
    auto bf_ofs = m_num_bf * m_orientation;

    vecxd ret = vecxd::Zero(m_num_bf);
    for (size_t i = 0; i < m_num_bf; i++)
        ret(i) = basisFunctions[bf_ofs + i];

    return ret;
}

matxd
reference_element::mass_matrix(void) const
{
    assert(m_geometric_order == 1);
    return m_mass_matrix;
}

matxd
reference_element::mass_matrix(size_t iQp) const
{
    assert(m_geometric_order > 1);
    return m_mass_matrices.block(m_num_bf*iQp, 0, m_num_bf, m_num_bf);
}

matxd
reference_element::mass_matrices(void) const
{
    assert(m_geometric_order > 1);
    return m_mass_matrices;
}

vec_quadpt_3d
reference_element::integration_points(void) const
{
    return m_quadpoints;
}













reference_elements_factory::reference_elements_factory(const entity_params& ep)
    : dim(ep.dim), tag(ep.tag), elemType(ep.etype), geom_order(ep.gorder),
      approx_order(ep.aorder)
{}

std::vector<reference_element>
reference_elements_factory::get_elements()
{
    /* Get integration points and weights on refelem */
    std::vector<double>     ips;
    std::vector<double>     iws;
    gmm::getIntegrationPoints(elemType, quadrature_name(2*approx_order), ips, iws);
    assert(ips.size() == 3*iws.size());

    auto num_qp = iws.size();

    int bf_nc, num_bfun_orient;
    std::vector<double> basisFunctions;
    gmm::getBasisFunctions(elemType, ips, basis_func_name(approx_order), bf_nc,
        basisFunctions, num_bfun_orient);
    assert(bf_nc == 1);
    
    auto num_bf = basisFunctions.size() / (num_qp * num_bfun_orient);
    assert( num_bf * num_qp * num_bfun_orient == basisFunctions.size() );

    int num_bgrad_orient;
    std::vector<double> basisGradients;
    gmm::getBasisFunctions(elemType, ips, basis_grad_name(approx_order), bf_nc,
        basisGradients, num_bgrad_orient);
    assert(bf_nc == 3);

    assert(num_bfun_orient == num_bgrad_orient);

    std::vector<reference_element> ret;

    for (size_t iO = 0; iO < num_bfun_orient; iO++)
    {
        /* Create one reference element per orientation */
        reference_element new_re;
        new_re.m_approx_order = approx_order;
        new_re.m_geometric_order = geom_order;
        new_re.m_dimension = dim;
        new_re.m_orientation = iO;
        new_re.m_parent_entity_tag = tag;
        new_re.m_basis_functions = vecxd::Zero(num_qp*num_bf);
        new_re.m_basis_gradients = matxd::Zero(num_qp*num_bf, 3);
        new_re.m_num_bf = num_bf;
        new_re.m_gmsh_elem_type = elemType;

        auto bf_base = num_bf * num_qp * iO;
        auto bg_base = 3 * bf_base;
        for (size_t iQp = 0; iQp < num_qp; iQp++)
        {
            for (size_t ibf = 0; ibf < num_bf; ibf++)
            {
                auto ofs = bf_base + num_bf*iQp + ibf;
                assert(ofs < basisFunctions.size());
                new_re.m_basis_functions(num_bf*iQp + ibf) =
                    basisFunctions[ofs];

                ofs = bg_base + 3*num_bf*iQp + 3*ibf + 0;
                assert(ofs < basisGradients.size());
                new_re.m_basis_gradients(num_bf*iQp + ibf, 0) =
                    basisGradients[ofs];
                
                ofs = bg_base + 3*num_bf*iQp + 3*ibf + 1;
                assert(ofs < basisGradients.size());
                new_re.m_basis_gradients(num_bf*iQp + ibf, 1) =
                    basisGradients[ofs];

                ofs = bg_base + 3*num_bf*iQp + 3*ibf + 2;
                assert(ofs < basisGradients.size());
                new_re.m_basis_gradients(num_bf*iQp + ibf, 2) =
                    basisGradients[ofs];
            }
        }

        for (size_t iQp = 0; iQp < num_qp; iQp++)
        {
            auto weight = iws[iQp];
            auto qpt_u = ips[3*iQp+0];
            auto qpt_v = ips[3*iQp+1];
            auto qpt_w = ips[3*iQp+2];
            auto pt = point_3d(qpt_u, qpt_v, qpt_w);
            new_re.m_quadpoints.push_back( quadrature_point_3d(pt, weight) );       
        }

        if (new_re.m_geometric_order == 1)
        {
            new_re.m_mass_matrix = matxd::Zero(num_bf, num_bf);
            for (size_t iQp = 0; iQp < num_qp; iQp++)
            {
                auto w = new_re.m_quadpoints[iQp].weight();
                vecxd phi = new_re.basis_functions(iQp);
                new_re.m_mass_matrix += w * phi * phi.transpose();
            }
        }
        else
        {
            new_re.m_mass_matrices = matxd::Zero(num_bf*num_qp, num_bf);
            for (size_t iQp = 0; iQp < num_qp; iQp++)
            {
                auto w = new_re.m_quadpoints[iQp].weight();
                vecxd phi = new_re.m_basis_functions.segment(num_bf*iQp, num_bf);
                new_re.m_mass_matrices.block(num_bf*iQp, 0, num_bf, num_bf) =
                    w * phi * phi.transpose();
            }
        }

        ret.push_back(std::move(new_re));
    }

    return ret;
}