/*************************************************************************
 * Copyright (c) 2011 AT&T Intellectual Property
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * https://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors: Details at https://graphviz.org
 *************************************************************************/

#include <algorithm>
#include <common/types.h>
#include <common/globals.h>
#include <sparse/general.h>
#include <math.h>
#include <sparse/SparseMatrix.h>
#include <mingle/edge_bundling.h>
#include <time.h>
#include <sparse/clustering.h>
#include <mingle/ink.h>
#include <mingle/agglomerative_bundling.h>
#include <string.h>
#include <vector>

#define SMALL 1.e-10

static double norm(int n, const double *x) {
  double res = 0;
  int i;

  for (i = 0; i < n; i++) res += x[i]*x[i];
  return sqrt(res);

}

static double sqr_dist(int dim, const double *x, const double *y) {
  int i;
  double res = 0;
  for (i = 0; i < dim; i++) res += (x[i] - y[i])*(x[i] - y[i]);
  return res;
}

static double dist(int dim, const double *x, const double *y) {
  return sqrt(sqr_dist(dim,x,y));
}

static pedge pedge_new(int np, int dim, const double *x) {
  pedge e;

  e.npoints = np;
  e.dim = dim;
  e.len = np;
  e.x.assign(x, &x[dim * e.len]);
  e.edge_length = dist(dim, &x[0 * dim], &x[(np - 1) * dim]);
  e.wgt = 1;
  return e;

}

pedge pedge_wgt_new(int np, int dim, double *x, double wgt) {
  pedge e;

  e.npoints = np;
  e.dim = dim;
  e.len = np;
  e.x.assign(x, &x[dim * e.len]);
  e.edge_length = dist(dim, &x[0 * dim], &x[(np - 1) * dim]);
  e.wgt = wgt;
  e.wgts = std::vector<double>(np - 1, wgt);
  return e;
}

void pedge_delete(pedge &) { }

static double edge_compatibility(const pedge &e1, const pedge &e2) {
  /* two edges are u1->v1, u2->v2.
     return 1 if two edges are exactly the same, 0 if they are very different.
   */
  double dist1, dist2, len1, len2;
  const int dim = e1.dim;
  bool flipped = false;

  const double *u1 = e1.x.data();
  const double *v1 = e1.x.data() + e1.npoints * dim - dim;
  const double *u2 = e2.x.data();
  const double *v2 = e2.x.data() + e2.npoints * dim - dim;
  dist1 = sqr_dist(dim, u1, u2) + sqr_dist(dim, v1, v2);
  dist2 =  sqr_dist(dim, u1, v2) + sqr_dist(dim, v1, u2);
  if (dist1 > dist2){
    std::swap(u2, v2);
    dist1 = dist2;
    flipped = true;
  }
  len1 = dist(dim, u1, v1);
  len2 = dist(dim, u2, v2);
  dist1 = std::max(0.1, dist1 / (len1 + len2 + 0.0001 * dist1));
  if (flipped){
    return -1/dist1; 
  } else {
    return 1/dist1; 
  }
}

static double edge_compatibility_full(const pedge &e1, const pedge &e2) {
  /* two edges are u1->v1, u2->v2.
     return 1 if two edges are exactly the same, 0 if they are very different.
     This is based on Holten and van Wijk's paper
   */
  double dist1, dist2, len1, len2, len;
  double tmp, ca, cp, cs;
  int dim = e1.dim, i;
  bool flipped = false;

  const double *u1 = e1.x.data();
  const double *v1 = e1.x.data() + e1.npoints * dim - dim;
  const double *u2 = e2.x.data();
  const double *v2 = e2.x.data() + e2.npoints * dim - dim;
  dist1 = sqr_dist(dim, u1, u2) + sqr_dist(dim, v1, v2);
  dist2 =  sqr_dist(dim, u1, v2) + sqr_dist(dim, v1, u2);
  if (dist1 > dist2){
    std::swap(u2, v2);
    dist1 = dist2;
    flipped = true;
  }
  len1 = std::max(dist(dim, u1, v1), SMALL);
  len2 = std::max(dist(dim, u2, v2), SMALL);
  len = 0.5*(len1+len2);

  /* angle compatibility */
  ca = 0;
  for (i = 0; i < dim; i++) 
    ca += (v1[i]-u1[i])*(v2[i]-u2[i]);
  ca = fabs(ca/(len1*len2));
  assert(ca > -0.001);

  /* scale compatibility */
  cs = 2 / (std::max(len1, len2) / len + len / std::min(len1, len2));
  assert(cs > -0.001 && cs < 1.001);
 
  /* position compatibility */
  cp = 0;
  for (i = 0; i < dim; i++) {
    tmp = .5*(v1[i]+u1[i])-.5*(v2[i]+u2[i]);
    cp += tmp*tmp;
  }
  cp = sqrt(cp);
  cp = len/(len + cp);
  assert(cp > -0.001 && cp < 1.001);

  /* visibility compatibility */

  dist1 = cp*ca*cs;
  if (flipped){
    return -dist1; 
  } else {
    return dist1; 
  }
}

static void fprint_rgb(FILE* fp, int r, int g, int b, int alpha){
  fprintf(fp, "#%02x%02x%02x%02x", r, g, b, alpha);
}

void pedge_export_gv(FILE *fp, int ne, const std::vector<pedge> &edges) {
  double t;
  int i, j, k, kk, dim, sta, sto;
  double maxwgt = 0, len, len_total0;
  int r, g, b;
  double tt1[3]={0.15,0.5,0.85};
  double tt2[4]={0.15,0.4,0.6,0.85};
  double *tt;

  fprintf(fp,"strict graph{\n");
  /* points */
  for (i = 0; i < ne; i++){
    const pedge &edge = edges[i];
    const std::vector<double> &x = edge.x;
    dim = edge.dim;
    sta = 0;
    sto = edge.npoints - 1;

    fprintf(fp, "%d [pos=\"", i);
    for (k = 0; k < dim; k++) {
      if (k != 0)  fprintf(fp, ",");
      fprintf(fp, "%f", x[sta*dim+k]);
    }
    fprintf(fp, "\"];\n");

    fprintf(fp, "%d [pos=\"", i + ne);
    for (k = 0; k < dim; k++) {
      if (k != 0)  fprintf(fp, ",");
      fprintf(fp, "%f", x[sto*dim+k]);
    }
    fprintf(fp, "\"];\n");

  }

  /* figure out max number of bundled original edges in a pedge */
  for (i = 0; i < ne; i++){
    const pedge &edge = edges[i];
    if (!edge.wgts.empty()) {
      for (j = 0; j < edge.npoints - 1; j++) {
        maxwgt = std::max(maxwgt, edge.wgts[j]);
      }
    }
  }

  /* spline and colors */
  for (i = 0; i < ne; i++){
    fprintf(fp,"%d -- %d [pos=\"", i, i + ne);
    const pedge &edge = edges[i];
    const std::vector<double> &x = edge.x;
    dim = edge.dim;
    /* splines */
    for (j = 0; j < edge.npoints; j++) {
      if (j != 0) {
	int mm = 3;
	fprintf(fp," ");
	/* there are ninterval+1 points, add 3*ninterval+2 points, get rid of internal ninternal-1 points,
	  make into 3*ninterval+4 points so that gviz spline rendering can work */
	if (j == 1 || j == edge.npoints - 1) {
	  // every interval gets 3 points inserted except the first and last one
	  tt = tt2;
	  mm = 4;
	} else {
	  tt = tt1;
	}
	for (kk = 1; kk <= mm; kk++){
	  t = tt[kk-1];
	  for (k = 0; k < dim; k++) {
	    if (k != 0) fprintf(fp,",");
	    fprintf(fp, "%f", (x[(j-1)*dim+k]*(1-t)+x[j*dim+k]*(t)));
	  }
	  fprintf(fp," ");
	}
      }
      if (j == 0 || j == edge.npoints - 1){
	for (k = 0; k < dim; k++) {
	  if (k != 0) fprintf(fp,",");
	  fprintf(fp, "%f", x[j*dim+k]);
	}
      }
    }
    /* colors based on how much bundling */
    if (!edge.wgts.empty()) {
      fprintf(fp, "\", wgts=\"");
      for (j = 0; j < edge.npoints - 1; j++){
	if (j != 0) fprintf(fp,",");
	fprintf(fp, "%f", edge.wgts[j]);
      }

      len_total0 = 0;
      fprintf(fp, "\", color=\"");
      for (j = 0; j < edge.npoints - 1; j++){
	len = 0;
	for (k = 0; k < dim; k++){
	  len += (edge.x[dim * j + k] - edge.x[dim * (j + 1) + k]) * (edge.x[dim * j + k] - edge.x[dim * (j + 1) + k]);
	}
	len = sqrt(len/k);
	len_total0 += len;
      }
      for (j = 0; j < edge.npoints - 1; j++) {
	len = 0;
	for (k = 0; k < dim; k++){
	  len += (edge.x[dim * j + k] - edge.x[dim * (j + 1) + k]) * (edge.x[dim * j + k] - edge.x[dim * (j + 1) + k]);
	}
	len = sqrt(len/k);
	t = edge.wgts[j] / maxwgt;
	// interpolate between red (t = 1) to blue (t = 0)
	r = 255*t; g = 0; b = 255*(1-t); b = 255*(1-t);
	if (j != 0) fprintf(fp,":");
	fprint_rgb(fp, r, g, b, 85);
	if (j < edge.npoints - 2) fprintf(fp, ";%f", len / len_total0);
      }

    }
    fprintf(fp, "\"];\n");
    
  }
  fprintf(fp,"}\n");
}

#ifdef DEBUG
static void pedge_print(char *comments, const pedge &e) {
  int i, j, dim;
  dim = e.dim;
  fprintf(stderr,"%s", comments);
  for (i = 0; i < e.npoints; i++){
    if (i > 0) fprintf(stderr,",");
    fprintf(stderr,"{");
    for (j = 0; j < dim; j++){
      if (j > 0) fprintf(stderr,",");
      fprintf(stderr, "%f", e.x[dim * i + j]);
    }
    fprintf(stderr,"}");
  }
  fprintf(stderr,"\n");
}
#endif

void pedge_wgts_realloc(pedge &e, int n) {
  int i;
  if (n <= e.npoints)
    return;
  e.x.resize(e.dim * n, 0);
  if (e.wgts.empty()) {
    e.wgts.resize(n - 1);
    for (i = 0; i < e.npoints; i++)
      e.wgts[i] = e.wgt;
  } else {
    e.wgts.resize(n - 1);
  }
  e.len = n;
}

void pedge_double(pedge &e) {
  /* double the number of points (more precisely, add a point between two points in the polyline */
  int npoints = e.npoints, len = e.len, i, dim = e.dim;
  int j, ii, ii2, np;

  assert(npoints >= 2);
  if (npoints*2-1 > len){
    len = 3*npoints;
    e.x.resize(dim * len, 0);
  }
  std::vector<double> &x = e.x;

  for (i = npoints - 1; i >= 0; i--){
    ii = 2*i;
    for (j = 0; j < dim; j++){
      x[dim*ii + j] = x[dim*i + j];
    }
  }

  for (i = 0; i < npoints - 1; i++){
    ii = 2*i;/* left and right interpolant of a new point */
    ii2 = 2*(i+1);
    for (j = 0; j < dim; j++){
      x[dim*(2*i + 1) + j] = 0.5*(x[dim*ii + j] + x[dim*ii2 + j]);
    }
  }
  e.len = len;
  np = e.npoints = 2 * e.npoints - 1;
  e.edge_length = dist(dim, &x.data()[0 * dim], &x.data()[(np - 1) * dim]);
}

static void edge_tension_force(std::vector<double> &force, const pedge &e) {
  const std::vector<double> &x = e.x;
  const int dim = e.dim;
  const int np = e.npoints;
  int i, left, right, j;
  double s;

  /* tension force = ((np-1)*||2x-xleft-xright||)/||e||, so the force is nominal and unitless
  */
  s =  (np - 1) / std::max(SMALL, e.edge_length);
  for (i = 1; i <= np - 2; i++){
    left = i - 1;
    right = i + 1;
    for (j = 0; j < dim; j++) force[i*dim + j] += s*(x[left*dim + j] - x[i*dim + j]);
    for (j = 0; j < dim; j++) force[i*dim + j] += s*(x[right*dim + j] - x[i*dim + j]);
  }
}

static void edge_attraction_force(double similarity, const pedge &e1,
                                  const pedge &e2, std::vector<double> &force) {
  /* attractive force from x2 applied to x1 */
  const std::vector<double> &x1 = e1.x, &x2 = e2.x;
  const int dim = e1.dim;
  const int np = e1.npoints;
  int i, j;
  double dist, s, ss;
  const double edge_length = e1.edge_length;

  assert(e1.npoints == e2.npoints);

  /* attractive force = 1/d where d = D/||e1|| is the relative distance, D is the distance between e1 and e2.
   so the force is nominal and unitless
  */
  if (similarity > 0){
    s = edge_length;
    s = similarity*edge_length;
    for (i = 1; i <= np - 2; i++){
      dist = sqr_dist(dim, &x1.data()[i*dim], &x2.data()[i*dim]);
      if (dist < SMALL) dist = SMALL;
      ss = s/(dist+0.1*edge_length*sqrt(dist));
      for (j = 0; j < dim; j++) force[i*dim + j] += ss*(x2[i*dim + j] - x1[i*dim + j]);
    }
  } else {/* clip e2 */
    s = -edge_length;
    s = -similarity*edge_length; 
    for (i = 1; i <= np - 2; i++){
      dist = sqr_dist(dim, &x1.data()[i*dim], &x2.data()[(np - 1 - i)*dim]);
      if (dist < SMALL) dist = SMALL;
      ss = s/(dist+0.1*edge_length*sqrt(dist));
      for (j = 0; j < dim; j++) force[i*dim + j] += ss*(x2[(np - 1 - i)*dim + j] - x1[i*dim + j]);
    }
  }

}

static void force_directed_edge_bundling(SparseMatrix A,
                                         std::vector<pedge> &edges, int maxit,
                                         double step0, double K) {
  int i, j, ne = A->n, k;
  int *ia = A->ia, *ja = A->ja, iter = 0;
  double *a = (double*) A->a;
  const int np = edges[0].npoints, dim = edges[0].dim;
  double step = step0;
  double fnorm_a, fnorm_t, edge_length, start;
  
  if (Verbose > 1)
    fprintf(stderr, "total interaction pairs = %d out of %d, avg neighbors per edge = %f\n",A->nz, A->m*A->m, A->nz/(double) A->m);

  std::vector<double> force_t(dim * np);
  std::vector<double> force_a(dim * np);
  while (step > 0.001 && iter < maxit){
    start = clock();
    iter++;
    for (i = 0; i < ne; i++){
      for (j = 0; j < dim*np; j++) {
	force_t[j] = 0.;
	force_a[j] = 0.;
      }
      pedge &e1 = edges[i];
      std::vector<double> &x = e1.x;
      edge_tension_force(force_t, e1);
      for (j = ia[i]; j < ia[i+1]; j++){
	const pedge &e2 = edges[ja[j]];
	edge_attraction_force(a[j], e1, e2, force_a);
      }
      fnorm_t = std::max(SMALL, norm(dim * (np - 2), &force_t.data()[dim]));
      fnorm_a = std::max(SMALL, norm(dim * (np - 2), &force_a.data()[dim]));
      edge_length = e1.edge_length;

      for (j = 1; j <= np - 2; j++){
	for (k = 0; k < dim; k++) {
	  x[j * dim + k] += step * edge_length
	                  * (force_t[j * dim + k] + K * force_a[j * dim+k])
	                  / hypot(fnorm_t, K * fnorm_a);
	}
      }
      
    }
    step = step*0.9;
  if (Verbose > 1)
    fprintf(stderr, "iter ==== %d cpu = %f npoints = %d\n",iter, ((double) (clock() - start))/CLOCKS_PER_SEC, np - 2);
  }
}

static void modularity_ink_bundling(int dim, int ne, SparseMatrix B,
                                    std::vector<pedge> &edges,
                                    double angle_param, double angle) {
  int *assignment = NULL, nclusters;
  double modularity;
  int *clusterp, *clusters;
  SparseMatrix D, C;
  point_t meet1, meet2;
  double ink0, ink1;
  int i, j, jj;

  SparseMatrix BB;

  /* B may contain negative entries */
  BB = SparseMatrix_copy(B);
  BB = SparseMatrix_apply_fun(BB, fabs);
  modularity_clustering(BB, true, 0, &nclusters, &assignment, &modularity);
  SparseMatrix_delete(BB);

  if (Verbose > 1) fprintf(stderr, "there are %d clusters, modularity = %f\n",nclusters, modularity);
  
  C = SparseMatrix_new(1, 1, 1, MATRIX_TYPE_PATTERN, FORMAT_COORD);
  
  for (i = 0; i < ne; i++){
    jj = assignment[i];
    SparseMatrix_coordinate_form_add_entry(C, jj, i, NULL);
  }
  
  D = SparseMatrix_from_coordinate_format(C);
  SparseMatrix_delete(C);
  clusterp = D->ia;
  clusters = D->ja;
  for (i = 0; i < nclusters; i++) {
    ink1 = ink(edges, clusterp[i + 1] - clusterp[i], &clusters[clusterp[i]],
               &ink0, &meet1, &meet2, angle_param, angle);
    if (Verbose > 1)
      fprintf(stderr,"nedges = %d ink0 = %f, ink1 = %f\n",clusterp[i+1] - clusterp[i], ink0, ink1);
    if (ink1 < ink0){
      for (j = clusterp[i]; j < clusterp[i+1]; j++){
	/* make this edge 5 points, insert two meeting points at 1 and 2, make 3 the last point */
	pedge_double(edges[clusters[j]]);
	pedge_double(edges[clusters[j]]);
	pedge &e = edges[clusters[j]];
	e.x[1 * dim] = meet1.x;
	e.x[1 * dim + 1] = meet1.y;
	e.x[2 * dim] = meet2.x;
	e.x[2 * dim + 1] = meet2.y;
	e.x[3 * dim] = e.x[4 * dim];
	e.x[3 * dim + 1] = e.x[4 * dim + 1];
	e.npoints = 4;
      }
    }
  }
  SparseMatrix_delete(D);
}

static SparseMatrix check_compatibility(SparseMatrix A, int ne,
                                        const std::vector<pedge> &edges,
                                        int compatibility_method, double tol) {
  /* go through the links and make sure edges are compatible */
  SparseMatrix B, C;
  int *ia, *ja, i, j, jj;
  double start;
  double dist;

  B = SparseMatrix_new(1, 1, 1, MATRIX_TYPE_REAL, FORMAT_COORD);
  ia = A->ia; ja = A->ja;
  start = clock();
  for (i = 0; i < ne; i++){
    for (j = ia[i]; j < ia[i+1]; j++){
      jj = ja[j];
      if (i == jj) continue;
      if (compatibility_method == COMPATIBILITY_DIST){
	dist = edge_compatibility_full(edges[i], edges[jj]);
      } else if (compatibility_method == COMPATIBILITY_FULL){
	dist = edge_compatibility(edges[i], edges[jj]);
      } 

      if (fabs(dist) > tol){
	B = SparseMatrix_coordinate_form_add_entry(B, i, jj, &dist);
	B = SparseMatrix_coordinate_form_add_entry(B, jj, i, &dist);
      }
    }
  }
  C = SparseMatrix_from_coordinate_format(B);
  SparseMatrix_delete(B);
  B = C;
  if (Verbose > 1)
    fprintf(stderr, "edge compatibilitu time = %f\n",((double) (clock() - start))/CLOCKS_PER_SEC);
  return B;
}

std::vector<pedge> edge_bundling(SparseMatrix A0, int dim,
                                 const std::vector<double> &x, int maxit_outer,
                                 double K, int method, int nneighbor,
                                 int compatibility_method, int max_recursion,
                                 double angle_param, double angle) {
  /* bundle edges.
     A: edge graph
     x: edge i is at {p,q}, 
     .  where p = x[2*dim*i : 2*dim*i+dim-1]
     .    and q = x[2*dim*i+dim : 2*dim*i+2*dim-1]
     maxit_outer: max outer iteration for force directed bundling. Every outer iter subdivide each edge segment into 2.
     K: nominal edge length in force directed bundling
     method: which method to use.
     nneighbor: number of neighbors to be used in forming nearest neighbor graph. Used only in agglomerative method
     compatibility_method: which method to use to calculate compatibility. Used only in force directed.
     max_recursion: used only in agglomerative method. Specify how many level of recursion to do to bundle bundled edges again

  */
  int ne = A0->m;
  SparseMatrix A = A0, B = NULL;
  int i;
  double tol = 0.001;
  int k;
  double step0 = 0.1, start = 0.0;
  int maxit = 10;

  assert(A->n == ne);
  std::vector<pedge> edges;
  edges.reserve(ne);

  for (i = 0; i < ne; i++){
    edges.emplace_back(pedge_new(2, dim, &x.data()[dim * 2 * i]));
  }

  A = SparseMatrix_symmetrize(A0, true);



  if (Verbose) start = clock();
  if (method == METHOD_INK){

    /* go through the links and make sure edges are compatible */
    B = check_compatibility(A, ne, edges, compatibility_method, tol);

    modularity_ink_bundling(dim, ne, B, edges, angle_param, angle);

  } else if (method == METHOD_INK_AGGLOMERATE){
    /* plan: merge a node with its neighbors if doing so improve. Form coarsening graph, repeat until no more ink saving */
    agglomerative_ink_bundling(dim, A, edges, nneighbor, max_recursion,
                               angle_param, angle);
  } else if (method == METHOD_FD){/* FD method */
    
    /* go through the links and make sure edges are compatible */
    B = check_compatibility(A, ne, edges, compatibility_method, tol);


    for (k = 0; k < maxit_outer; k++){
      for (i = 0; i < ne; i++){
	pedge_double(edges[i]);
      }
      step0 /= 2;
      force_directed_edge_bundling(B, edges, maxit, step0, K);
    }
    
  } else if (method == METHOD_NONE){
    ;
  } else {
    assert(0);
  }
  if (Verbose)
    fprintf(stderr, "total edge bundling cpu = %f\n",((double) (clock() - start))/CLOCKS_PER_SEC);
  if (B != A) SparseMatrix_delete(B);
  if (A != A0) SparseMatrix_delete(A);

  return edges;
}