/*
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/** \file
 * \ingroup bmesh
 *
 * Advanced smoothing.
 */

#include "MEM_guardedalloc.h"

#include "BLI_math.h"

#include "eigen_capi.h"

#include "bmesh.h"

#include "intern/bmesh_operators_private.h" /* own include */

// #define SMOOTH_LAPLACIAN_AREA_FACTOR 4.0f  /* UNUSED */
// #define SMOOTH_LAPLACIAN_EDGE_FACTOR 2.0f  /* UNUSED */
#define SMOOTH_LAPLACIAN_MAX_EDGE_PERCENTAGE 1.8f
#define SMOOTH_LAPLACIAN_MIN_EDGE_PERCENTAGE 0.15f

struct BLaplacianSystem {
  float *eweights;      /* Length weights per Edge */
  float (*fweights)[3]; /* Cotangent weights per face */
  float *ring_areas;    /* Total area per ring*/
  float *vlengths;      /* Total sum of lengths(edges) per vertice*/
  float *vweights;      /* Total sum of weights per vertice*/
  int numEdges;         /* Number of edges*/
  int numFaces;         /* Number of faces*/
  int numVerts;         /* Number of verts*/
  short *zerola;        /* Is zero area or length*/

  /* Pointers to data*/
  BMesh *bm;
  BMOperator *op;
  LinearSolver *context;

  /*Data*/
  float min_area;
};
typedef struct BLaplacianSystem LaplacianSystem;

static bool vert_is_boundary(BMVert *v);
static LaplacianSystem *init_laplacian_system(int a_numEdges, int a_numFaces, int a_numVerts);
static void init_laplacian_matrix(LaplacianSystem *sys);
static void delete_laplacian_system(LaplacianSystem *sys);
static void delete_void_pointer(void *data);
static void fill_laplacian_matrix(LaplacianSystem *sys);
static void memset_laplacian_system(LaplacianSystem *sys, int val);
static void validate_solution(
    LaplacianSystem *sys, int usex, int usey, int usez, int preserve_volume);
static void volume_preservation(
    BMOperator *op, float vini, float vend, int usex, int usey, int usez);

static void delete_void_pointer(void *data)
{
  if (data) {
    MEM_freeN(data);
  }
}

static void delete_laplacian_system(LaplacianSystem *sys)
{
  delete_void_pointer(sys->eweights);
  delete_void_pointer(sys->fweights);
  delete_void_pointer(sys->ring_areas);
  delete_void_pointer(sys->vlengths);
  delete_void_pointer(sys->vweights);
  delete_void_pointer(sys->zerola);
  if (sys->context) {
    EIG_linear_solver_delete(sys->context);
  }
  sys->bm = NULL;
  sys->op = NULL;
  MEM_freeN(sys);
}

static void memset_laplacian_system(LaplacianSystem *sys, int val)
{
  memset(sys->eweights, val, sizeof(float) * sys->numEdges);
  memset(sys->fweights, val, sizeof(float) * sys->numFaces * 3);
  memset(sys->ring_areas, val, sizeof(float) * sys->numVerts);
  memset(sys->vlengths, val, sizeof(float) * sys->numVerts);
  memset(sys->vweights, val, sizeof(float) * sys->numVerts);
  memset(sys->zerola, val, sizeof(short) * sys->numVerts);
}

static LaplacianSystem *init_laplacian_system(int a_numEdges, int a_numFaces, int a_numVerts)
{
  LaplacianSystem *sys;
  sys = MEM_callocN(sizeof(LaplacianSystem), "ModLaplSmoothSystem");
  sys->numEdges = a_numEdges;
  sys->numFaces = a_numFaces;
  sys->numVerts = a_numVerts;

  sys->eweights = MEM_callocN(sizeof(float) * sys->numEdges, "ModLaplSmoothEWeight");
  if (!sys->eweights) {
    delete_laplacian_system(sys);
    return NULL;
  }

  sys->fweights = MEM_callocN(sizeof(float[3]) * sys->numFaces, "ModLaplSmoothFWeight");
  if (!sys->fweights) {
    delete_laplacian_system(sys);
    return NULL;
  }

  sys->ring_areas = MEM_callocN(sizeof(float) * sys->numVerts, "ModLaplSmoothRingAreas");
  if (!sys->ring_areas) {
    delete_laplacian_system(sys);
    return NULL;
  }

  sys->vlengths = MEM_callocN(sizeof(float) * sys->numVerts, "ModLaplSmoothVlengths");
  if (!sys->vlengths) {
    delete_laplacian_system(sys);
    return NULL;
  }

  sys->vweights = MEM_callocN(sizeof(float) * sys->numVerts, "ModLaplSmoothVweights");
  if (!sys->vweights) {
    delete_laplacian_system(sys);
    return NULL;
  }

  sys->zerola = MEM_callocN(sizeof(short) * sys->numVerts, "ModLaplSmoothZeloa");
  if (!sys->zerola) {
    delete_laplacian_system(sys);
    return NULL;
  }

  return sys;
}

/**
 * Compute weight between vertice v_i and all your neighbors
 * weight between v_i and v_neighbor
 * <pre>
 * Wij = cot(alpha) + cot(beta) / (4.0 * total area of all faces  * sum all weight)
 *
 *        v_i *
 *          / | \
 *         /  |  \
 *  v_beta*   |   * v_alpha
 *         \  |  /
 *          \ | /
 *            * v_neighbor
 * </pre>
 */

static void init_laplacian_matrix(LaplacianSystem *sys)
{
  float areaf;
  float *v1, *v2, *v3, *v4;
  float w1, w2, w3, w4;
  int i, j;
  bool has_4_vert;
  uint idv1, idv2, idv3, idv4, idv[4];
  BMEdge *e;
  BMFace *f;
  BMIter eiter;
  BMIter fiter;
  BMIter vi;
  BMVert *vn;
  BMVert *vf[4];

  BM_ITER_MESH_INDEX (e, &eiter, sys->bm, BM_EDGES_OF_MESH, i) {
    if (!BM_elem_flag_test(e, BM_ELEM_SELECT) && BM_edge_is_boundary(e)) {
      v1 = e->v1->co;
      v2 = e->v2->co;
      idv1 = BM_elem_index_get(e->v1);
      idv2 = BM_elem_index_get(e->v2);

      w1 = len_v3v3(v1, v2);
      if (w1 > sys->min_area) {
        w1 = 1.0f / w1;
        sys->eweights[i] = w1;
        sys->vlengths[idv1] += w1;
        sys->vlengths[idv2] += w1;
      }
      else {
        sys->zerola[idv1] = 1;
        sys->zerola[idv2] = 1;
      }
    }
  }

  BM_ITER_MESH_INDEX (f, &fiter, sys->bm, BM_FACES_OF_MESH, i) {
    if (BM_elem_flag_test(f, BM_ELEM_SELECT)) {

      BM_ITER_ELEM_INDEX (vn, &vi, f, BM_VERTS_OF_FACE, j) {
        vf[j] = vn;
      }
      has_4_vert = (j == 4) ? 1 : 0;
      idv1 = BM_elem_index_get(vf[0]);
      idv2 = BM_elem_index_get(vf[1]);
      idv3 = BM_elem_index_get(vf[2]);
      idv4 = has_4_vert ? BM_elem_index_get(vf[3]) : 0;

      v1 = vf[0]->co;
      v2 = vf[1]->co;
      v3 = vf[2]->co;
      v4 = has_4_vert ? vf[3]->co : NULL;

      if (has_4_vert) {
        areaf = area_quad_v3(v1, v2, v3, v4);
      }
      else {
        areaf = area_tri_v3(v1, v2, v3);
      }

      if (fabsf(areaf) < sys->min_area) {
        sys->zerola[idv1] = 1;
        sys->zerola[idv2] = 1;
        sys->zerola[idv3] = 1;
        if (has_4_vert) {
          sys->zerola[idv4] = 1;
        }
      }

      sys->ring_areas[idv1] += areaf;
      sys->ring_areas[idv2] += areaf;
      sys->ring_areas[idv3] += areaf;
      if (has_4_vert) {
        sys->ring_areas[idv4] += areaf;
      }

      if (has_4_vert) {

        idv[0] = idv1;
        idv[1] = idv2;
        idv[2] = idv3;
        idv[3] = idv4;

        for (j = 0; j < 4; j++) {
          idv1 = idv[j];
          idv2 = idv[(j + 1) % 4];
          idv3 = idv[(j + 2) % 4];
          idv4 = idv[(j + 3) % 4];

          v1 = vf[j]->co;
          v2 = vf[(j + 1) % 4]->co;
          v3 = vf[(j + 2) % 4]->co;
          v4 = vf[(j + 3) % 4]->co;

          w2 = cotangent_tri_weight_v3(v4, v1, v2) + cotangent_tri_weight_v3(v3, v1, v2);
          w3 = cotangent_tri_weight_v3(v2, v3, v1) + cotangent_tri_weight_v3(v4, v1, v3);
          w4 = cotangent_tri_weight_v3(v2, v4, v1) + cotangent_tri_weight_v3(v3, v4, v1);

          sys->vweights[idv1] += (w2 + w3 + w4) / 4.0f;
        }
      }
      else {
        w1 = cotangent_tri_weight_v3(v1, v2, v3);
        w2 = cotangent_tri_weight_v3(v2, v3, v1);
        w3 = cotangent_tri_weight_v3(v3, v1, v2);

        sys->fweights[i][0] += w1;
        sys->fweights[i][1] += w2;
        sys->fweights[i][2] += w3;

        sys->vweights[idv1] += w2 + w3;
        sys->vweights[idv2] += w1 + w3;
        sys->vweights[idv3] += w1 + w2;
      }
    }
  }
}

static void fill_laplacian_matrix(LaplacianSystem *sys)
{
  float *v1, *v2, *v3, *v4;
  float w2, w3, w4;
  int i, j;
  bool has_4_vert;
  uint idv1, idv2, idv3, idv4, idv[4];

  BMEdge *e;
  BMFace *f;
  BMIter eiter;
  BMIter fiter;
  BMIter vi;
  BMVert *vn;
  BMVert *vf[4];

  BM_ITER_MESH_INDEX (f, &fiter, sys->bm, BM_FACES_OF_MESH, i) {
    if (BM_elem_flag_test(f, BM_ELEM_SELECT)) {
      BM_ITER_ELEM_INDEX (vn, &vi, f, BM_VERTS_OF_FACE, j) {
        vf[j] = vn;
      }
      has_4_vert = (j == 4) ? 1 : 0;
      if (has_4_vert) {
        idv[0] = BM_elem_index_get(vf[0]);
        idv[1] = BM_elem_index_get(vf[1]);
        idv[2] = BM_elem_index_get(vf[2]);
        idv[3] = BM_elem_index_get(vf[3]);
        for (j = 0; j < 4; j++) {
          idv1 = idv[j];
          idv2 = idv[(j + 1) % 4];
          idv3 = idv[(j + 2) % 4];
          idv4 = idv[(j + 3) % 4];

          v1 = vf[j]->co;
          v2 = vf[(j + 1) % 4]->co;
          v3 = vf[(j + 2) % 4]->co;
          v4 = vf[(j + 3) % 4]->co;

          w2 = cotangent_tri_weight_v3(v4, v1, v2) + cotangent_tri_weight_v3(v3, v1, v2);
          w3 = cotangent_tri_weight_v3(v2, v3, v1) + cotangent_tri_weight_v3(v4, v1, v3);
          w4 = cotangent_tri_weight_v3(v2, v4, v1) + cotangent_tri_weight_v3(v3, v4, v1);

          w2 = w2 / 4.0f;
          w3 = w3 / 4.0f;
          w4 = w4 / 4.0f;

          if (!vert_is_boundary(vf[j]) && sys->zerola[idv1] == 0) {
            EIG_linear_solver_matrix_add(sys->context, idv1, idv2, w2 * sys->vweights[idv1]);
            EIG_linear_solver_matrix_add(sys->context, idv1, idv3, w3 * sys->vweights[idv1]);
            EIG_linear_solver_matrix_add(sys->context, idv1, idv4, w4 * sys->vweights[idv1]);
          }
        }
      }
      else {
        idv1 = BM_elem_index_get(vf[0]);
        idv2 = BM_elem_index_get(vf[1]);
        idv3 = BM_elem_index_get(vf[2]);
        /* Is ring if number of faces == number of edges around vertice. */
        if (!vert_is_boundary(vf[0]) && sys->zerola[idv1] == 0) {
          EIG_linear_solver_matrix_add(
              sys->context, idv1, idv2, sys->fweights[i][2] * sys->vweights[idv1]);
          EIG_linear_solver_matrix_add(
              sys->context, idv1, idv3, sys->fweights[i][1] * sys->vweights[idv1]);
        }
        if (!vert_is_boundary(vf[1]) && sys->zerola[idv2] == 0) {
          EIG_linear_solver_matrix_add(
              sys->context, idv2, idv1, sys->fweights[i][2] * sys->vweights[idv2]);
          EIG_linear_solver_matrix_add(
              sys->context, idv2, idv3, sys->fweights[i][0] * sys->vweights[idv2]);
        }
        if (!vert_is_boundary(vf[2]) && sys->zerola[idv3] == 0) {
          EIG_linear_solver_matrix_add(
              sys->context, idv3, idv1, sys->fweights[i][1] * sys->vweights[idv3]);
          EIG_linear_solver_matrix_add(
              sys->context, idv3, idv2, sys->fweights[i][0] * sys->vweights[idv3]);
        }
      }
    }
  }
  BM_ITER_MESH_INDEX (e, &eiter, sys->bm, BM_EDGES_OF_MESH, i) {
    if (!BM_elem_flag_test(e, BM_ELEM_SELECT) && BM_edge_is_boundary(e)) {
      v1 = e->v1->co;
      v2 = e->v2->co;
      idv1 = BM_elem_index_get(e->v1);
      idv2 = BM_elem_index_get(e->v2);
      if (sys->zerola[idv1] == 0 && sys->zerola[idv2] == 0) {
        EIG_linear_solver_matrix_add(
            sys->context, idv1, idv2, sys->eweights[i] * sys->vlengths[idv1]);
        EIG_linear_solver_matrix_add(
            sys->context, idv2, idv1, sys->eweights[i] * sys->vlengths[idv2]);
      }
    }
  }
}

static bool vert_is_boundary(BMVert *v)
{
  BMEdge *ed;
  BMFace *f;
  BMIter ei;
  BMIter fi;
  BM_ITER_ELEM (ed, &ei, v, BM_EDGES_OF_VERT) {
    if (BM_edge_is_boundary(ed)) {
      return 1;
    }
  }
  BM_ITER_ELEM (f, &fi, v, BM_FACES_OF_VERT) {
    if (!BM_elem_flag_test(f, BM_ELEM_SELECT)) {
      return 1;
    }
  }
  return 0;
}

static void volume_preservation(
    BMOperator *op, float vini, float vend, int usex, int usey, int usez)
{
  float beta;
  BMOIter siter;
  BMVert *v;

  if (vend != 0.0f) {
    beta = pow(vini / vend, 1.0f / 3.0f);
    BMO_ITER (v, &siter, op->slots_in, "verts", BM_VERT) {
      if (usex) {
        v->co[0] *= beta;
      }
      if (usey) {
        v->co[1] *= beta;
      }
      if (usez) {
        v->co[2] *= beta;
      }
    }
  }
}

static void validate_solution(
    LaplacianSystem *sys, int usex, int usey, int usez, int preserve_volume)
{
  int m_vertex_id;
  float leni, lene;
  float vini, vend;
  float *vi1, *vi2, ve1[3], ve2[3];
  uint idv1, idv2;
  BMOIter siter;
  BMVert *v;
  BMEdge *e;
  BMIter eiter;

  BM_ITER_MESH (e, &eiter, sys->bm, BM_EDGES_OF_MESH) {
    idv1 = BM_elem_index_get(e->v1);
    idv2 = BM_elem_index_get(e->v2);
    vi1 = e->v1->co;
    vi2 = e->v2->co;
    ve1[0] = EIG_linear_solver_variable_get(sys->context, 0, idv1);
    ve1[1] = EIG_linear_solver_variable_get(sys->context, 1, idv1);
    ve1[2] = EIG_linear_solver_variable_get(sys->context, 2, idv1);
    ve2[0] = EIG_linear_solver_variable_get(sys->context, 0, idv2);
    ve2[1] = EIG_linear_solver_variable_get(sys->context, 1, idv2);
    ve2[2] = EIG_linear_solver_variable_get(sys->context, 2, idv2);
    leni = len_v3v3(vi1, vi2);
    lene = len_v3v3(ve1, ve2);
    if (lene > leni * SMOOTH_LAPLACIAN_MAX_EDGE_PERCENTAGE ||
        lene < leni * SMOOTH_LAPLACIAN_MIN_EDGE_PERCENTAGE) {
      sys->zerola[idv1] = 1;
      sys->zerola[idv2] = 1;
    }
  }

  if (preserve_volume) {
    vini = BM_mesh_calc_volume(sys->bm, false);
  }
  BMO_ITER (v, &siter, sys->op->slots_in, "verts", BM_VERT) {
    m_vertex_id = BM_elem_index_get(v);
    if (sys->zerola[m_vertex_id] == 0) {
      if (usex) {
        v->co[0] = EIG_linear_solver_variable_get(sys->context, 0, m_vertex_id);
      }
      if (usey) {
        v->co[1] = EIG_linear_solver_variable_get(sys->context, 1, m_vertex_id);
      }
      if (usez) {
        v->co[2] = EIG_linear_solver_variable_get(sys->context, 2, m_vertex_id);
      }
    }
  }
  if (preserve_volume) {
    vend = BM_mesh_calc_volume(sys->bm, false);
    volume_preservation(sys->op, vini, vend, usex, usey, usez);
  }
}

void bmo_smooth_laplacian_vert_exec(BMesh *bm, BMOperator *op)
{
  int i;
  int m_vertex_id;
  bool usex, usey, usez, preserve_volume;
  float lambda_factor, lambda_border;
  float w;
  BMOIter siter;
  BMVert *v;
  LaplacianSystem *sys;

  if (bm->totface == 0) {
    return;
  }
  sys = init_laplacian_system(bm->totedge, bm->totface, bm->totvert);
  if (!sys) {
    return;
  }
  sys->bm = bm;
  sys->op = op;

  memset_laplacian_system(sys, 0);

  BM_mesh_elem_index_ensure(bm, BM_VERT);
  lambda_factor = BMO_slot_float_get(op->slots_in, "lambda_factor");
  lambda_border = BMO_slot_float_get(op->slots_in, "lambda_border");
  sys->min_area = 0.00001f;
  usex = BMO_slot_bool_get(op->slots_in, "use_x");
  usey = BMO_slot_bool_get(op->slots_in, "use_y");
  usez = BMO_slot_bool_get(op->slots_in, "use_z");
  preserve_volume = BMO_slot_bool_get(op->slots_in, "preserve_volume");

  sys->context = EIG_linear_least_squares_solver_new(bm->totvert, bm->totvert, 3);

  for (i = 0; i < bm->totvert; i++) {
    EIG_linear_solver_variable_lock(sys->context, i);
  }
  BMO_ITER (v, &siter, op->slots_in, "verts", BM_VERT) {
    m_vertex_id = BM_elem_index_get(v);
    EIG_linear_solver_variable_unlock(sys->context, m_vertex_id);
    EIG_linear_solver_variable_set(sys->context, 0, m_vertex_id, v->co[0]);
    EIG_linear_solver_variable_set(sys->context, 1, m_vertex_id, v->co[1]);
    EIG_linear_solver_variable_set(sys->context, 2, m_vertex_id, v->co[2]);
  }

  init_laplacian_matrix(sys);
  BMO_ITER (v, &siter, op->slots_in, "verts", BM_VERT) {
    m_vertex_id = BM_elem_index_get(v);
    EIG_linear_solver_right_hand_side_add(sys->context, 0, m_vertex_id, v->co[0]);
    EIG_linear_solver_right_hand_side_add(sys->context, 1, m_vertex_id, v->co[1]);
    EIG_linear_solver_right_hand_side_add(sys->context, 2, m_vertex_id, v->co[2]);
    i = m_vertex_id;
    if (sys->zerola[i] == 0) {
      w = sys->vweights[i] * sys->ring_areas[i];
      sys->vweights[i] = (w == 0.0f) ? 0.0f : -lambda_factor / (4.0f * w);
      w = sys->vlengths[i];
      sys->vlengths[i] = (w == 0.0f) ? 0.0f : -lambda_border * 2.0f / w;

      if (!vert_is_boundary(v)) {
        EIG_linear_solver_matrix_add(
            sys->context, i, i, 1.0f + lambda_factor / (4.0f * sys->ring_areas[i]));
      }
      else {
        EIG_linear_solver_matrix_add(sys->context, i, i, 1.0f + lambda_border * 2.0f);
      }
    }
    else {
      EIG_linear_solver_matrix_add(sys->context, i, i, 1.0f);
    }
  }
  fill_laplacian_matrix(sys);

  if (EIG_linear_solver_solve(sys->context)) {
    validate_solution(sys, usex, usey, usez, preserve_volume);
  }

  delete_laplacian_system(sys);
}
