#include "network.hpp"

class nanException: public std::exception
{
  virtual const char* what() const throw()
  {
    return "The linear problem returned NaN.";
  }
} nanex;

class nofuseException: public std::exception
{
  virtual const char* what() const throw()
  {
    return "No valid fuse was available to break.";
  }
} nofuseex;

problem::problem(const graph& G, unsigned axis, cholmod_sparse *vcmat, cholmod_common *c) : G(G), voltcurmat(vcmat), c(c), axis(axis) {
  b = CHOL_F(zeros)(G.vertices.size(), 1, CHOLMOD_REAL, c);
  for (unsigned i = 0; i < G.edges.size(); i++) {
    graph::coordinate v0 = G.vertices[G.edges[i].v[0]].r;
    graph::coordinate v1 = G.vertices[G.edges[i].v[1]].r;

    if (G.edges[i].crossings[axis]) {
      bool ind;
      if (axis == 1) {
        ind = v0.y < v1.y ? 0 : 1;
      } else {
        ind = v0.x < v1.x ? 0 : 1;
      }

      ((double *)b->x)[G.edges[i].v[ind]] += 1.0;
      ((double *)b->x)[G.edges[i].v[!ind]] -= 1.0;
    }
  }

  unsigned nnz = G.vertices.size() + G.edges.size();

  cholmod_triplet *t = CHOL_F(allocate_triplet)(G.vertices.size(), G.vertices.size(), nnz,  1, CHOLMOD_REAL, c);

  for (unsigned i = 0; i < G.vertices.size(); i++) {
    ((CHOL_INT *)t->i)[i] = i;
    ((CHOL_INT *)t->j)[i] = i;
    ((double *)t->x)[i] = 0.0;
  }

  unsigned terms = G.vertices.size();

  std::unordered_map<std::array<unsigned, 2>, unsigned> known_edges;

  for (unsigned i = 0; i < G.edges.size(); i++) {
    unsigned v0 = G.edges[i].v[0];
    unsigned v1 = G.edges[i].v[1];

    ((double *)t->x)[v0]++;
    ((double *)t->x)[v1]++;

    unsigned s0 = v0 < v1 ? v0 : v1;
    unsigned s1 = v0 < v1 ? v1 : v0;

    auto it = known_edges.find({s0, s1});

    if (it == known_edges.end()) {
      ((CHOL_INT *)t->i)[terms] = s1;
      ((CHOL_INT *)t->j)[terms] = s0;
      ((double *)t->x)[terms] = -1.0;

      known_edges[{s0, s1}] = terms;
      terms++;
    } else {
      ((double *)t->x)[it->second] -= 1.0;
    }
  }

  ((double *)t->x)[0]++;

  t->nnz = terms;

  cholmod_sparse *laplacian = CHOL_F(triplet_to_sparse)(t, terms, c);
  CHOL_F(free_triplet)(&t, c);
  factor = CHOL_F(analyze)(laplacian, c);
  CHOL_F(factorize)(laplacian, factor, c);
  CHOL_F(free_sparse)(&laplacian, c);
}

problem::problem(const graph& G, unsigned axis, cholmod_common *c) : problem(G, axis, NULL, c) {
  cholmod_triplet *t = CHOL_F(allocate_triplet)(G.edges.size(), G.vertices.size(), 2 * G.edges.size(), 0, CHOLMOD_REAL, c);

  t->nnz = 2 * G.edges.size();

  for (unsigned i = 0; i < G.edges.size(); i++) {
    ((CHOL_INT *)t->i)[2 * i] = i;
    ((CHOL_INT *)t->j)[2 * i] = G.edges[i].v[0];
    ((double *)t->x)[2 * i] = 1.0;

    ((CHOL_INT *)t->i)[2 * i + 1] = i;
    ((CHOL_INT *)t->j)[2 * i + 1] = G.edges[i].v[1];
    ((double *)t->x)[2 * i + 1] = -1.0;
  }

  voltcurmat = CHOL_F(triplet_to_sparse)(t, 2 * G.edges.size(), c);

  CHOL_F(free_triplet)(&t, c);
}

problem::problem(const problem& other) : G(other.G), c(other.c), axis(other.axis) {
  b = CHOL_F(copy_dense)(other.b, c);
  factor = CHOL_F(copy_factor)(other.factor, c);
  voltcurmat = CHOL_F(copy_sparse)(other.voltcurmat, c);
} 

problem::~problem() {
  CHOL_F(free_dense)(&b, c);
  CHOL_F(free_factor)(&factor, c);
  CHOL_F(free_sparse)(&voltcurmat, c);
}

current_info problem::solve(std::vector<bool>& fuses) {
  cholmod_dense *x = CHOL_F(solve)(CHOLMOD_A, factor, b, c);

  if (((double *)x->x)[0] != ((double *)x->x)[0]) {
    throw nanex;
  }

  cholmod_dense *y = CHOL_F(allocate_dense)(G.edges.size(), 1, G.edges.size(), CHOLMOD_REAL, c);

  double alpha[2] = {1, 0};
  double beta[2] = {0, 0};
  CHOL_F(sdmult)(voltcurmat, 0, alpha, beta, x, y, c);

  std::vector<double> currents(G.edges.size());

  double total_current = 0;

  for (int i = 0; i < G.edges.size(); i++) {
    if (fuses[i]) {
      currents[i] = 0;
    } else {
      currents[i] = ((double *)y->x)[i];

      graph::coordinate v0 = G.vertices[G.edges[i].v[0]].r;
      graph::coordinate v1 = G.vertices[G.edges[i].v[1]].r;

      if (G.edges[i].crossings[axis]) {
        bool comp;
        if (axis == 1) {
          comp = v0.y > v1.y;
        } else {
          comp = v0.x > v1.x;
        }

        if (comp) {
          currents[i] += 1.0;
          total_current += currents[i];
        } else {
          currents[i] -= 1.0;
          total_current -= currents[i];
        }
      }
    }
  }

  CHOL_F(free_dense)(&x, c);
  CHOL_F(free_dense)(&y, c);

  return {total_current, currents};
}

void problem::break_edge(unsigned e, bool unbreak) {
  unsigned v0 = G.edges[e].v[0];
  unsigned v1 = G.edges[e].v[1];

  unsigned n = factor->n;

  cholmod_sparse *update_mat = CHOL_F(allocate_sparse)(n, n, 2, true, true, 0, CHOLMOD_REAL, c);

  unsigned s1, s2;
  s1 = v0 < v1 ? v0 : v1;
  s2 = v0 < v1 ? v1 : v0;

  CHOL_INT *pp = (CHOL_INT *)update_mat->p;
  CHOL_INT *ii = (CHOL_INT *)update_mat->i;
  double *xx = (double *)update_mat->x;

  for (unsigned i = 0; i <= s1; i++) {
    pp[i] = 0;
  }

  for (unsigned i = s1 + 1; i <= n; i++) {
    pp[i] = 2;
  }

  ii[0] = s1;
  ii[1] = s2;
  xx[0] = 1;
  xx[1] = -1;

  cholmod_sparse *perm_update_mat = CHOL_F(submatrix)(
      update_mat, (CHOL_INT *)factor->Perm, factor->n, NULL, -1, true, true, c);

  CHOL_F(updown)(unbreak, perm_update_mat, factor, c);

  CHOL_F(free_sparse)(&perm_update_mat, c);
  CHOL_F(free_sparse)(&update_mat, c);

  graph::coordinate r0 = G.vertices[v0].r;
  graph::coordinate r1 = G.vertices[v1].r;

  if (G.edges[e].crossings[axis]) {
    bool ind;
    if (axis == 1) {
      ind = r0.y < r1.y ? unbreak : !unbreak;
    } else {
      ind = r0.x < r1.x ? unbreak : !unbreak;
    }

    ((double *)b->x)[G.edges[e].v[ind]] -= 1.0;
    ((double *)b->x)[G.edges[e].v[!ind]] += 1.0;
  }
}

network::network(const graph& G) : G(G), fuses(G.edges.size()), thresholds(G.edges.size()) {}

network::network(const network& n) : G(n.G), fuses(n.fuses), thresholds(n.thresholds) {}

void network::set_thresholds(double beta, std::mt19937& rng) {
  if (beta == 0.0) {
    /* zero beta doesn't make any sense computationally, we interpret it as "infinity" */
    for (long double& threshold : thresholds) {
      threshold = 1.0;
    }
  } else {
    std::uniform_real_distribution<long double> d(0.0, 1.0);

    for (long double& threshold : thresholds) {
      threshold = std::numeric_limits<long double>::lowest();

      while (threshold == std::numeric_limits<long double>::lowest()) {
        threshold = logl(d(rng)) / (long double)beta;
      }
    }
  }
}

fuse_network::fuse_network(const graph& G, cholmod_common *c) : network(G), ohm(G, 1, c) {
}

void fuse_network::fracture(hooks& m, double cutoff) {
  m.pre_fracture(*this);

  while (true) {
    current_info ci = ohm.solve(fuses);

    if (ci.conductivity < cutoff * G.vertices.size()) {
      break;
    }

    unsigned max_pos = UINT_MAX;
    long double max_val = std::numeric_limits<long double>::lowest();

    for (unsigned i = 0; i < G.edges.size(); i++) {
      if (!fuses[i] && fabs(ci.currents[i]) > cutoff) {
        long double val = logl(fabs(ci.currents[i])) - thresholds[i];
        if (val > max_val) {
          max_val = val;
          max_pos = i;
        }
      }
    }

    if (max_pos == UINT_MAX)  {
      throw nofuseex;
    }

    fuses[max_pos] = true;
    ohm.break_edge(max_pos);

    m.bond_broken(*this, ci, max_pos);
  }

  m.post_fracture(*this);
}

elastic_network::elastic_network(const graph& G, cholmod_common *c) : network(G), hook_x(G, 1, c), hook_y(G, 0, c) {
}

elastic_network::elastic_network(const elastic_network& n) : network(n), hook_x(n.hook_x), hook_y(n.hook_y) {
}

void elastic_network::fracture(hooks& m, double weight, double cutoff) {
  m.pre_fracture(*this);

  while (true) {
    current_info cx = hook_x.solve(fuses);
    current_info cy = hook_y.solve(fuses);

    current_info ctot;

    ctot.conductivity = sqrt(pow((1 - weight) * cx.conductivity, 2) + pow(weight * cy.conductivity, 2));

    if (ctot.conductivity < cutoff * G.vertices.size()) {
      break;
    }

    ctot.currents.resize(G.edges.size());
    for (unsigned i = 0; i < G.edges.size(); i++) {
      ctot.currents[i] = sqrt(pow((1 - weight) * cx.currents[i], 2) + pow(weight * cy.currents[i], 2));
    }

    unsigned max_pos = UINT_MAX;
    long double max_val = std::numeric_limits<long double>::lowest();

    for (unsigned i = 0; i < G.edges.size(); i++) {
      if (!fuses[i] && fabs(ctot.currents[i]) > cutoff) {
        long double val = logl(fabs(ctot.currents[i])) - thresholds[i];
        if (val > max_val) {
          max_val = val;
          max_pos = i;
        }
      }
    }

    if (max_pos == UINT_MAX)  {
      throw nofuseex;
    }

    fuses[max_pos] = true;
    hook_x.break_edge(max_pos);
    hook_y.break_edge(max_pos);

    m.bond_broken(*this, ctot, max_pos);
  }

  m.post_fracture(*this);
}