#include "network.hpp"
#include <sstream>
#include <iostream>

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

network::network(const graph& G, bool two_sides, cholmod_common* c)
    : px(G, 0, c), py(G, 1, c), two_sides(two_sides), G(G),
    C(G.dual_vertices.size()), fuses(G.edges.size()), backbone(G.edges.size()),
    thresholds(G.edges.size()) {}

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;
      }
    }
  }
}

void network::fracture(hooks& m) {
  m.pre_fracture(*this);

  while (true) {
    const double min_cond = 1.0 / G.edges.size();

    current_info c = this->get_current_info();

    if ((c.conductivity[0] < min_cond || c.conductivity[1] < min_cond) && two_sides) {
      break;
    } else if (c.conductivity[0] < min_cond && !two_sides) {
      break;
    }

    if (px.sol.conductivity[0] > min_cond) {
      this->update_backbone(px.sol.currents, c.conductivity);
    } else {
      this->update_backbone(py.sol.currents, c.conductivity);
    }

    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 (!backbone[i]) {
        long double val = logl(c.currents[i]) - thresholds[i];

        if (val > max_val) {
          max_val = val;
          max_pos = i;
        }
      }
    }

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

    m.bond_broken(*this, c, max_pos);
    this->break_edge(max_pos);
  }

  m.post_fracture(*this);
}

void network::get_cycle_edges_helper(std::set<unsigned>& cycle_edges,
    std::set<unsigned>& seen_vertices, unsigned v_prev,
    unsigned v_cur) const {
  seen_vertices.insert(v_cur);
  for (unsigned ei : G.dual_vertices[v_cur].ne) {
    if (fuses[ei]) {
      const std::array<unsigned, 2>& e = G.dual_edges[ei].v;
      unsigned vn = e[0] == v_cur ? e[1] : e[0];

      if (vn != v_prev) {
        if (seen_vertices.contains(vn)) {
          cycle_edges.insert(ei);
        } else {
          this->get_cycle_edges_helper(cycle_edges, seen_vertices, v_cur, vn);
        }
      }
    }
  }
}

std::set<unsigned> network::get_cycle_edges(unsigned v0) const {
  std::set<unsigned> seen_vertices;
  std::set<unsigned> cycle_edges;

  this->get_cycle_edges_helper(cycle_edges, seen_vertices, v0, v0);

  return cycle_edges;
}

bool network::find_cycle_helper(std::array<unsigned, 2>& sig, const std::set<unsigned>& cycle_edges,
    unsigned vPrev, unsigned vCur, unsigned vEnd) const {
  for (unsigned ei : G.dual_vertices[vCur].ne) {
    if (fuses[ei]) {
      if (!cycle_edges.contains(ei)) {
        const std::array<unsigned, 2>& e = G.dual_edges[ei].v;
        unsigned vn = e[0] == vCur ? e[1] : e[0];
        if (vn != vPrev) {
          if (vn == vEnd) {
            if (G.dual_edges[ei].crossings[0])
              sig[0]++;
            if (G.dual_edges[ei].crossings[1])
              sig[1]++;
            return true;
          } else {
            if (this->find_cycle_helper(sig, cycle_edges, vCur, vn, vEnd)) {
              if (G.dual_edges[ei].crossings[0])
                sig[0]++;
              if (G.dual_edges[ei].crossings[1])
                sig[1]++;
              return true;
            }
          }
        }
      }
    }
  }

  return false;
}

std::array<unsigned, 2> network::find_cycle(const std::set<unsigned>& cycle_edges, unsigned v0,
    unsigned v1) const {
  std::array<unsigned, 2> sig = {0, 0};
  this->find_cycle_helper(sig, cycle_edges, v0, v0, v1);
  return sig;
}

bool network::get_cycle_helper(std::array<unsigned, 2>& sig, std::set<unsigned>& edges,
    const std::set<unsigned>& cycle_edges, unsigned vPrev, unsigned vCur,
    unsigned vEnd) const {
  for (unsigned ei : G.dual_vertices[vCur].ne) {
    if (fuses[ei]) {
      if (!cycle_edges.contains(ei)) {
        const std::array<unsigned, 2>& e = G.dual_edges[ei].v;
        unsigned vn = e[0] == vCur ? e[1] : e[0];
        if (vn != vPrev) {
          if (vn == vEnd) {
            edges.insert(ei);
            if (G.dual_edges[ei].crossings[0])
              sig[0]++;
            if (G.dual_edges[ei].crossings[1])
              sig[1]++;
            return true;
          } else {
            if (this->get_cycle_helper(sig, edges, cycle_edges, vCur, vn, vEnd)) {
              edges.insert(ei);
              if (G.dual_edges[ei].crossings[0])
                sig[0]++;
              if (G.dual_edges[ei].crossings[1])
                sig[1]++;
              return true;
            }
          }
        }
      }
    }
  }

  return false;
}

std::pair<std::array<unsigned, 2>, std::set<unsigned>>
network::get_cycle(const std::set<unsigned>& cycle_edges, unsigned v0, unsigned v1) const {
  std::set<unsigned> edges;
  std::array<unsigned, 2> sig = {0, 0};
  this->get_cycle_helper(sig, edges, cycle_edges, v0, v0, v1);
  return {sig, edges};
}

void network::get_cluster_edges_helper(std::set<unsigned>& seen_edges, unsigned v) const {
  for (unsigned ei : G.vertices[v].ne) {
    if (!backbone[ei]) {
      if (!seen_edges.contains(ei)) {
        const std::array<unsigned, 2>& e = G.edges[ei].v;
        unsigned vn = e[0] == v ? e[1] : e[0];

        seen_edges.insert(ei);
        this->get_cluster_edges_helper(seen_edges, vn);
      }
    }
  }
}

std::set<unsigned> network::get_cluster_edges(unsigned v0) const {
  std::set<unsigned> cluster_edges;
  this->get_cluster_edges_helper(cluster_edges, v0);
  return cluster_edges;
}

void network::get_tie_flaps_helper(std::set<unsigned>& added_edges, unsigned v0,
    unsigned vCur) const {
  for (unsigned ei : G.vertices[vCur].ne) {
    if (!backbone[ei]) {
      if (!added_edges.contains(ei)) {
        const std::array<unsigned, 2>& e = G.edges[ei].v;
        unsigned vn = e[0] == vCur ? e[1] : e[0];

        added_edges.insert(ei);
        if (vn != v0) {
          this->get_tie_flaps_helper(added_edges, v0, vn);
        }
      }
    }
  }
}

std::list<std::set<unsigned>> network::get_tie_flaps(unsigned v0) const {
  std::list<std::set<unsigned>> tie_flaps;

  for (unsigned ei : G.vertices[v0].ne) {
    if (!backbone[ei]) {
      bool seen_edge = false;
      for (const std::set<unsigned>& flap : tie_flaps) {
        if (flap.contains(ei)) {
          seen_edge = true;
          break;
        }
      }

      if (!seen_edge) {
        tie_flaps.push_back({ei});

        const std::array<unsigned, 2>& e = G.edges[ei].v;
        unsigned vn = e[0] == v0 ? e[1] : e[0];

        this->get_tie_flaps_helper(tie_flaps.back(), v0, vn);
      }
    }
  }

  return tie_flaps;
}

void network::update_backbone(const std::vector<double>& c, const std::array<double, 2>& cc) {
  bool done_x = cc[0] < 1.0 / G.edges.size();
  bool done_y = cc[1] < 1.0 / G.edges.size();

  auto cycle_check_current = [done_x, done_y](std::array<unsigned, 2> c) {
    return (c[0] % 2 == 0 && c[1] % 2 == 0) ||
      ((c[0] % 2 == 0 && done_x) || (c[1] % 2 == 0 && done_y));
  };
  auto cycle_check_wrap_x = [](std::array<unsigned, 2> c) {
    return (c[0] % 2 == 0 && c[1] % 2 == 1);
  };
  auto cycle_check_wrap_y = [](std::array<unsigned, 2> c) {
    return (c[0] % 2 == 1 && c[1] % 2 == 0);
  };

  // First, we will check for lollipops!
  for (unsigned i = 0; i < G.edges.size(); i++) {
    if ((!backbone[i]) && C.same_component(G.dual_edges[i].v[0], G.dual_edges[i].v[1])) {
      if ((!seen_pairs_x.contains({G.dual_edges[i].v[0], G.dual_edges[i].v[1]}) && !done_x) || (!seen_pairs_y.contains({G.dual_edges[i].v[0], G.dual_edges[i].v[1]}) && !done_y)) {
        // This is a candidate lollipop stem. First, we will identify any
        // cycles in the dual cluster that impinges on the stem and mark them
        // by an edge that uniquely severs each.
        std::set<unsigned> cedges = this->get_cycle_edges(G.dual_edges[i].v[0]);

        // Now, we create a crossing signature for each cycle. First, we do it
        // for the new cycle that would form by removing the stem.
        std::array<unsigned, 2> base_sig =
          this->find_cycle(cedges, G.dual_edges[i].v[0], G.dual_edges[i].v[1]);
        if (G.dual_edges[i].crossings[0])
          base_sig[0]++;
        if (G.dual_edges[i].crossings[1])
          base_sig[1]++;

        // Then, we do it for each of the existing cycles we found in the
        // previous step.
        std::list<std::array<unsigned, 2>> all_sigs = {base_sig};
        for (unsigned e : cedges) {
          std::array<unsigned, 2> new_sig_1 =
            this->find_cycle(cedges, G.dual_edges[i].v[0], G.dual_edges[e].v[0]);
          std::array<unsigned, 2> new_sig_2 =
            this->find_cycle(cedges, G.dual_edges[i].v[1], G.dual_edges[e].v[1]);
          std::array<unsigned, 2> new_sig = {new_sig_1[0] + new_sig_2[0],
            new_sig_1[1] + new_sig_2[1]};

          if (G.dual_edges[i].crossings[0])
            new_sig[0]++;
          if (G.dual_edges[i].crossings[1])
            new_sig[1]++;
          if (G.dual_edges[e].crossings[0])
            new_sig[0]++;
          if (G.dual_edges[e].crossings[1])
            new_sig[1]++;

          all_sigs.push_back(new_sig);
        }

        // Now, having found all cycles involving the candidate stem, we check
        // if any of them spans the torus and therefore can carry current.
        if (std::any_of(all_sigs.begin(), all_sigs.end(), cycle_check_current)) {
          // If none does, we remove it from the backbone!
          seen_pairs_x[{G.dual_edges[i].v[0], G.dual_edges[i].v[1]}] = true;
          seen_pairs_y[{G.dual_edges[i].v[0], G.dual_edges[i].v[1]}] = true;
          backbone[i] = true;

          // We're not done yet: we've severed the stem but the pop remains! We
          // check each side of the lollipop and sum up all of the currents
          // that span an edge. Any component without a sufficiently large net
          // current must be disconnected from the current-carrying cluster and
          // can also be removed from the backbone.
          for (unsigned j = 0; j < 2; j++) {
            std::set<unsigned> cluster_edges = this->get_cluster_edges(G.edges[i].v[j]);
            std::array<double, 2> total_current = {0.0, 0.0};

            for (unsigned e : cluster_edges) {
              if (G.edges[e].crossings[0]) {
                if (G.vertices[G.edges[e].v[0]].r.x < G.vertices[G.edges[e].v[1]].r.x) {
                  total_current[0] += c[e];
                } else {
                  total_current[0] -= c[e];
                }
              }
              if (G.edges[e].crossings[1]) {
                if (G.vertices[G.edges[e].v[0]].r.y < G.vertices[G.edges[e].v[1]].r.y) {
                  total_current[1] += c[e];
                } else {
                  total_current[1] -= c[e];
                }
              }
            }

            if (fabs(total_current[0]) < 1.0 / G.edges.size() &&
                fabs(total_current[1]) < 1.0 / G.edges.size()) {
              for (unsigned e : cluster_edges) {
                backbone[e] = true;
              }
            }
          }
        } else if (std::any_of(all_sigs.begin(), all_sigs.end(), cycle_check_wrap_x)) {
          // If the bond separates a wrapping path, breaking it would end the
          // fracture and therefore it will never be removed from the backbone
          // in this manner. We can skip it in the future.
          seen_pairs_x[{G.dual_edges[i].v[0], G.dual_edges[i].v[1]}] = true;
        } else  if (std::any_of(all_sigs.begin(), all_sigs.end(), cycle_check_wrap_y)) {
          seen_pairs_y[{G.dual_edges[i].v[0], G.dual_edges[i].v[1]}] = true;
        }
      }
    }
  }

  // Now, we will check for bow ties!
  std::set<unsigned> bb_verts;

  // First we get a list of all vertices that remain in the backbone.
  for (unsigned i = 0; i < G.edges.size(); i++) {
    if (!backbone[i]) {
      bb_verts.insert(G.edges[i].v[0]);
      bb_verts.insert(G.edges[i].v[1]);
    }
  }

  for (unsigned v : bb_verts) {
    bool found_pair = false;
    unsigned d1, d2, p1, p2;
    // For each vertex, we check to see if two of the faces adjacent to the
    // vertex are in the same component of the dual lattice and if so, we make
    // sure they do not belong to a contiguously connected series of such
    // faces.
    for (unsigned i = 0; i < G.vertices[v].nd.size(); i++) {
      for (unsigned j = 2; j < G.vertices[v].nd.size() - 1; j++) {
        unsigned l = (i + j) % G.vertices[v].nd.size();
        if (C.same_component(G.vertices[v].nd[i], G.vertices[v].nd[l])) {
          unsigned il = i < l ? i : l;
          unsigned ig = i < l ? l : i;
          if ((!seen_pairs_x.contains({G.vertices[v].nd[il], G.vertices[v].nd[ig]}) && !done_x) || (!seen_pairs_y.contains({G.vertices[v].nd[il], G.vertices[v].nd[ig]}) && !done_y)) {
            bool any_intervening1 = false;
            bool any_intervening2 = false;
            unsigned ie1 = 0;
            unsigned ie2 = 0;

            // yuck, think of something more elegant?
            for (unsigned k = il + 1; k < ig; k++) {
              if (!C.same_component(G.vertices[v].nd[i], G.vertices[v].nd[k])) {
                any_intervening1 = true;
                break;
              }
            }
            for (unsigned k = (ig + 1); k % G.vertices[v].nd.size() != il; k++) {
              if (!C.same_component(G.vertices[v].nd[i],
                                    G.vertices[v].nd[k % G.vertices[v].nd.size()])) {
                any_intervening2 = true;
                break;
              }
            }
            if (any_intervening2 && !any_intervening1) {
              for (unsigned k = il + 1; k <= ig; k++) {
                if (!backbone[G.vertices[v].ne[k]]) {
                  ie1++;
                }
              }
            }
            if (any_intervening1 && !any_intervening2) {
              for (unsigned k = ig + 1; k % G.vertices[v].nd.size() != il + 1; k++) {
                if (!backbone[G.vertices[v].ne[k % G.vertices[v].nd.size()]]) {
                  ie2++;
                }
              }
            }

            // If all these conditions are true, we process the pair. The same
            // cycle analysis as in the lollipop case is done.
            if ((any_intervening1 && any_intervening2) || (any_intervening1 && ie2 > 1) ||
                (any_intervening2 && ie1 > 1)) {
              found_pair = true;
              p1 = il;
              p2 = ig;
              d1 = G.vertices[v].nd[il];
              d2 = G.vertices[v].nd[ig];

              std::set<unsigned> cedges = this->get_cycle_edges(d1);

              std::array<unsigned, 2> base_sig = this->find_cycle(cedges, d1, d2);
              for (unsigned k = p1; k < p2; k++) {
                if (G.dual_edges[G.vertices[v].ne[k]].crossings[0])
                  base_sig[0]++;
                if (G.dual_edges[G.vertices[v].ne[k]].crossings[1])
                  base_sig[1]++;
              }

              std::list<std::array<unsigned, 2>> all_sigs = {base_sig};
              for (unsigned e : cedges) {
                std::array<unsigned, 2> new_sig_1 =
                    this->find_cycle(cedges, d1, G.dual_edges[e].v[0]);
                std::array<unsigned, 2> new_sig_2 =
                    this->find_cycle(cedges, d2, G.dual_edges[e].v[1]);
                std::array<unsigned, 2> new_sig = {new_sig_1[0] + new_sig_2[0],
                                                   new_sig_1[1] + new_sig_2[1]};

                for (unsigned k = p1; k < p2; k++) {
                  if (G.dual_edges[G.vertices[v].ne[k]].crossings[0])
                    new_sig[0]++;
                  if (G.dual_edges[G.vertices[v].ne[k]].crossings[1])
                    new_sig[1]++;
                }
                if (G.dual_edges[e].crossings[0])
                  new_sig[0]++;
                if (G.dual_edges[e].crossings[1])
                  new_sig[1]++;

                all_sigs.push_back(new_sig);
              }

              if (std::any_of(all_sigs.begin(), all_sigs.end(), cycle_check_current)) {
                // one of our pairs qualifies for trimming!
                seen_pairs_x[{d1, d2}] = true;
                seen_pairs_y[{d1, d2}] = true;

                // We separate the flaps of the bowtie (there may be more than
                // two!) into distinct sets.
                std::list<std::set<unsigned>> flaps = this->get_tie_flaps(v);

                // All the bonds in each flap without current are removed from
                // the backbone.
                for (const std::set<unsigned>& flap : flaps) {
                  std::array<double, 2> total_current = {0.0, 0.0};
                  for (unsigned e : flap) {
                    if (G.edges[e].crossings[0]) {
                      if (G.vertices[G.edges[e].v[0]].r.x < G.vertices[G.edges[e].v[1]].r.x) {
                        total_current[0] += c[e];
                      } else {
                        total_current[0] -= c[e];
                      }
                    }
                    if (G.edges[e].crossings[1]) {
                      if (G.vertices[G.edges[e].v[0]].r.y < G.vertices[G.edges[e].v[1]].r.y) {
                        total_current[1] += c[e];
                      } else {
                        total_current[1] -= c[e];
                      }
                    }
                  }

                  if (fabs(total_current[0]) < 1.0 / G.edges.size() &&
                      fabs(total_current[1]) < 1.0 / G.edges.size()) {
                    for (unsigned e : flap) {
                      backbone[e] = true;
                    }
                  }
                }
              } else if (std::any_of(all_sigs.begin(), all_sigs.end(), cycle_check_wrap_x)) {
                seen_pairs_x[{d1, d2}] = true;
              } else if (std::any_of(all_sigs.begin(), all_sigs.end(), cycle_check_wrap_y)) {
                seen_pairs_y[{d1, d2}] = true;
              }
            }
          }
        }
      }
    }
  }
}

void network::break_edge(unsigned e, bool unbreak) {
  fuses[e] = !unbreak;
  backbone[e] = !unbreak;
  C.add_bond(G.dual_edges[e]);
  px.break_edge(e, unbreak);
  py.break_edge(e, unbreak);
}

std::string network::write() {
  std::string output;

  current_info c = this->get_current_info();

  output += "\"fuses\"->{";
  for (unsigned i = 0; i < G.edges.size(); i++) {
    if (!fuses[i]) {
      output += std::to_string(i) + ",";
    }
  }
  output.pop_back();
  output += "},\"backbone\"->{";
  for (unsigned i = 0; i < G.edges.size(); i++) {
    if (!backbone[i]) {
      output += std::to_string(i) + ",";
    }
  }
  output.pop_back();
  output += "},\"thresholds\"->{";
  for (const long double& t : thresholds) {
    output += std::to_string(t) + ",";
  }
  output.pop_back();
  output += "},\"conductivity\"->{" + std::to_string(c.conductivity[0]) + "," +
            std::to_string(c.conductivity[1]);
  output += "},\"currents\"->{";
  for (const double& t : c.currents) {
    output += std::to_string(t) + ",";
  }
  output.pop_back();
  output += "}," + G.write();

  return output;
};

fuse_network::fuse_network(const graph& G, cholmod_common* c) : network(G, false, c), weight(1.0) {}

fuse_network::fuse_network(const fuse_network& n) : network(n), weight(1.0) {}

current_info fuse_network::get_current_info() {
  px.solve(fuses);
  py.solve(fuses);

  bool done_x = px.sol.conductivity[0] < 1.0 / G.edges.size();

  current_info ctot;
  ctot.currents.resize(G.edges.size());
  ctot.conductivity = {px.sol.conductivity[0], py.sol.conductivity[1]};

  if (!done_x) {
    for (unsigned i = 0; i < G.edges.size(); i++) {
      ctot.currents[i] = fabs(px.sol.currents[i] / px.sol.conductivity[0]);
    }
  }

  return ctot;
}

elastic_network::elastic_network(const graph& G, cholmod_common* c, double weight)
    : network(G, true, c), weight(weight) {}

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

current_info elastic_network::get_current_info() {
  px.solve(fuses);
  py.solve(fuses);

  bool done_x = px.sol.conductivity[0] < 1.0 / G.edges.size();
  bool done_y = py.sol.conductivity[1] < 1.0 / G.edges.size();

  current_info ctot;
  ctot.currents.resize(G.edges.size());
  ctot.conductivity[0] = px.sol.conductivity[0];
  ctot.conductivity[1] = py.sol.conductivity[1];

  if (done_x && !done_y) {
    for (unsigned i = 0; i < G.edges.size(); i++) {
      ctot.currents[i] = weight * fabs(py.sol.currents[i]) / py.sol.conductivity[1];
    }
  } else if (done_y && !done_x) {
    for (unsigned i = 0; i < G.edges.size(); i++) {
      ctot.currents[i] = (1 - weight) * fabs(px.sol.currents[i]) / px.sol.conductivity[0];
    }
  } else if (!done_x && !done_y) {
    for (unsigned i = 0; i < G.edges.size(); i++) {
      ctot.currents[i] = sqrt(pow((1 - weight) * px.sol.currents[i] / px.sol.conductivity[0], 2) +
                              pow(weight * py.sol.currents[i] / py.sol.conductivity[1], 2));
    }
  }

  return ctot;
}

percolation_network::percolation_network(const graph& G, cholmod_common* c) : network(G, true, c) {}

percolation_network::percolation_network(const percolation_network& n) : network(n) {}

current_info percolation_network::get_current_info() {
  current_info ctot;
  ctot.currents.resize(G.edges.size(), 0.0);

  px.solve(fuses);
  py.solve(fuses);

  ctot.conductivity = {px.sol.conductivity[0], py.sol.conductivity[1]};

  for (unsigned i = 0; i < G.edges.size(); i++) {
    if (fabs(px.sol.currents[i]) > CURRENT_CUTOFF || fabs(py.sol.currents[i]) > CURRENT_CUTOFF) {
      ctot.currents[i] = 1.0;
    }
  }

  return ctot;
}