#include "maxwell_interface.h"

namespace maxwell {

void
field::zero()
{
    Ex = vecxd::Zero(num_dofs);
    Ey = vecxd::Zero(num_dofs);
    Ez = vecxd::Zero(num_dofs);
    Hx = vecxd::Zero(num_dofs);
    Hy = vecxd::Zero(num_dofs);
    Hz = vecxd::Zero(num_dofs);
}

void
field::resize(size_t p_num_dofs)
{
    Ex = vecxd::Zero(p_num_dofs);
    Ey = vecxd::Zero(p_num_dofs);
    Ez = vecxd::Zero(p_num_dofs);
    Hx = vecxd::Zero(p_num_dofs);
    Hy = vecxd::Zero(p_num_dofs);
    Hz = vecxd::Zero(p_num_dofs);
    num_dofs = p_num_dofs;
}

#ifdef ENABLE_GPU_SOLVER

pinned_field::pinned_field()
    : Ex(nullptr), Ey(nullptr), Ez(nullptr),
      Hx(nullptr), Hy(nullptr), Hz(nullptr),
      num_dofs(0)
{}

void
pinned_field::zero()
{
    memset(Ex, '\0', num_dofs*sizeof(double));
    memset(Ey, '\0', num_dofs*sizeof(double));
    memset(Ez, '\0', num_dofs*sizeof(double));
    memset(Hx, '\0', num_dofs*sizeof(double));
    memset(Hy, '\0', num_dofs*sizeof(double));
    memset(Hz, '\0', num_dofs*sizeof(double));
}

void
pinned_field::resize(size_t sz)
{
    if (Ex and sz != num_dofs) checkGPU( gpuFreeHost(Ex) );
    if (Ey and sz != num_dofs) checkGPU( gpuFreeHost(Ey) );
    if (Ez and sz != num_dofs) checkGPU( gpuFreeHost(Ez) );
    if (Hx and sz != num_dofs) checkGPU( gpuFreeHost(Hx) );
    if (Hy and sz != num_dofs) checkGPU( gpuFreeHost(Hy) );
    if (Hz and sz != num_dofs) checkGPU( gpuFreeHost(Hz) );
    
    zero();

    checkGPU( gpuMallocHost((void**)&Ex, sz*sizeof(double)) );
    checkGPU( gpuMallocHost((void**)&Ey, sz*sizeof(double)) );
    checkGPU( gpuMallocHost((void**)&Ez, sz*sizeof(double)) );
    checkGPU( gpuMallocHost((void**)&Hx, sz*sizeof(double)) );
    checkGPU( gpuMallocHost((void**)&Hy, sz*sizeof(double)) );
    checkGPU( gpuMallocHost((void**)&Hz, sz*sizeof(double)) );
    num_dofs = sz;
}

pinned_field::~pinned_field()
{
    if (Ex) checkGPU( gpuFreeHost(Ex) );
    if (Ey) checkGPU( gpuFreeHost(Ey) );
    if (Ez) checkGPU( gpuFreeHost(Ez) );
    if (Hx) checkGPU( gpuFreeHost(Hx) );
    if (Hy) checkGPU( gpuFreeHost(Hy) );
    if (Hz) checkGPU( gpuFreeHost(Hz) );
}

void
field_gpu::zero()
{
    Ex.zero();
    Ey.zero();
    Ez.zero();
    Hx.zero();
    Hy.zero();
    Hz.zero();
}
void
field_gpu::resize(size_t p_num_dofs)
{
    Ex.resize(p_num_dofs);
    Ey.resize(p_num_dofs);
    Ez.resize(p_num_dofs);
    Hx.resize(p_num_dofs);
    Hy.resize(p_num_dofs);
    Hz.resize(p_num_dofs);
    zero();
    num_dofs = p_num_dofs;
}

field_gpu::raw_ptrs
field_gpu::data(void)
{
    raw_ptrs ret;

    ret.num_dofs = num_dofs;
    ret.Ex = Ex.data();
    ret.Ey = Ey.data();
    ret.Ez = Ez.data();
    ret.Hx = Hx.data();
    ret.Hy = Hy.data();
    ret.Hz = Hz.data();

    return ret;
}

field_gpu::const_raw_ptrs
field_gpu::data(void) const
{
    const_raw_ptrs ret;

    ret.num_dofs = num_dofs;
    ret.Ex = Ex.data();
    ret.Ey = Ey.data();
    ret.Ez = Ez.data();
    ret.Hx = Hx.data();
    ret.Hy = Hy.data();
    ret.Hz = Hz.data();

    return ret;
}

void
field_gpu::copyin(const field& emf)
{
    Ex.copyin(emf.Ex.data(), emf.Ex.size());
    Ey.copyin(emf.Ey.data(), emf.Ey.size());
    Ez.copyin(emf.Ez.data(), emf.Ez.size());
    Hx.copyin(emf.Hx.data(), emf.Hx.size());
    Hy.copyin(emf.Hy.data(), emf.Hy.size());
    Hz.copyin(emf.Hz.data(), emf.Hz.size());
}

void
field_gpu::copyin(const field& emf, const stream& st)
{
    Ex.copyin(emf.Ex.data(), emf.Ex.size(), st);
    Ey.copyin(emf.Ey.data(), emf.Ey.size(), st);
    Ez.copyin(emf.Ez.data(), emf.Ez.size(), st);
    Hx.copyin(emf.Hx.data(), emf.Hx.size(), st);
    Hy.copyin(emf.Hy.data(), emf.Hy.size(), st);
    Hz.copyin(emf.Hz.data(), emf.Hz.size(), st);
}

void
field_gpu::copyin(const pinned_field& emf, const stream& st)
{
    Ex.copyin(emf.Ex, emf.num_dofs, st);
    Ey.copyin(emf.Ey, emf.num_dofs, st);
    Ez.copyin(emf.Ez, emf.num_dofs, st);
    Hx.copyin(emf.Hx, emf.num_dofs, st);
    Hy.copyin(emf.Hy, emf.num_dofs, st);
    Hz.copyin(emf.Hz, emf.num_dofs, st);
}

void
field_gpu::copyout(field& emf) const
{
    if (num_dofs != emf.num_dofs)
        emf.resize(num_dofs);

    Ex.copyout(emf.Ex.data());
    Ey.copyout(emf.Ey.data());
    Ez.copyout(emf.Ez.data());
    Hx.copyout(emf.Hx.data());
    Hy.copyout(emf.Hy.data());
    Hz.copyout(emf.Hz.data());
}



material_params_gpu::raw_ptrs
material_params_gpu::data(void)
{
    material_params_gpu::raw_ptrs ret;
    ret.num_dofs            = num_dofs;
    ret.num_fluxes          = num_fluxes;
    ret.inv_epsilon         = inv_epsilon.data();
    ret.inv_mu              = inv_mu.data();
    ret.sigma               = sigma.data();
    ret.sigma_over_epsilon  = sigma_over_epsilon.data();
    ret.aE                  = aE.data();
    ret.bE                  = bE.data();
    ret.aH                  = aH.data();
    ret.bH                  = bH.data();
    ret.bc_coeffs           = bc_coeffs.data();
    return ret;
}

material_params_gpu::const_raw_ptrs
material_params_gpu::data(void) const
{
    material_params_gpu::const_raw_ptrs ret;
    ret.num_dofs            = num_dofs;
    ret.num_fluxes          = num_fluxes;
    ret.inv_epsilon         = inv_epsilon.data();
    ret.inv_mu              = inv_mu.data();
    ret.sigma               = sigma.data();
    ret.sigma_over_epsilon  = sigma_over_epsilon.data();
    ret.aE                  = aE.data();
    ret.bE                  = bE.data();
    ret.aH                  = aH.data();
    ret.bH                  = bH.data();
    ret.bc_coeffs           = bc_coeffs.data();
    return ret;
}

void
material_params_gpu::copyin(const material_params& mp)
{
    num_dofs    = mp.num_dofs;
    num_fluxes  = mp.num_fluxes;

    inv_epsilon.copyin(mp.inv_epsilon.data(), mp.inv_epsilon.size());
    inv_mu.copyin(mp.inv_mu.data(), mp.inv_mu.size());
    sigma.copyin(mp.sigma.data(), mp.sigma.size());
    sigma_over_epsilon.copyin(mp.sigma_over_epsilon.data(), mp.sigma_over_epsilon.size());
    aE.copyin(mp.aE.data(), mp.aE.size());
    bE.copyin(mp.bE.data(), mp.bE.size());
    aH.copyin(mp.aH.data(), mp.aH.size());
    bH.copyin(mp.bH.data(), mp.bH.size());
    bc_coeffs.copyin(mp.bc_coeffs.data(), mp.bc_coeffs.size());
}

#endif /* ENABLE_GPU_SOLVER */





#define SEC_BCONDS              "bndconds"
#define BC_KIND                 "kind"
#define BC_SOURCE               "source"
#define BCOND_TYPE_PEC          "pec"
#define BCOND_TYPE_PMC          "pmc"
#define BCOND_TYPE_IMPEDANCE    "impedance"
#define BCOND_TYPE_EPLANEW      "plane_wave_E"
#define BCOND_TYPE_HPLANEW      "plane_wave_H"
#define BCOND_TYPE_EFIELD       "E_field"
#define BCOND_TYPE_HFIELD       "H_field"
#define BCOND_TYPE_SURFCURR     "surface_current"

#define SEC_IFCONDS             "ifaceconds"
#define IFC_KIND                "kind"
#define IFC_SOURCE              "source"
#define IFCOND_TYPE_EFIELD      "E_field"
#define IFCOND_TYPE_HFIELD      "H_field"
#define IFCOND_TYPE_SURFCURR    "surface_current"


#define SEC_MATERIALS           "materials"
#define MAT_EPS                 "epsilon"
#define MAT_EPS0                "eps0"
#define MAT_MU                  "mu"
#define MAT_MU0                 "mu0"
#define MAT_SIGMA               "sigma"



parameter_loader::parameter_loader()
{
    lua["const"][MAT_EPS0] = 8.8541878128e-12;
    lua["const"][MAT_MU0]  = 4e-7*M_PI;
    lua["postpro"]["E"] = lua.create_table();
    lua["postpro"]["E"]["silo_mode"] = "nodal";
    lua["postpro"]["H"] = lua.create_table();
    lua["postpro"]["H"]["silo_mode"] = "nodal";
    lua["postpro"]["J"] = lua.create_table();
    lua["postpro"]["J"]["silo_mode"] = "nodal";
}

bool
parameter_loader::validate_materials(const std::string& mpn, int tag) const
{
    auto mfun = lua[SEC_MATERIALS][mpn];
    if ( mfun.valid() )
        return true;

    auto mparams = lua[SEC_MATERIALS][tag];
    if ( mparams.valid() )
    {
        auto matparams_mpn = lua[SEC_MATERIALS][tag][mpn];
        if ( matparams_mpn.valid() )
            return true;
    }

    std::cout << "[CONFIG] '" SEC_MATERIALS ".[" << tag << "]." << mpn;
    std::cout << "' not defined and '" SEC_MATERIALS "." << mpn;
    std::cout << "(tag,x,y,z)' not present." << std::endl;
    return false;
}

bool
parameter_loader::validate_boundary_condition(const model&, int tag) const
{
    auto bc = lua[SEC_BCONDS][tag];
    if (not bc.valid())
        return true;

    if ( not bc[BC_KIND].valid() )
    {
        std::cout << "[CONFIG] '" BC_KIND "' not specified on interface ";
        std::cout << tag << std::endl;
        return false;
    }

    std::string bc_kind = bc[BC_KIND];

    if (bc_kind == BCOND_TYPE_PEC)
        return true;

    if (bc_kind == BCOND_TYPE_PMC)
        return true;

    if (bc_kind == BCOND_TYPE_IMPEDANCE)
        return true;

    if (bc_kind == BCOND_TYPE_EPLANEW)
    {
        if (bc[BC_SOURCE].valid())
            return true;
        
        std::cout << "[CONFIG] '" BC_SOURCE "' not specified for plane ";
        std::cout << "wave condition on surface " << tag << std::endl;
        return false;
    }

    std::cout << "[CONFIG] boundary condition not implemented on ";
    std::cout << "surface " << tag << std::endl;
    return false;
}

bool
parameter_loader::validate_interface_condition(const model&, int tag) const
{
    auto ic = lua[SEC_IFCONDS][tag];
    if (not ic.valid())
        return true;

    if ( not ic[IFC_KIND].valid() )
    {
        std::cout << "[CONFIG '" IFC_KIND "' not specified on interface ";
        std::cout << tag << std::endl;
        return false;
    }

    std::string ic_kind = ic[IFC_KIND];

    if (ic_kind == IFCOND_TYPE_SURFCURR)
    {
        if (ic[IFC_SOURCE].valid())
            return true;
        
        std::cout << "[CONFIG] '" IFC_SOURCE "' not specified for surface ";
        std::cout << "current condition on surface " << tag << std::endl;
        return false;
    }

    std::cout << "[CONFIG] interface condition not implemented on ";
    std::cout << "surface " << tag << std::endl;
    return false;
}

bool
parameter_loader::validate_conditions(const model& mod) const
{
    bool success = true;
    auto bds = mod.boundary_descriptors();
    std::sort(bds.begin(), bds.end());
    bds.erase( std::unique(bds.begin(), bds.end()), bds.end() );
    
    auto ft_none = [](const boundary_descriptor& bd) -> bool {
        return bd.type == face_type::NONE;
    };

    bds.erase( std::remove_if(bds.begin(), bds.end(), ft_none), bds.end() );

    for (auto& bd : bds)
    {
        switch (bd.type)
        {
            case face_type::NONE:
                assert(false);
                break;

            case face_type::BOUNDARY:
                if (not validate_boundary_condition(mod, bd.gmsh_entity) )
                    success = false;
                break;

            case face_type::INTERFACE:
                if (not validate_interface_condition(mod, bd.gmsh_entity) )
                    success = false;
                break;
        }
    }

    return success;
}

bool
parameter_loader::validate_model_params(const model& mod) const
{
    bool success = true;
    for (auto& e : mod)
    {
        auto tag = e.gmsh_tag();
        bool eps_valid = validate_materials(MAT_EPS, tag);
        bool mu_valid = validate_materials(MAT_MU, tag);
        bool sigma_valid = validate_materials(MAT_SIGMA, tag);
        if ( (not eps_valid) or (not mu_valid) or (not sigma_valid) )
            success = false;
    }

    if ( not validate_conditions(mod) )
        success = false;

    return success;
}

bool
parameter_loader::initial_Efield_defined(void) const
{
    auto eic = lua["electric_initial_condition"];
    return eic.valid();
}

vec3d
parameter_loader::initial_Efield(const point_3d& pt) const
{
    vec3d ret;
    sol::tie(ret(0), ret(1), ret(2)) =
        lua["electric_initial_condition"](pt.x(), pt.y(), pt.z());
    return ret;
}

bool
parameter_loader::initial_Hfield_defined(void) const
{
    auto mic = lua["magnetic_initial_condition"];
    return mic.valid();
}

vec3d
parameter_loader::initial_Hfield(const point_3d& pt) const
{
    vec3d ret;
    sol::tie(ret(0), ret(1), ret(2)) =
        lua["magnetic_initial_condition"](pt.x(), pt.y(), pt.z());
    return ret;
}

double
parameter_loader::epsilon(int tag, const point_3d& pt) const
{
    double eps0 = lua["const"]["eps0"];
    auto eps_dom = lua[SEC_MATERIALS][tag][MAT_EPS];
    if (eps_dom.valid())
    {
        double ret = eps_dom;
        return eps0*ret;
    }

    double ret = lua[SEC_MATERIALS][MAT_EPS](tag, pt.x(), pt.y(), pt.z());
    return eps0*ret;
}

double
parameter_loader::mu(int tag, const point_3d& pt) const
{
    double mu0 = lua["const"]["mu0"];
    auto mu_dom = lua[SEC_MATERIALS][tag][MAT_MU];
    if (mu_dom.valid())
    {
        double ret = mu_dom;
        return mu0*ret;
    }

    double ret = lua[SEC_MATERIALS][MAT_MU](tag, pt.x(), pt.y(), pt.z());
    return mu0*ret;
}

double
parameter_loader::sigma(int tag, const point_3d& pt) const
{
    auto sigma_dom = lua[SEC_MATERIALS][tag][MAT_SIGMA];
    if (sigma_dom.valid())
        return sigma_dom;

    return lua[SEC_MATERIALS][MAT_SIGMA](tag, pt.x(), pt.y(), pt.z());
}

boundary_condition
parameter_loader::boundary_type(int tag) const
{
    auto bnd_data = lua[SEC_BCONDS][tag];
    if (not bnd_data.valid())
        return boundary_condition::UNSPECIFIED;

    std::string kind_str = bnd_data[BC_KIND];

    if (kind_str == BCOND_TYPE_PEC)
        return boundary_condition::PEC;

    if (kind_str == BCOND_TYPE_PMC)
        return boundary_condition::PMC;

    if (kind_str == BCOND_TYPE_IMPEDANCE)
        return boundary_condition::IMPEDANCE;

    if (kind_str == BCOND_TYPE_EPLANEW)
        return boundary_condition::PLANE_WAVE_E;

    if (kind_str == BCOND_TYPE_HPLANEW)
        return boundary_condition::PLANE_WAVE_H;

    if (kind_str == BCOND_TYPE_EFIELD)
        return boundary_condition::E_FIELD;
    
    if (kind_str == BCOND_TYPE_HFIELD)
        return boundary_condition::H_FIELD;

    if (kind_str == BCOND_TYPE_SURFCURR)
        return boundary_condition::SURFACE_CURRENT;

    return boundary_condition::UNSPECIFIED;
}

interface_condition
parameter_loader::interface_type(int tag) const
{
    auto if_data = lua[SEC_IFCONDS][tag];
    if (not if_data.valid())
        return interface_condition::UNSPECIFIED;
    
    auto kind = if_data[IFC_KIND];

    std::string kind_str = kind;

    if (kind_str == IFCOND_TYPE_EFIELD)
        return interface_condition::E_FIELD;
    
    if (kind_str == IFCOND_TYPE_HFIELD)
        return interface_condition::H_FIELD;

    if (kind_str == IFCOND_TYPE_SURFCURR)
        return interface_condition::SURFACE_CURRENT;

    return interface_condition::UNSPECIFIED;
}

vec3d
parameter_loader::eval_boundary_source(int tag, const point_3d& pt, double t) const
{
    vec3d ret;
    sol::tie(ret(0), ret(1), ret(2)) =
        lua[SEC_BCONDS][tag][BC_SOURCE](tag, pt.x(), pt.y(), pt.z(), t);
    return ret;
}

vec3d
parameter_loader::eval_interface_source(int tag, const point_3d& pt, double t) const
{
    vec3d ret;
    sol::tie(ret(0), ret(1), ret(2)) =
        lua[SEC_IFCONDS][tag][IFC_SOURCE](tag, pt.x(), pt.y(), pt.z(), t);
    return ret;
}

field_export_mode
parameter_loader::postpro_fieldExportMode(const std::string& field) const
{
    auto pp = lua["postpro"][field];
    if (not pp.valid())
    {
        std::cout << "[CONFIG] warning: invalid field '" << field;
        std::cout << "' in 'postpro' table" << std::endl;
        return field_export_mode::NONE;
    }

    std::string silo_mode = pp["silo_mode"];

    if (silo_mode == "nodal")
        return field_export_mode::NODAL;

    if (silo_mode == "zonal")
        return field_export_mode::ZONAL;

    if (silo_mode != "none")
    {
        std::cout << "[CONFIG] warning: invalid export mode '" << silo_mode;
        std::cout << "'" << std::endl;
    }

    return field_export_mode::NONE;
}

} // namespace maxwell