#include "measurements.hpp"

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

void update_distribution_data(std::string id, const std::vector<unsigned int>& 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();
}

void update_field_data(std::string id, double tot, const std::vector<double>& data, 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;
  uint64_t tot_old = 0;
  uint64_t tot_2_old = 0;
  std::vector<uint64_t> data_old(data.size(), 0);
  std::vector<uint64_t> data_old_2(data.size(), 0);

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

    file.close();
  }

  std::ofstream file_out(filename);

  file_out <<std::fixed<< N_old + 1 << " " << (uint64_t)(tot_old + tot) << " " << (uint64_t)(tot_2_old + pow(tot, 2))<< "\n";
  for (unsigned int i = 0; i < data.size(); i++) {
    file_out << (uint64_t)(data_old[i] + data[i]) << " ";
  }
  file_out <<std::fixed<< "\n";
  for (unsigned int i = 0; i < data.size(); i++) {
    file_out << (uint64_t)(data_old_2[i] + pow(data[i], 2)) << " ";
  }

  file_out.close();
}

std::vector<fftw_complex> data_transform(unsigned int L, const std::vector<double>& 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] = data[i];
  }

  fftw_execute(forward_plan);

  std::vector<fftw_complex> output(pow(L, 2));

  for (unsigned int i = 0; i < pow(L, 2); i++) {
    output[i][0] = fftw_forward_out[i][0];
    output[i][1] = fftw_forward_out[i][1];
  }

  return output;
}

std::vector<double> correlation(unsigned int L, 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 < pow(L, 2); 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);

  std::vector<double> output(pow(L / 2 + 1, 2));

  for (unsigned int i = 0; i < pow(L / 2 + 1, 2); i++) {
    output[i] = fftw_reverse_out[L * (i / (L / 2 + 1)) + i % (L / 2 + 1)] / pow(L, 2);
  }

  return output;
}

ma::ma(unsigned int N, unsigned int L, double beta) : L(L), G(2 * pow(L / 2, 2)), bin_counts(log2(L) + 1, 0), N(N), beta(beta) {
  ad.resize(pow(L, 2), 0);
  cd.resize(pow(L, 2), 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_data("na", ad, N, L, beta);
  update_distribution_data("nc", cd, N, L, beta);
  update_distribution_data("bc", bin_counts, N, L, beta);

}

void ma::pre_fracture(const network &) {
  lv = 0;
  as = 0;
  avalanches.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) {
    ad[as]++;
    as = 0;
    lv = cur.first / fabs(cur.second[i]) * net.thresholds[i];
    avalanches.push_back({});
  } else {
    as++;
    avalanches.back().push_back(i);
  }

  boost::add_edge(net.G.dual_edges[i][0], net.G.dual_edges[i][1], 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];
    }
  }

  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) {
        bin_counts[be]++;
      }
    }
  }

  bin_counts[log2(L)]++;

  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);
  std::vector<double> Ccc = correlation(L, t_crack_damage, t_crack_damage, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  update_field_data("Ccc", damage_tot, Ccc, L, beta);

  for (auto e = avalanches.back().rbegin(); e != avalanches.back().rend(); e++) {
    boost::remove_edge(n.G.dual_edges[*e][0], n.G.dual_edges[*e][1], G);
  }

  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) {
      cd[size - 1]++;
      t_crack_damage = data_transform(L, crack_damage, forward_plan, fftw_forward_in, fftw_forward_out);
      Ccc = correlation(L, t_crack_damage, t_crack_damage, reverse_plan, fftw_reverse_in, fftw_reverse_out);

      update_field_data("Cclcl", size, Ccc, L, beta);
    }
  }

}