#include "graph.hpp"
#include <cstring>

#define JC_VORONOI_IMPLEMENTATION
#define JCV_REAL_TYPE double
#define JCV_ATAN2 atan2
#define JCV_FLT_MAX 1.7976931348623157E+308
#undef assert
#define assert_error_report_helper(cond) "assertion failed: "
#define assert(cond) {if(!(cond)) { std::cerr << assert_error_report_helper(cond) "\n"; throw assert_error_report_helper(cond); } }
#include <jc_voronoi.h>

// actual mod for floats
double mod(double a, double b) {
  if (a >= 0) {
    return fmod(a, b);
  } else {
    return fmod(a + b * ceil(-a / b), b);
  }
}

graph::graph(unsigned int Nx, unsigned int Ny) {
  L = {(double)Nx, (double)Ny};

  unsigned int ne = Nx * Ny;
  unsigned int nv = ne / 2;

  vertices.resize(nv);
  edges.reserve(ne);

  dual_vertices.resize(nv);
  dual_edges.reserve(ne);

  for (unsigned int i = 0; i < nv; i++) {
    vertices[i].r.x = (double)((1 + i / (Nx / 2)) % 2 + 2 * (i % (Nx / 2)));
    vertices[i].r.y = (double)(i / (Nx / 2));
    vertices[i].polygon = {
      {vertices[i].r.x - 0.5, vertices[i].r.y},
      {vertices[i].r.x, vertices[i].r.y - 0.5},
      {vertices[i].r.x + 0.5, vertices[i].r.y},
      {vertices[i].r.x, vertices[i].r.y + 0.5}
    };

    dual_vertices[i].r.x = (double)((i / (Nx / 2)) % 2 + 2 * (i % (Nx / 2)));
    dual_vertices[i].r.y = (double)(i / (Nx / 2));
    dual_vertices[i].polygon = {
      {dual_vertices[i].r.x - 0.5, vertices[i].r.y},
      {dual_vertices[i].r.x, vertices[i].r.y - 0.5},
      {dual_vertices[i].r.x + 0.5, vertices[i].r.y},
      {dual_vertices[i].r.x, vertices[i].r.y + 0.5}
    };
  }

  for (unsigned int x = 0; x < Ny; x++) {
    for (unsigned int y = 0; y < Nx; y++) {
      unsigned int v1 = (Nx * x) / 2 + ((y + x % 2) / 2) % (Nx / 2);
      unsigned int v2 = ((Nx * (x + 1)) / 2 + ((y + (x + 1) % 2) / 2) % (Nx / 2)) % nv;

      edges.push_back({{v1, v2}, {0.5 + (double)y, 0.5 + (double)x}});

      unsigned int dv1 = (Nx * x) / 2 + ((y + (x + 1) % 2) / 2) % (Nx / 2);
      unsigned int dv2 = ((Nx * (x + 1)) / 2 + ((y + x % 2) / 2) % (Nx / 2)) % nv;

      dual_edges.push_back({{dv1, dv2}, {0.5 + (double)y, 0.5 + (double)x}});
    }
  }
}

namespace std
{
    template<typename T, size_t N>
    struct hash<array<T, N> >
    {
        typedef array<T, N> argument_type;
        typedef size_t result_type;

        result_type operator()(const argument_type& a) const
        {
            hash<T> hasher;
            result_type h = 0;
            for (result_type i = 0; i < N; ++i)
            {
                h = h * 31 + hasher(a[i]);
            }
            return h;
        }
    };
}

class eulerException: public std::exception
{
  virtual const char* what() const throw()
  {
    return "The voronoi tessellation generated has the incorrect Euler characteristic for a torus and is malformed.";
  }
} eulerex;

graph::graph(double Lx, double Ly, std::mt19937& rng, double relax, double step) {
  L = {Lx, Ly};

  // randomly choose N to be floor(Lx * Ly / 2) or ceil(Lx * Ly / 2) with
  // probability proportional to the distance from each
  std::uniform_real_distribution<double> d(0.0, 1.0);
  unsigned int N = round(Lx * Ly / 2 + d(rng) - 0.5);

  unsigned int nv = N;
  vertices.resize(nv);

  // the coordinates of the lattice, from which a delaunay triangulation
  // and corresponding voronoi tessellation will be built. Everyone is in the
  // rectangle (0, 0) (Lx, Ly)
  for (unsigned int i = 0; i < nv; i++) {
    vertices[i] = {{L.x * d(rng), L.y * d(rng)}};
  }

  // set up the voronoi objects
  jcv_diagram diagram;
  memset(&diagram, 0, sizeof(jcv_diagram));
  jcv_rect bounds = {{-Lx, -Ly}, {2 * Lx, 2 * Ly}};
  std::vector<jcv_point> points(9 * nv);

  double rstep = sqrt(step);

  double max_difference = std::numeric_limits<double>::max();
  while (max_difference > pow(relax, 2) * 2 * N) {
    double cur_diff = 0;
    // to make the resulting tessellation periodic, we tile eight (!) copies of
    // the original points for a total of nine. note that the index of each
    // point quotient 9 is equal to the original index (we'll use this to
    // translate back)
    for (unsigned int i = 0; i < nv; i++) {
      const vertex& v = vertices[i];

      points[9 * i + 0] = {v.r.x + 0.0, v.r.y + 0.0};
      points[9 * i + 1] = {v.r.x + L.x, v.r.y + 0.0};
      points[9 * i + 2] = {v.r.x - L.x, v.r.y + 0.0};
      points[9 * i + 3] = {v.r.x + 0.0, v.r.y + L.y};
      points[9 * i + 4] = {v.r.x + 0.0, v.r.y - L.y};
      points[9 * i + 5] = {v.r.x + L.x, v.r.y + L.y};
      points[9 * i + 6] = {v.r.x - L.x, v.r.y + L.y};
      points[9 * i + 7] = {v.r.x + L.x, v.r.y - L.y};
      points[9 * i + 8] = {v.r.x - L.x, v.r.y - L.y};
    }

    // run voronoi
    jcv_diagram_generate(9 * nv, points.data(), &bounds, &diagram);

    // relax the sites by moving the vertices towards their centroids
    const jcv_site* sites = jcv_diagram_get_sites(&diagram);
    for (int i = 0; i < diagram.numsites; i++) {
      const jcv_site* site = &sites[i];
      unsigned int ind = site->index;
      if (ind % 9 == 0) {
        double Cx = 0.0;
        double Cy = 0.0;
        double A = 0.0;

        const jcv_graphedge* e = site->edges;
        while (e) {
          double dA = e->pos[0].x * e->pos[1].y - e->pos[1].x * e->pos[0].y;

          A += dA;
          Cx += (e->pos[0].x + e->pos[1].x) * dA;
          Cy += (e->pos[0].y + e->pos[1].y) * dA;

          e = e->next;
        }

        A /= 2;
        Cx /= 6 * A;
        Cy /= 6 * A;

        double dx = Cx - vertices[ind / 9].r.x;
        double dy = Cy - vertices[ind / 9].r.y;

        double dist = pow(dx, 2) + pow(dy, 2);
        if (dist > cur_diff) {
          cur_diff = dist;
        }
        vertices[ind / 9] = {{mod(vertices[ind / 9].r.x + rstep * dx, L.x), mod(vertices[ind / 9].r.y + rstep * dy, L.y)}};
      }
    }
    max_difference = cur_diff;
    jcv_diagram_free(&diagram);
    memset(&diagram, 0, sizeof(jcv_diagram));
  }

  for (unsigned int i = 0; i < nv; i++) {
    const vertex& v = vertices[i];
    points[9 * i + 0] = {v.r.x + 0.0, v.r.y + 0.0};
    points[9 * i + 1] = {v.r.x + L.x, v.r.y + 0.0};
    points[9 * i + 2] = {v.r.x - L.x, v.r.y + 0.0};
    points[9 * i + 3] = {v.r.x + 0.0, v.r.y + L.y};
    points[9 * i + 4] = {v.r.x + 0.0, v.r.y - L.y};
    points[9 * i + 5] = {v.r.x + L.x, v.r.y + L.y};
    points[9 * i + 6] = {v.r.x - L.x, v.r.y + L.y};
    points[9 * i + 7] = {v.r.x + L.x, v.r.y - L.y};
    points[9 * i + 8] = {v.r.x - L.x, v.r.y - L.y};
  }
  jcv_diagram_generate(9 * nv, points.data(), &bounds, &diagram);

  const jcv_site* sites = jcv_diagram_get_sites(&diagram);

  struct paircomp {
    bool operator() (const std::array<unsigned int, 2>& lhs, const std::array<unsigned int, 2>& rhs) const
    {if (lhs[0] != lhs[1]) return lhs[0] < lhs[1];
      else return rhs[0] < rhs[1];
    }
  };

  std::unordered_map<std::array<unsigned int, 3>, unsigned int> known_vertices;

  for (int i = 0; i < diagram.numsites; i++) {
    const jcv_site* site = &sites[i];

    // we only care about processing the cells of our original, central sites
    if (site->index % 9 == 0) {
      unsigned int i1 = (unsigned int)(site->index / 9);
      const jcv_graphedge* e = site->edges;
      const jcv_graphedge* ep = site->edges;
      while (ep->next) {
        ep = ep->next;
      }
      // for each edge on the site's cell
      while(e) {
        // assess whether the edge needs to be added. only neighboring
        // sites whose index is larger than ours are given edges, so we
        // don't double up later
        const jcv_site* neighbor = e->neighbor;
        unsigned int i2 = (unsigned int)(neighbor->index / 9);

        vertices[i1].polygon.push_back({e->pos[0].x, e->pos[0].y});

        unsigned int i3p = (unsigned int)(ep->neighbor->index / 9);
        std::array<unsigned int, 3> t1 = {i1, i2, i3p};
        std::sort(t1.begin(), t1.end());

        auto it1 = known_vertices.find(t1);

        unsigned int vi1;

        if (it1 == known_vertices.end()) {
          vi1 = dual_vertices.size();
          dual_vertices.push_back({{mod(e->pos[0].x, L.x), mod(e->pos[0].y, L.y)},{{site->p.x, site->p.y}}});
          known_vertices[t1] = vi1;
        } else {
          vi1 = it1->second;
          dual_vertices[vi1].polygon.push_back({site->p.x, site->p.y});
        }
        
        if (i1 < i2) {
          edges.push_back({{i1, i2},
              {mod((site->p.x + neighbor->p.x) / 2, L.x),
               mod((site->p.y + neighbor->p.y) / 2, L.y)}
               });

          jcv_graphedge *en;
          if (e->next == NULL) {
            en = site->edges;
          } else {
            en = e->next;
          }

          unsigned int i3n = (unsigned int)(en->neighbor->index / 9);
          std::array<unsigned int, 3> t2 = {i1, i2, i3n};
          std::sort(t2.begin(), t2.end());

          auto it2 = known_vertices.find(t2);

          unsigned int vi2;

          if (it2 == known_vertices.end()) {
            vi2 = dual_vertices.size();
            dual_vertices.push_back({{mod(e->pos[1].x, L.x), mod(e->pos[1].y, L.y)},{}});
            known_vertices[t2] = vi2;
          } else {
            vi2 = it2->second;
          }

          dual_edges.push_back({{vi1, vi2},
              {mod((e->pos[0].x + e->pos[1].x) / 2, L.x),
               mod((e->pos[0].y + e->pos[1].y) / 2, L.y)}
               });
        }

        ep = e;
        e = e->next;
      }
    }
  }

  if (edges.size() != vertices.size() + dual_vertices.size()) {
    throw eulerex;
  }

  for (vertex &v : dual_vertices) {
    if (fabs(v.polygon[0].x - v.polygon[1].x) > L.x / 2) {
      if (v.polygon[0].x < L.x / 2) {
        v.polygon[0].x += L.x;
      } else {
        v.polygon[1].x += L.x;
      }
    }

    if (fabs(v.polygon[0].y - v.polygon[1].y) > L.y / 2) {
      if (v.polygon[0].y < L.y / 2) {
        v.polygon[0].y += L.y;
      } else {
        v.polygon[1].y += L.y;
      }
    }

    if (fabs(v.polygon[2].x - v.polygon[1].x) > L.x / 2) {
      if (v.polygon[2].x < L.x / 2) {
        v.polygon[2].x += L.x;
      } else {
        v.polygon[1].x += L.x;
      }
    }

    if (fabs(v.polygon[2].y - v.polygon[1].y) > L.y / 2) {
      if (v.polygon[2].y < L.y / 2) {
        v.polygon[2].y += L.y;
      } else {
        v.polygon[1].y += L.y;
      }
    }

    if (fabs(v.polygon[2].x - v.polygon[0].x) > L.x / 2) {
      if (v.polygon[2].x < L.x / 2) {
        v.polygon[2].x += L.x;
      } else {
        v.polygon[0].x += L.x;
      }
    }

    if (fabs(v.polygon[2].y - v.polygon[0].y) > L.y / 2) {
      if (v.polygon[2].y < L.y / 2) {
        v.polygon[2].y += L.y;
      } else {
        v.polygon[0].y += L.y;
      }
    }
  }

  jcv_diagram_free(&diagram);
}