#include <iostream>
#include <fstream>
#include <cassert>

#include "mfem.hpp"

extern "C" {
#include "xf.h"
#include "xf_All.h"
#include "xf_Mesh.h"
#include "xf_Basis.h"
#include "xf_Data.h"
#include "xf_State.h"
#include "xf_Math.h"
}

// Convert XFlow shapes to MFEM geometry
mfem::Geometry::Type ShapeToGeometry(enum xfe_ShapeType shape)
{
  switch (shape)
  {
  case xfe_Point         : return mfem::Geometry::POINT;
  case xfe_Segment       : return mfem::Geometry::SEGMENT;
  case xfe_Triangle      : return mfem::Geometry::TRIANGLE;
  case xfe_Quadrilateral : return mfem::Geometry::SQUARE;
  case xfe_Hexahedron    : return mfem::Geometry::CUBE;
  case xfe_Tetrahedron   : return mfem::Geometry::TETRAHEDRON;
  default: throw std::runtime_error("Unknown shape");
  }
}

void ReverseIntArray(const int *input, int size, int *output)
{
  for (int i = 0; i < size; i++) output[input[i]] = i;
}

mfem::GridFunction convertXFVector(const xf_Mesh *xflow_mesh,
                                   mfem::Mesh *mesh,
                                   const xf_Vector *V)
{
  int ierr;
  const int vdim = V->StateRank;
  const int dim = xflow_mesh->Dim;

  // Make sure the vector is interpolated
  if (V->Order == NULL) throw std::runtime_error("Data is not interpolated");

  // ... and has element linkage
  if (V->Linkage != xfe_LinkageGlobElem) throw std::runtime_error("Is not implemented for non-element-linked data");

  // ... and that the number of element groups equals the number of arrays
  if (V->nArray != xflow_mesh->nElemGroup) throw std::runtime_error("Need number of elements groups to match number of arrays");

  // Find max order
  int max_order = 0;
  for (int array = 0; array < V->nArray; array++)
  {
    if (V->vOrder == NULL)
    {
      const int order = V->Order[array];
      max_order = order > max_order ? order : max_order;
    }
    else
    {
      for (int comp = 0; comp < V->nComp[array]; comp++)
      {
        const int order = V->vOrder[array][comp];
        max_order = order > max_order ? order : max_order;
      }
    }
  }

  enum xfe_ShapeType elem_shape;
  ierr = xf_Error(xf_Basis2Shape(xflow_mesh->ElemGroup[0].QBasis, &elem_shape));
  if (ierr != xf_OK) throw std::runtime_error("An error occured");

  // elem_basis is the basis of the solution _not_ necessarily that of the mesh
  enum xfe_BasisType elem_basis;
  ierr = xf_Error(xf_Shape2UniformLagrange(elem_shape, &elem_basis));
  if (ierr != xf_OK) throw std::runtime_error("An error occured");

  int nn;
  ierr = xf_Error(xf_Order2nNode(elem_basis, max_order, &nn));
  if (ierr != xf_OK) throw std::runtime_error("An error occured");

  mfem::L2_FECollection *fec = new mfem::L2_FECollection(max_order, xflow_mesh->Dim, mfem::BasisType::ClosedUniform);
  mfem::FiniteElementSpace *fes = new mfem::FiniteElementSpace(mesh, fec, vdim, mfem::Ordering::byVDIM);
  mfem::GridFunction gf(fes);
  // TODO make sure fec gets deleted somehow...
  // gf.MakeOwner(fec); // gf will destroy 'fec' and 'fes'

  // Create a temporary data set
  xf_DataSet *DataSet;
  ierr = xf_Error(xf_CreateDataSet(&DataSet));
  if (ierr != xf_OK) throw std::runtime_error("An error occured");

  // This is the same operation as filling the mesh Nodes
  // GridFunction, but it cannot be reused since xflow does not expose
  // these as an xf_Vector.
  mfem::Array<int> vdofs;
  xf_Matrix *T = NULL;
  for (int egrp = 0; egrp < V->nArray; egrp++)
  {
    xf_GenArray &GA = V->GenArray[egrp];

    for (int elem = 0; elem < GA.n; elem++)
    {
      const enum xfe_BasisType basis = (V->Basis != NULL) ? V->Basis[egrp] : V->vBasis[egrp][elem];
      const int order = (V->vOrder != NULL) ? V->vOrder[egrp][elem] : V->Order[egrp];

      fes->GetElementVDofs(elem, vdofs);
      const mfem::FiniteElement *fe = fes->GetFE(elem);
      const int ndof = fe->GetDof();

      // Assumption: element data (so, use elem_shape from above)
      mfem::Array<double> values;
      if (order != max_order)
      {
        // Variable order - need to find a transfer matrix in XFlow
        ierr = xf_Error(xf_FindTransferMatrix(DataSet, basis, order, elem_basis, max_order, &T));
        if (ierr != xf_OK) throw std::runtime_error("An error occured");

        int ndof_e;
        ierr = xf_Error(xf_Order2nNode(basis, order, &ndof_e));
        if (ierr != xf_OK) throw std::runtime_error("An error occured");

        values.SetSize(ndof_e * ndof);

        xf_MxM_Set(T->GenArray[0].rValue[0], GA.rValue[elem], ndof, ndof_e, vdim, values);
      }
      else
      {
        values.MakeRef(GA.rValue[elem], ndof * vdim);
      }

      // The ordering of MFEM grid functions is _almost_ the same as
      // XFlow (both store lexicographically) but MFEM stores vdofs
      // byNODES so we need to loop over the dimensions manually:
      for (int i = 0; i < ndof; i++) {
        for (int vd = 0; vd < vdim; vd++) {
          gf(vdofs[i+ndof*vd]) = values[i * vdim + vd];
          // This is the syntax if we had to map:
          // gf(dofs[dof_map[i]]*vdim + vd) = values[i * dim + vd];
        }
      }
      // If the ordering where the same we could simply use this:
      // gf.SetSubVector(vdofs, values);
    }
  }

  ierr = xf_Error(xf_DestroyDataSet(DataSet));
  if (ierr != xf_OK) throw std::runtime_error("An error occured");

  return gf;
}

class NodeHash
{
private:
  std::vector<int> forward;
  std::vector<int> reverse;  // indexed by node

public:
  NodeHash() : forward(), reverse() { }

  int add(int node) // Returns the tag
  {
    if (reverse.size() < node+1)
    {
      reverse.resize(node+1, -1);
    }

    int tag = reverse[node];
    if (tag < 0)
    {
      reverse[node] = forward.size();
      forward.push_back(node);
    }

    return tag;
  }

  int get(int tag) const
  {
    return forward[tag];
  }

  int size() const
  {
    return forward.size();
  }
};

class XFMesh : public mfem::Mesh
{
public:
  XFMesh(xf_Mesh &xflow_mesh);
};

XFMesh::XFMesh(xf_Mesh &xflow_mesh) : mfem::Mesh()
{
  int ierr;
  std::vector<int> dof_map_rev, verts, fvec;

  // Count elements
  int nElemTot = 0;
  for (int egrp = 0; egrp < xflow_mesh.nElemGroup; egrp++)
  {
    nElemTot += xflow_mesh.ElemGroup[egrp].nElem;
  }

  // Count boundary elements
  int nBdrElem = 0;
  for (int bfgrp = 0; bfgrp < xflow_mesh.nBFaceGroup; bfgrp++)
  {
    nBdrElem += xflow_mesh.BFaceGroup[bfgrp].nBFace;
  }

  const int dim = xflow_mesh.Dim;

  // Check whether or not QOrder changes. If so, we cannot
  // represent the mesh nodes by a simple H1 grid function.
  int pQOrder = xflow_mesh.ElemGroup[0].QOrder;
  bool QOrderChanges = false;
  int MaxQOrder = pQOrder;
  for (int egrp = 1; egrp < xflow_mesh.nElemGroup; egrp++)
  {
    QOrderChanges |= (xflow_mesh.ElemGroup[egrp].QOrder != pQOrder);
    MaxQOrder = max(xflow_mesh.ElemGroup[egrp].QOrder, MaxQOrder);
  }
  if (QOrderChanges)
  {
    std::cout << "QOrder changes - will not curve mesh. In the future this function could determine the high order nodes on the low order elements." << std::endl;
  }

  // Create a FiniteElementCollection early on for dof_map access
  // dof_map = Cartesian to local H1 dof map
  mfem::H1_FECollection *fec = new mfem::H1_FECollection(MaxQOrder, dim);

  // Node hash to get the vertex ordering
  NodeHash nh;
  for (int egrp = 0; egrp < xflow_mesh.nElemGroup; egrp++)
  {
    const xf_ElemGroup &EG = xflow_mesh.ElemGroup[egrp];

    enum xfe_ShapeType shape;
    ierr = xf_Basis2Shape(EG.QBasis, &shape);
    if (ierr != xf_OK) throw std::runtime_error("An error occured");

    const int geom_type = ShapeToGeometry(shape);
    const int *dof_map = fec->GetDofMap(geom_type);
    const mfem::FiniteElement *fe = fec->FiniteElementForGeometry(geom_type);
    const int num_dof = fe->GetDof();
    dof_map_rev.resize(num_dof);
    ReverseIntArray(dof_map, num_dof, dof_map_rev.data());

    const int num_verts = mfem::Geometry::NumVerts[geom_type];
    for (int elem = 0; elem < EG.nElem; elem++)
    {
      for (int i = 0; i < num_verts; i++)
      {
        nh.add(EG.Node[elem][dof_map_rev[i]]);
      }
    }
  }

  InitMesh(dim, dim, nh.size(), nElemTot, nBdrElem);

  // 1. Add the vertices
  for (int i = 0; i < nh.size(); i++)
  {
    const int node = nh.get(i);
    AddVertex(xflow_mesh.Coord[node]);
  }
  NumOfVertices = nh.size();

  // 2. Add the elements
  int egrp_offset = 0;
  for (int egrp = 0; egrp < xflow_mesh.nElemGroup; egrp++)
  {
    const xf_ElemGroup &EG = xflow_mesh.ElemGroup[egrp];
    mfem::Element **elements_egrp = elements.GetData() + egrp_offset;

    enum xfe_ShapeType shape;
    ierr = xf_Basis2Shape(EG.QBasis, &shape);
    if (ierr != xf_OK) throw std::runtime_error("An error occured");

    switch (shape)
    {
    case xfe_Segment:
      for (int e = 0; e < EG.nElem; e++) elements_egrp[e] = new mfem::Segment();
      break;
    case xfe_Quadrilateral:
      for (int e = 0; e < EG.nElem; e++) elements_egrp[e] = new mfem::Quadrilateral();
      break;
    case xfe_Triangle:
      for (int e = 0; e < EG.nElem; e++) elements_egrp[e] = new mfem::Triangle();
      break;
    case xfe_Hexahedron:
      for (int e = 0; e < EG.nElem; e++) elements_egrp[e] = new mfem::Hexahedron();
      break;
    case xfe_Tetrahedron:
      for (int e = 0; e < EG.nElem; e++) elements_egrp[e] = new mfem::Tetrahedron();
      break;
    }

    const int geom_type = ShapeToGeometry(shape);
    const int *dof_map = fec->GetDofMap(geom_type);
    const mfem::FiniteElement *fe = fec->FiniteElementForGeometry(geom_type);
    const int num_dof = fe->GetDof();
    dof_map_rev.resize(num_dof);
    ReverseIntArray(dof_map, num_dof, dof_map_rev.data());

    const int num_verts = mfem::Geometry::NumVerts[geom_type];
    verts.resize(num_verts);
    for (int elem = 0; elem < EG.nElem; elem++)
    {
      for (int i = 0; i < num_verts; i++)
      {
        verts[i] = nh.add(EG.Node[elem][dof_map_rev[i]]);
      }
      elements_egrp[elem]->SetVertices(verts.data());
      elements_egrp[elem]->SetAttribute(1 + egrp);
    }

    egrp_offset += EG.nElem;
  }
  NumOfElements = egrp_offset;

  // 3. Add the boundary elements
  int bfgrp_offset = 0;
  for (int bfgrp = 0; bfgrp < xflow_mesh.nBFaceGroup; bfgrp++)
  {
    xf_BFaceGroup &BFG = xflow_mesh.BFaceGroup[bfgrp];

    mfem::Element **boundary_bfgrp = boundary.GetData() + bfgrp_offset;

    for (int bface = 0; bface < BFG.nBFace; bface++)
    {
      // Every element in the boundary face group may have a different
      // shape or derive from an element of different QOrder, hence
      // the allocation needs to happen inside the loop over bface.

      const xf_BFace &BF = BFG.BFace[bface];
      const xf_ElemGroup &EG = xflow_mesh.ElemGroup[BF.ElemGroup];

      enum xfe_ShapeType elem_shape, face_shape;
      ierr = xf_Basis2Shape(EG.QBasis, &elem_shape);
      if (ierr != xf_OK) throw std::runtime_error("An error occured");
      ierr = xf_FaceShape(elem_shape, BF.Face, &face_shape);
      if (ierr != xf_OK) throw std::runtime_error("An error occured");

      switch (face_shape)
      {
      case xfe_Point:
        boundary_bfgrp[bface] = new mfem::Point();
        break;
      case xfe_Segment:
        boundary_bfgrp[bface] = new mfem::Segment();
        break;
      case xfe_Quadrilateral:
        boundary_bfgrp[bface] = new mfem::Quadrilateral();
        break;
      case xfe_Triangle:
        boundary_bfgrp[bface] = new mfem::Triangle();
        break;
      }

      const int geom_type = ShapeToGeometry(face_shape);
      const mfem::FiniteElement *fe = fec->FiniteElementForGeometry(geom_type);
      int num_dof;
      if (geom_type != mfem::Geometry::POINT)
      {
        const int *dof_map = fec->GetDofMap(geom_type);
        num_dof = fe->GetDof();
        ReverseIntArray(dof_map, num_dof, dof_map_rev.data());
      }
      else
      {
        const int dof_map[1] = {0};
        num_dof = 1;
        ReverseIntArray(dof_map, num_dof, dof_map_rev.data());
      }
      dof_map_rev.resize(num_dof);

      const int num_verts = mfem::Geometry::NumVerts[geom_type];

      int nqnode;
      fvec.resize(num_dof);
      ierr = xf_NodesOnFace(EG.QBasis, EG.QOrder, BF.Face, &nqnode, fvec.data());
      if (ierr != xf_OK) throw std::runtime_error("An error occured");

      verts.resize(num_verts);
      for (int i = 0; i < num_verts; i++)
      {
        verts[i] = nh.add(EG.Node[BF.Elem][fvec[dof_map_rev[i]]]);
      }
      boundary_bfgrp[bface]->SetVertices(verts.data());
      boundary_bfgrp[bface]->SetAttribute(1 + bfgrp);
    }

    bfgrp_offset += BFG.nBFace;
  }
  NumOfBdrElements = bfgrp_offset;

  FinalizeTopology();

  if (MaxQOrder == 1) return;

  // --- Set curvature to MaxQOrder ---

  // The method below works, but it does extra work
  // SetCurvature(MaxQOrder, false, dim, mfem::Ordering::byVDIM);

  // Adding back the Q1 nodes manually is much more efficient
  mfem::FiniteElementSpace *fes = new mfem::FiniteElementSpace(this, fec, dim,
                                                               mfem::Ordering::byVDIM);
  Nodes = new mfem::GridFunction(fes);
  Nodes->MakeOwner(fec); // Nodes will destroy 'fec' and 'fes'
  own_nodes = 1;

  mfem::Array<int> dofs;
  for (int egrp = 0; egrp < xflow_mesh.nElemGroup; egrp++)
  {
    xf_ElemGroup &EG = xflow_mesh.ElemGroup[egrp];
    enum xfe_ShapeType shape;

    ierr = xf_Basis2Shape(EG.QBasis, &shape);
    if (ierr != xf_OK) throw std::runtime_error("An error occured");

    const int geom_type = ShapeToGeometry(shape);

    const int *dof_map = fec->GetDofMap(geom_type);
    for (int elem = 0; elem < EG.nElem; elem++)
    {
      fes->GetElementDofs(elem, dofs);
      const mfem::FiniteElement *fe = fes->GetFE(elem);
      const int ndof = fe->GetDof();
      for (int i = 0; i < ndof; i++) {
        for (int vd = 0; vd < dim; vd++) {
          (*Nodes)(dofs[dof_map[i]]*dim + vd) = xflow_mesh.Coord[0][EG.Node[elem][i] * dim + vd];
        }
      }
    }
  }
}

int main(int argc, char *argv[])
{
  int ierr;
  xf_All *All;
  const char *xf_mesh_file = "NULL";
  const char *xf_data_file = "NULL";
  const char *xf_xfa_file = "NULL";
  const char *convert_list_char = "";
  bool write = false;
  bool plot = true;

  /// GLVis options
  char vishost[] = "localhost";
  int  visport   = 19916;

  mfem::OptionsParser args(argc, argv);
  args.AddOption(&xf_mesh_file, "-m", "--mesh",
                 "Mesh file to use.");
  args.AddOption(&xf_data_file, "-d", "--data",
                 "Data file to use.");
  args.AddOption(&xf_xfa_file, "-x", "--xfa",
                 "Xfa file to use.");
  args.AddOption(&convert_list_char, "-c", "--convert",
                 "Comma-separated list of data names to convert.");
  args.AddOption(&write, "-w", "--write", "-no-w", "--no-write",
                 "Enable or disable output writing.");
  args.AddOption(&plot, "-p", "--plot", "-no-p", "--no-plot",
                 "Enable or disable plotting.");
  args.Parse();
  if (!args.Good())
  {
    args.PrintUsage(std::cout);
    return 1;
  }
  args.PrintOptions(std::cout);

  /* Create .xfa structure */
  xf_Call(xf_CreateAll(&All, xfe_False));

  if (strcmp(xf_mesh_file, "NULL"))
  {
    // Mesh file passed
    std::cout << "Using mesh " << xf_mesh_file << std::endl;
    xf_Call(xf_ReadGriFile(xf_mesh_file, NULL, All->Mesh));
  }
  else if (strcmp(xf_xfa_file, "NULL"))
  {
    // xfa file passed
    std::cout << "Using xfa " << xf_xfa_file << std::endl;
    xf_Call(xf_ReadAllBinary(xf_xfa_file, All));
    if (!strcmp(xf_data_file, "NULL"))
    {
      std::cout << "with " << xf_xfa_file;
      ierr = xf_Error(xf_ReadAllBinary(xf_xfa_file, All));
      if (ierr != xf_OK) throw std::runtime_error("An error occured");
    }
  }

  if (strcmp(xf_data_file, "NULL"))
  {
    // Additional data file passed
    std::cout << "Using data file " << xf_data_file << std::endl;

    xf_DataSet *DataSet;
    xf_Call(xf_CreateDataSet(&DataSet));
    xf_Call(xf_ReadDataSetBinary(All->Mesh, NULL, xf_data_file, DataSet));
    xf_Call(xf_DataSetMerge(DataSet, All->DataSet));
    xf_Call(xf_DestroyDataSet(DataSet));
  }

  std::cout << "Converting mesh" << std::flush;

  // --- Convert mesh ---
  XFMesh mesh(*All->Mesh);

  std::cout << " ... done" << std::endl;

  std::cout << "Reading data" << std::flush;

  std::string convert_list(convert_list_char);
  int num_converted = 0;
  std::vector<mfem::GridFunction> gf_list;
  std::vector<std::string> gf_names;
  for (xf_Data *D = All->DataSet->Head; D != NULL; D = D->Next)
  {
    std::string data_title(D->Title);
    const std::size_t pos = convert_list.find(data_title);
    if (pos == std::string::npos) continue;
    std::cout << "Converting " << data_title << std::endl;
    if (D->Type == xfe_Vector)
    {
      xf_Vector *V = (xf_Vector *) D->Data;
      if (V->Linkage == xfe_LinkageGlobElem)
      {
        gf_list.emplace_back(convertXFVector(All->Mesh, &mesh, V));
        gf_names.emplace_back(std::string(D->Title));
        num_converted++;
      }
    }
    else if (D->Type == xfe_VectorGroup)
    {
      xf_VectorGroup *VG = (xf_VectorGroup *) D->Data;
      for (int i = 0; i < VG->nVector; i++)
      {
        if (VG->Vector[i]->Linkage == xfe_LinkageGlobElem)
        {
          gf_list.emplace_back(convertXFVector(All->Mesh, &mesh, VG->Vector[i]));
          gf_names.emplace_back(std::string(D->Title) + "_Vector" + std::to_string(i));
          num_converted++;
        }
      }
    }
    else
    {
      mfem::mfem_error("Type not supported");
    }
  }

  std::cout << " ... done - number converted = " << num_converted << std::endl;

  if (write)
  {
    // Save mesh
    std::ofstream ofs("xflow.mesh");
    ofs.precision(8);
    mesh.Print(ofs);

    // Save gridfunctions
    for (int i = 0; i < gf_list.size(); i++)
    {
      std::stringstream ss;
      ss << gf_names[i] << ".gf";
      std::ofstream ofs(ss.str());
      ofs.precision(8);
      gf_list[i].Save(ofs);
    }
  }

  if (plot)
  {
    if (gf_list.size() == 0)
    {
      // Plot only the mesh
      mfem::socketstream sout(vishost, visport);
      sout.precision(8);
      sout << "mesh\n" << mesh << std::flush;
    }
    else
    {
      // Loop over gridfunctions
      int Wx = 0, Wy = 0; // window position
      const int Ww = 350, Wh = 350; // window size
      int offx = Ww+10; // window offsets
      for (int i = 0; i < gf_list.size(); i++)
      {
        mfem::socketstream sout(vishost, visport);
        if (gf_list[i].VectorDim() == 2)
        {
          sout << "solution\n" << mesh << gf_list[i] << std::flush;
          sout << "window_title '" << gf_names[i] << "'\n"
               << "window_geometry "
               << Wx << " " << Wy << " " << Ww << " " << Wh << "\n" << std::flush;
        }
        else
        {
          std::cout << "WARNING: Skipping " << gf_names[i] << " because it has more than one vector output. Use -write and plot manually with GLVis using the -gc option." << std::endl;
        }
      }
    }
  }

  /* Destroy .xfa structure */
  xf_Call(xf_DestroyAll(All));

  return 0;
}