#include "measurements.hpp"
#include <iostream>

template <class T>
bool is_shorter(std::list<T> l1, std::list<T> l2) {
  return l1.size() < l2.size();
}

bool trivial(boost::detail::edge_desc_impl<boost::undirected_tag,unsigned long>) {
  return true;
}

std::list<unsigned int> find_minimal_crack(const Graph& G, const network& n) {
  Graph Gtmp(n.G.vertices.size());
  std::list<unsigned int> removed_edges;

  class add_tree_edges : public boost::default_dfs_visitor {
    public:
      Graph& G;
      std::list<unsigned int>& E;

      add_tree_edges(Graph& G, std::list<unsigned int>& E) : G(G), E(E) {}

      void tree_edge(boost::graph_traits<Graph>::edge_descriptor e, const Graph& g) {
        boost::add_edge(boost::source(e, g), boost::target(e, g), g[e], G);
      }

      void back_edge(boost::graph_traits<Graph>::edge_descriptor e, const Graph& g) {
        if (!(boost::edge(boost::source(e, g), boost::target(e, g), G).second)) {
          E.push_back(g[e].index);
        }
      }
  };

  add_tree_edges ate(Gtmp, removed_edges);
  boost::depth_first_search(G, visitor(ate));

  class find_cycle : public boost::default_dfs_visitor {
    public:
      std::list<unsigned int>& E;
      unsigned int end;
      struct done{};

      find_cycle(std::list<unsigned int>& E, unsigned int end) : E(E), end(end) {}

      void discover_vertex(boost::graph_traits<Graph>::vertex_descriptor v, const Graph& g) {
        if (v == end) {
          throw done{};
        }
      }

      void examine_edge(boost::graph_traits<Graph>::edge_descriptor e, const Graph& g) {
        E.push_back(g[e].index);
      }

      void finish_edge(boost::graph_traits<Graph>::edge_descriptor e, const Graph& g) {
        E.erase(std::find(E.begin(), E.end(), g[e].index));
      }
  };

  std::list<std::list<unsigned int>> cycles;

  for (auto edge : removed_edges) {
    std::list<unsigned int> cycle = {edge};
    find_cycle vis(cycle, n.G.dual_edges[edge][1]);
    std::vector<boost::default_color_type> new_color_map(boost::num_vertices(Gtmp));
    try {
    boost::depth_first_visit(Gtmp, n.G.dual_edges[edge][0], vis, boost::make_iterator_property_map(new_color_map.begin(), boost::get(boost::vertex_index, Gtmp), new_color_map[0]));
    } catch(find_cycle::done const&) {
      cycles.push_back(cycle);
    }
  }

  if (cycles.size() > 1) {
    std::list<std::valarray<uint8_t>> bool_cycles;
    for (auto cycle : cycles) {
      std::valarray<uint8_t> bool_cycle(n.G.edges.size());
      for (auto v : cycle) {
        bool_cycle[v] = 1;
      }
      bool_cycles.push_back(bool_cycle);
    }

    // generate all possible cycles by taking xor of the edge sets of the known cycles
    for (auto it1 = bool_cycles.begin(); it1 != std::prev(bool_cycles.end()); it1++) {
      for (auto it2 = std::next(it1); it2 != bool_cycles.end(); it2++) {
        std::valarray<uint8_t> new_bool_cycle = (*it1) ^ (*it2);
        std::list<unsigned int> new_cycle;
        unsigned int pos = 0;
        for (uint8_t included : new_bool_cycle) {
          if (included) {
            new_cycle.push_back(pos);
          }
          pos++;
        }
        cycles.push_back(new_cycle);
      }
    }

    // find the cycle representing the crack by counting boundary crossings
    for (auto cycle : cycles) {
      std::array<unsigned int, 2> crossing_count{0,0};

      for (auto edge : cycle) {
        double dx = fabs(n.G.vertices[n.G.dual_edges[edge][0]].r.x - n.G.vertices[n.G.dual_edges[edge][1]].r.x);
        if (dx > n.G.L.x / 2) {
          crossing_count[0]++;
        }
        double dy = fabs(n.G.vertices[n.G.dual_edges[edge][0]].r.y - n.G.vertices[n.G.dual_edges[edge][1]].r.y);
        if (dy > n.G.L.y / 2) {
          crossing_count[1]++;
        }
      }

      if (crossing_count[0] % 2 == 1 && crossing_count[1] % 2 == 0) {
        return cycle;
      }
    }
  } else {
    return cycles.front();
  }

  exit(5);
}

void update_distribution_file(std::string id, const std::vector<uint64_t>& data, unsigned int N, unsigned int Lx, unsigned int Ly, double beta) {
  std::string filename = "fracture_" + std::to_string(Lx) + "_" + std::to_string(Ly) + "_" + std::to_string(beta) + "_" + id + ".dat";
  std::ifstream file(filename);

  uint64_t N_old = 0;
  std::vector<uint64_t> data_old(data.size(), 0);

  if (file.is_open()) {
    file >> N_old;
    for (unsigned int i = 0; i < data.size(); i++) {
      uint64_t num;
      file >> num;
      data_old[i] = num;
    }

    file.close();
  }

  std::ofstream file_out(filename);

  file_out <<std::fixed<< N_old + N << "\n";
  for (unsigned int i = 0; i < data.size(); i++) {
    file_out <<std::fixed<< data_old[i] + data[i] << " ";
  }

  file_out.close();
}

template <class T>
void update_field_file(std::string id, const std::vector<T>& data, unsigned int N, unsigned int Lx, unsigned int Ly, double beta) {
  std::string filename = "fracture_" + std::to_string(Lx) + "_" + std::to_string(Ly) + "_" + std::to_string(beta) + "_" + id + ".dat";
  std::ifstream file(filename);

  uint64_t N_old = 0;
  std::vector<uint64_t> data_old(data.size(), 0);

  if (file.is_open()) {
    file >> N_old;
    for (unsigned int i = 0; i < data.size(); i++) {
      file >> data_old[i];
    }
    file.close();
  }

  std::ofstream file_out(filename);

  file_out <<std::fixed<< N_old + N <<  "\n";
  for (unsigned int i = 0; i < data.size(); i++) {
    file_out << data_old[i] + data[i] << " ";
  }

  file_out.close();
}

template <class T>
std::vector<fftw_complex> data_transform(unsigned int Lx, unsigned int Ly, const std::vector<T>& data, fftw_plan forward_plan, double *fftw_forward_in, fftw_complex *fftw_forward_out) {
  for (unsigned int i = 0; i < Lx * Ly; i++) {
    fftw_forward_in[i] = (double)data[i];
  }

  fftw_execute(forward_plan);

  std::vector<fftw_complex> output(Lx * (Ly / 2 + 1));

  for (unsigned int i = 0; i < Lx * (Ly / 2 + 1); i++) {
    output[i][0] = fftw_forward_out[i][0];
    output[i][1] = fftw_forward_out[i][1];
  }

  return output;
}

template <class T>
void correlation(unsigned int Lx, unsigned int Ly, std::vector<T>& data, const std::vector<fftw_complex>& tx1, const std::vector<fftw_complex>& tx2, fftw_plan reverse_plan, fftw_complex *fftw_reverse_in, double *fftw_reverse_out) {
  for (unsigned int i = 0; i < Lx * (Ly / 2 + 1); i++) {
    fftw_reverse_in[i][0] = tx1[i][0] * tx2[i][0] + tx1[i][1] * tx2[i][1];
    fftw_reverse_in[i][1] = tx1[i][0] * tx2[i][1] - tx1[i][1] * tx2[i][0];
  }

  fftw_execute(reverse_plan);

  for (unsigned int i = 0; i < (Lx / 2 + 1) * (Ly / 2 + 1); i++) {
    data[i] += (T)(fftw_reverse_out[Lx * (i / (Lx / 2 + 1)) + i % (Lx / 2 + 1)] / (Lx * Ly));
  }
}

template <class T>
void autocorrelation(unsigned int Lx, unsigned int Ly, std::vector<T>& out_data, fftw_plan forward_plan, double *fftw_forward_in, fftw_complex *fftw_forward_out, fftw_plan reverse_plan, fftw_complex *fftw_reverse_in, double *fftw_reverse_out) {
  fftw_execute(forward_plan);

  for (unsigned int i = 0; i < Ly * (Lx / 2 + 1); i++) {
    fftw_reverse_in[i][0] = pow(fftw_forward_out[i][0], 2) + pow(fftw_forward_out[i][1], 2);
    fftw_reverse_in[i][1] = 0.0;
  }

  fftw_execute(reverse_plan);

  for (unsigned int i = 0; i < (Lx / 2 + 1) * (Ly / 2 + 1); i++) {
    out_data[i] += (T)(fftw_reverse_out[Lx * (i / (Lx / 2 + 1)) + i % (Lx / 2 + 1)] / (Lx * Ly));
  }
}

ma::ma(unsigned int Lx, unsigned int Ly, double beta) : Lx(Lx), Ly(Ly), beta(beta), G(Lx * Ly / 2),
  sc(Lx * Ly, 0),
  sa(Lx * Ly, 0),
  sd(Lx * Ly, 0),
  sb(log2(Ly) + 1, 0),
  Ccc((Lx / 2 + 1) * (Ly / 2 + 1), 0),
  Css((Lx / 2 + 1) * (Ly / 2 + 1), 0),
  Cmm((Lx / 2 + 1) * (Ly / 2 + 1), 0),
  Caa((Lx / 2 + 1) * (Ly / 2 + 1), 0),
  Cdd((Lx / 2 + 1) * (Ly / 2 + 1), 0),
  Cll((Lx / 2 + 1) * (Ly / 2 + 1), 0),
  Cee((Lx / 2 + 1) * (Ly / 2 + 1), 0)
{
  N = 0;
  Nc = 0;
  Na = 0;

  // FFTW setup for correlation functions
  fftw_set_timelimit(1);

  fftw_forward_in = (double *)fftw_malloc(Lx * Ly * sizeof(double));
  fftw_forward_out = (fftw_complex *)fftw_malloc(Lx * Ly * sizeof(fftw_complex));
  fftw_reverse_in = (fftw_complex *)fftw_malloc(Lx * Ly * sizeof(fftw_complex));
  fftw_reverse_out = (double *)fftw_malloc(Lx * Ly * sizeof(double));

  forward_plan = fftw_plan_dft_r2c_2d(Ly, Lx, fftw_forward_in, fftw_forward_out, 0);
  reverse_plan = fftw_plan_dft_c2r_2d(Ly, Lx, fftw_reverse_in, fftw_reverse_out, 0);
}

ma::~ma() {
  // clean up FFTW objects
  fftw_free(fftw_forward_in);
  fftw_free(fftw_forward_out);
  fftw_free(fftw_reverse_in);
  fftw_free(fftw_reverse_out);
  fftw_destroy_plan(forward_plan);
  fftw_destroy_plan(reverse_plan);
  fftw_cleanup();

  update_distribution_file("sa", sa, Na, Lx, Ly, beta);
  update_distribution_file("sc", sc, Nc, Lx, Ly, beta);
  update_distribution_file("sd", sd, N, Lx, Ly, beta);
  update_distribution_file("sb", sb, N, Lx, Ly, beta);

  update_field_file("Ccc", Ccc, Nc, Lx, Ly, beta);
  update_field_file("Css", Css, N, Lx, Ly, beta);
  update_field_file("Cmm", Cmm, N, Lx, Ly, beta);
  update_field_file("Cdd", Cdd, N, Lx, Ly, beta);
  update_field_file("Caa", Caa, Na, Lx, Ly, beta);
  update_field_file("Cll", Cll, N, Lx, Ly, beta);
  update_field_file("Cee", Cee, N, Lx, Ly, beta);
}

void ma::pre_fracture(const network &) {
  lv = 0;
  avalanches = {{}};
  boost::remove_edge_if(trivial, G);
}

void ma::bond_broken(const network& net, const current_info& cur, unsigned int i) {
  double c = cur.conductivity / fabs(cur.currents[i]) * net.thresholds[i];
  if (c > lv && avalanches.back().size() > 0) {
    sa[avalanches.back().size() - 1]++;
    Na++;

    std::fill_n(fftw_forward_in, net.G.edges.size(), 0.0);

    for (auto e : avalanches.back()) {
      fftw_forward_in[e] = 1.0;
    }

    autocorrelation(Lx, Ly, Caa, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

    lv = c;
    avalanches.push_back({i});
  } else {
    avalanches.back().push_back(i);
  }

  boost::add_edge(net.G.dual_edges[i][0], net.G.dual_edges[i][1], {i}, G);
}

void ma::post_fracture(network &n) {
  std::vector<unsigned int> component(boost::num_vertices(G));
  unsigned int num = boost::connected_components(G, &component[0]);

  std::list<unsigned int> crack = find_minimal_crack(G, n);

  unsigned int crack_component = component[n.G.dual_edges[crack.front()][0]];

  // non-spanning clusters
  for (unsigned int i = 0; i < num; i++) {
    if (i != crack_component) {
      unsigned int size = 0;

      for (unsigned int j = 0; j < n.G.edges.size(); j++) {
        if (component[n.G.dual_edges[j][0]] == i && n.fuses[j]) {
          size++;
          fftw_forward_in[j] = 1.0;
        } else{
          fftw_forward_in[j] = 0.0;
        }
      }

      if (size > 0) {
        sc[size - 1]++;
        autocorrelation(Lx, Ly, Ccc, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);
        Nc++;
      }
    }
  }

  // bin counting
  for (unsigned int be = 0; be < log2(Ly); be++) {
    unsigned int bin = pow(2, be);

    for (unsigned int i = 0; i < Lx * Ly / pow(bin, 2); i++) {
      bool in_bin = false;
      for (unsigned int j = 0; j < pow(bin, 2); j++) {
        unsigned int edge = Lx * (bin * ((i * bin) / Lx) + j / bin) + (i * bin) % Lx + j % bin;
        if (!n.fuses[edge] && crack_component == component[n.G.dual_edges[edge][0]]) {
          in_bin = true;
          break;
        }
      }

      if (in_bin) {
        sb[be]++;
      }
    }
  }

  sb[log2(Ly)]++;

  for (unsigned int i = 0; i < n.G.edges.size(); i++) {
    if (n.fuses[i] && component[n.G.dual_edges[i][0]] == crack_component)  {
      fftw_forward_in[i] = 1.0;
    } else {
      fftw_forward_in[i] = 0.0;
    }
  }

  autocorrelation(Lx, Ly, Cmm, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  // crack surface correlations
  std::fill_n(fftw_forward_in, n.G.edges.size(), 0.0);

  for (auto edge : crack) {
    fftw_forward_in[edge] = 1.0;
  }

  autocorrelation(Lx, Ly, Css, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  std::function<bool(unsigned int)> inCrack = [&](unsigned int i) -> bool {
    return (bool)fftw_forward_in[i];
  };

  for (auto avalanche : avalanches) {
    if (avalanche.end() != std::find_if(avalanche.begin(), avalanche.end(), inCrack)) {
      for (auto edge : avalanche) {
        fftw_forward_in[edge] = 1.0;
      }
    }
  }

  autocorrelation(Lx, Ly, Cee, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  std::fill_n(fftw_forward_in, n.G.edges.size(), 0.0);

  // rewind the last avalanche
  for (auto e : avalanches.back()) {
    boost::remove_edge(n.G.dual_edges[e][0], n.G.dual_edges[e][1], G);
    n.break_edge(e, true);
    fftw_forward_in[e] = 1;
  }

  autocorrelation(Lx, Ly, Cll, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  // damage size distribution
  unsigned int total_broken = 0;

  for (unsigned int i = 0; i < n.G.edges.size(); i++) {
    if (n.fuses[i]) { 
      total_broken++;
      fftw_forward_in[i] = 1.0;
    } else {
      fftw_forward_in[i] = 0.0;
    }
  }

  autocorrelation(Lx, Ly, Cdd, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  sd[total_broken]++;

  N++;
}