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

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

  fftw_execute(forward_plan);

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

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

template <class T>
void autocorrelation(unsigned int Mx, unsigned int My, 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 < My * (Mx / 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 < (Mx / 2 + 1) * (My / 2 + 1); i++) {
    out_data[i] += (T)(fftw_reverse_out[Mx * (i / (Mx / 2 + 1)) + i % (Mx / 2 + 1)] / (Mx * My));
  }
}

unsigned int edge_r_to_ind(graph::coordinate r, double Lx, double Ly, unsigned int Mx, unsigned int My) {
  return floor((Mx * r.x) / Lx) + Mx * floor((My * r.y) / Ly);
}

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

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

  fftw_forward_in = (double *)fftw_malloc(Mx * My * sizeof(double));
  fftw_forward_out = (fftw_complex *)fftw_malloc(Mx * My * sizeof(fftw_complex));
  fftw_reverse_in = (fftw_complex *)fftw_malloc(Mx * My * sizeof(fftw_complex));
  fftw_reverse_out = (double *)fftw_malloc(Mx * My * sizeof(double));

  forward_plan = fftw_plan_dft_r2c_2d(My, Mx, fftw_forward_in, fftw_forward_out, 0);
  reverse_plan = fftw_plan_dft_c2r_2d(My, Mx, 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, Mx, My);
  update_field_file("Css", Css, N, Lx, Ly, beta, Mx, My);
  update_field_file("Cmm", Cmm, N, Lx, Ly, beta, Mx, My);
  update_field_file("Cdd", Cdd, N, Lx, Ly, beta, Mx, My);
  update_field_file("Caa", Caa, Na, Lx, Ly, beta, Mx, My);
  update_field_file("Cll", Cll, N, Lx, Ly, beta, Mx, My);
  update_field_file("Cee", Cee, N, Lx, Ly, beta, Mx, My);
}

void ma::pre_fracture(const network&) {
  lv = std::numeric_limits<long double>::lowest();
  avalanches = {{}};
  boost::remove_edge_if(trivial, G);
}

void ma::bond_broken(const network& net, const current_info& cur, unsigned int i) {
  long double c = logl(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, Mx * My, 0.0);

    for (auto e : avalanches.back()) {
      fftw_forward_in[edge_r_to_ind(net.G.edges[e].r, Lx, Ly, Mx, My)] += 1.0;
    }

    autocorrelation(Mx, My, 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].v[0], net.G.dual_edges[i].v[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);

  // crack surface correlations
  std::fill_n(fftw_forward_in, Mx * My, 0.0);

  for (auto edge : crack) {
    fftw_forward_in[edge_r_to_ind(n.G.dual_vertices[n.G.dual_edges[edge].v[0]].r, Lx, Ly, Mx, My)] = 0.5;
    fftw_forward_in[edge_r_to_ind(n.G.dual_vertices[n.G.dual_edges[edge].v[1]].r, Lx, Ly, Mx, My)] = 0.5;
  }

  autocorrelation(Mx, My, Css, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

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

  std::vector<std::list<unsigned int>> components(num);

  for (unsigned int i = 0; i < n.G.dual_vertices.size(); i++) {
    components[component[i]].push_back(i);
  }

  // non-spanning clusters
  std::fill_n(fftw_forward_in, Mx * My, 0.0);
  for (unsigned int i = 0; i < num; i++) {
    if (i != crack_component && components[i].size() > 0) {
      for (auto it = components[i].begin(); it != components[i].end(); it++) {
        fftw_forward_in[edge_r_to_ind(n.G.dual_vertices[*it].r, Lx, Ly, Mx, My)] += 1.0;
      }

      sc[components[i].size() - 1]++;
      autocorrelation(Mx, My, Ccc, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);
      Nc++;

      for (auto it = components[i].begin(); it != components[i].end(); it++) {
        fftw_forward_in[edge_r_to_ind(n.G.dual_vertices[*it].r, Lx, Ly, Mx, My)] = 0.0;
      }
    }
  }

  unsigned int max_factor = log2(Mx < My ? Mx : My) + 1;
  std::vector<std::valarray<unsigned int>> bins(max_factor);
  for (unsigned int i = 0; i < max_factor; i++) {
    bins[i].resize(Mx * My / (unsigned int)pow(2, 2 * i));
  }

  for (auto it = components[crack_component].begin(); it != components[crack_component].end(); it++) {
    for (unsigned int i = 0; i < max_factor; i++) {
      bins[i][edge_r_to_ind(n.G.dual_vertices[*it].r, Lx, Ly, Mx / pow(2, i), My / pow(2, i))] = 1;
    }
  }

  for (unsigned int i =0 ; i < max_factor; i++) {
    sb[i] += bins[i].sum();
  }

  // spanning cluster
  std::fill_n(fftw_forward_in, Mx * My, 0.0);

  for (unsigned int i = 0; i < n.G.dual_vertices.size(); i++) {
    if (component[i] == crack_component)  {
      fftw_forward_in[edge_r_to_ind(n.G.dual_vertices[i].r, Lx, Ly, Mx, My)] += 1.0;
    }
  }

  autocorrelation(Mx, My, Cmm, 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 component[n.G.dual_edges[i].v[0]] == crack_component;
  };

  for (auto avalanche : avalanches) {
    if (avalanche.end() != std::find_if(avalanche.begin(), avalanche.end(), inCrack)) {
      for (auto edge : avalanche) {
        fftw_forward_in[edge_r_to_ind(n.G.edges[edge].r, Lx, Ly, Mx, My)] += 1.0;
      }
    }
  }

  autocorrelation(Mx, My, Cee, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  std::fill_n(fftw_forward_in, Mx * My, 0.0);

  // rewind the last avalanche
  for (auto e : avalanches.back()) {
    boost::remove_edge(n.G.dual_edges[e].v[0], n.G.dual_edges[e].v[1], G);
    n.break_edge(e, true);
    fftw_forward_in[edge_r_to_ind(n.G.edges[e].r, Lx, Ly, Mx, My)] += 1.0;
  }

  autocorrelation(Mx, My, 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;

  std::fill_n(fftw_forward_in, Mx * My, 0.0);

  for (unsigned int i = 0; i < n.G.edges.size(); i++) {
    if (n.fuses[i]) { 
      total_broken++;
      fftw_forward_in[edge_r_to_ind(n.G.edges[i].r, Lx, Ly, Mx, My)] += 1.0;
    }
  }

  autocorrelation(Mx, My, Cdd, forward_plan, fftw_forward_in, fftw_forward_out, reverse_plan, fftw_reverse_in, fftw_reverse_out);

  sd[total_broken]++;

  N++;
}