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

void update_distribution_file(std::string id, const std::vector<uint64_t>& data, unsigned int N, unsigned int L, double beta) {
  std::string filename = "fracture_" + std::to_string(L) + "_" + 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 L, double beta) {
  std::string filename = "fracture_" + std::to_string(L) + "_" + 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 L, const std::vector<T>& data, fftw_plan forward_plan, double *fftw_forward_in, fftw_complex *fftw_forward_out) {
  for (unsigned int i = 0; i < pow(L, 2); i++) {
    fftw_forward_in[i] = (double)data[i];
  }

  fftw_execute(forward_plan);

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

  for (unsigned int i = 0; i < L * (L / 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 L, 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 < L * (L / 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 < pow(L / 2 + 1, 2); i++) {
    data[i] += (T)(fftw_reverse_out[L * (i / (L / 2 + 1)) + i % (L / 2 + 1)] / pow(L, 2));
  }
}

template <class T, class U>
void autocorrelation(unsigned int L, std::vector<T>& out_data, const std::vector<U>& 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) {
  for (unsigned int i = 0; i < pow(L, 2); i++) {
    fftw_forward_in[i] = (double)data[i];
  }

  fftw_execute(forward_plan);

  for (unsigned int i = 0; i < L * (L / 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 < pow(L / 2 + 1, 2); i++) {
    out_data[i] += (T)(fftw_reverse_out[L * (i / (L / 2 + 1)) + i % (L / 2 + 1)] / pow(L, 2));
  }
}

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

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

  fftw_forward_in = (double *)fftw_malloc(pow(L, 2) * sizeof(double));
  fftw_forward_out = (fftw_complex *)fftw_malloc(pow(L, 2) * sizeof(fftw_complex));
  fftw_reverse_in = (fftw_complex *)fftw_malloc(pow(L, 2) * sizeof(fftw_complex));
  fftw_reverse_out = (double *)fftw_malloc(pow(L, 2) * sizeof(double));

  forward_plan = fftw_plan_dft_r2c_2d(L, L, fftw_forward_in, fftw_forward_out, 0);
  reverse_plan = fftw_plan_dft_c2r_2d(L, L, 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, N, L, beta);
  update_distribution_file("sc", sc, N, L, beta);
  update_distribution_file("sd", sd, N, L, beta);
  update_distribution_file("sb", sb, N, L, beta);

  update_field_file("Ccc", Ccc, Nc, L, beta);
  update_field_file("Css", Css, N, L, beta);
  update_field_file("Cdd", Cdd, N, L, beta);
  update_field_file("Caa", Caa, Na, L, beta);

}

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

void ma::bond_broken(const network& net, const std::pair<double, std::vector<double>>& cur, unsigned int i) {
  if (cur.first / fabs(cur.second[i]) * net.thresholds[i] > lv) {
    sa[avalanche.size()]++;
    Na++;

    std::vector<bool>avalanche_damage(pow(L, 2), 0);

    for (auto it = avalanche.begin(); it != avalanche.end(); it++) {
      avalanche_damage[*it] = 1;
    }

    autocorrelation(L, Caa, avalanche_damage, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

    lv = cur.first / fabs(cur.second[i]) * net.thresholds[i];
    avalanche.clear();
  } else {
    avalanche.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::vector<unsigned int> comp_sizes(num, 0);

  for (unsigned int c : component) {
    comp_sizes[c]++;
  }

  unsigned int max_i = 0;
  unsigned int max_size = 0;

  for (unsigned int i = 0; i < num; i++) {
    if (comp_sizes[i] > max_size) {
      max_i = i;
      max_size = comp_sizes[i];
    }
  }

  unsigned int vertex_in_largest = 0;

  while (component[vertex_in_largest] != max_i) {
    vertex_in_largest++;
  }

  Graph Gtmp(2 * pow(L / 2, 2));
  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:
      unsigned int end;
      std::list<unsigned int>& V;
      std::list<unsigned int>& E;
      struct done{};

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

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

      void finish_vertex(boost::graph_traits<Graph>::vertex_descriptor v, const Graph& g) {
        V.remove(v);
      }

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

  if (removed_edges.size() == 0) {
    std::pair<double, std::vector<double>> currents = n.currents();
    std::cout << currents.first << "\n";
    for (unsigned int i = 0; i < L; i++) {
      for (unsigned int j = 0; j < L; j++) {
        std::cout << n.fuses[i * L + j] << " ";
      }
      std::cout << "\n";
    }
  }

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

  for (auto edge : removed_edges) {
    std::list<unsigned int> cycle = {edge};
    std::list<unsigned int> vcycle = {};
    find_cycle vis(vcycle, 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);
    }
  }

  std::list<unsigned int> crack;

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

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

    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.edges[edge][0]].r.x - n.G.vertices[n.G.edges[edge][1]].r.x);
        if (dx > 0.5) {
          crossing_count[0]++;
        }
        double dy = fabs(n.G.vertices[n.G.edges[edge][0]].r.y - n.G.vertices[n.G.edges[edge][1]].r.y);
        if (dy > 0.5) {
          crossing_count[1]++;
        }
      }

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

  for (auto edge : crack) {
    std::cout << edge << " ";
  }
  std::cout << "\n";
//  std::cout << max_size << " " << boost::num_edges(Gtmp) << " " << removed_edges.size() << "\n";

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

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

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

  sb[log2(L)]++;

  // crack surface correlations
  std::vector<double> crack_damage(pow(L, 2), 0.0);

  double damage_tot = 0;
  for (unsigned int i = 0; i < pow(L, 2); i++) {
    if (!n.fuses[i] && max_i == component[n.G.dual_edges[i][0]]) {
      damage_tot++;
      crack_damage[i] = 1.0;
    }
  }

  std::vector<fftw_complex> t_crack_damage = data_transform(L, crack_damage, forward_plan, fftw_forward_in, fftw_forward_out);
  correlation(L, Css, t_crack_damage, t_crack_damage, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  std::vector<bool> last_crack(pow(L, 2), 0);

  // rewind the last avalanche
  for (auto e = avalanche.rbegin(); e != avalanche.rend(); e++) {
    boost::remove_edge(n.G.dual_edges[*e][0], n.G.dual_edges[*e][1], G);
    n.add_edge(*e);
    last_crack[*e] = 1;
  }


  // cluster size distribution and cluster-cluster correlation
  num = boost::connected_components(G, &component[0]);

  for (unsigned int i = 0; i < num; i++) {
    double size = 0;
    std::fill(crack_damage.begin(), crack_damage.end(), 0.0);

    for (unsigned int j = 0; j < pow(L, 2); j++) {
      if (component[n.G.edges[j][0]] == i && n.fuses[j]) {
        size++;
        crack_damage[j] = 1.0;
      }
    }

    if (size > 0) {
      sc[size - 1]++;
      t_crack_damage = data_transform(L, crack_damage, forward_plan, fftw_forward_in, fftw_forward_out);
      correlation(L, Ccc, t_crack_damage, t_crack_damage, reverse_plan, fftw_reverse_in, fftw_reverse_out);
      Nc++;
    }
  }

  // damage-damage correlations
  std::vector<fftw_complex> t_damage = data_transform(L, n.fuses, forward_plan, fftw_forward_in, fftw_forward_out);
  correlation(L, Cdd, t_damage, t_damage, 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++;
    }
  }

  sd[total_broken]++;

}