#include <mpi.h>
#include <stdio.h>
#include <algorithm>
#include <vector>
#include "kruskal.h"
#include <math.h>
#include <iostream>
#include <stdlib.h>
#include <utility>
#include <time.h>
#include <stddef.h>
using namespace std;

template<typename T>
struct tree_element {
 tree_element* parent;
 T elem;
 int operation_counter;
 int rank;
 tree_element(T e) {
   operation_counter = 0;
   elem = e;
   parent = this;
   rank = 0;
 }
 tree_element() {
   elem = 0;
   parent = this;
   rank = 0;
 }
 tree_element* Union(tree_element& e) {
   tree_element* a = find_set(this);
   tree_element* b = find_set(&e);
   e.operation_counter++;
   if(a->rank > b->rank) {
     b->parent = a;
   } else {
     a->parent = b;
     if(rank == e.rank) {
         b->rank++;
     }
   }
   return parent;
 }
 static tree_element* find_set(tree_element* e);
 void print() {
   tree_element* current = this;
   printf("%d\n", current->elem);
   while(current != current->parent) {
     current = current->parent;
     printf("%d\n", current->elem);
   }
 }
};

template<typename T>
tree_element<T>* tree_element<T>::find_set(tree_element<T>* e)  {
 e->operation_counter++;
 if (e != e->parent) {
   e->parent = find_set(e->parent);
   e->operation_counter++;
 }
 return e->parent;
}



struct Edge{
  int in;
  int out;
  double cost;
  int used;
  int range;
  int inconsistent;
};

struct Node {
  double x;
  double y;
  int id;
};

double getDistance(Node& n1, Node& n2) {
  return sqrt((n1.x - n2.x) * (n1.x - n2.x) + (n1.y - n2.y) * (n1.y - n2.y));
}

long long int operation_counter = 0;

bool operator<(const Edge &e1, const Edge &e2) {
  operation_counter++;
    return e1.cost < e2.cost;
}

void null_element_usage(vector<Edge> edges_vector) {
  std::vector<Edge>::iterator it = edges_vector.begin();
  for(;it!=edges_vector.end(); it++) {
    it->used = false;
  }
}

void null_element_usage(Edge* edges_vector, int size) {
  for(int i = 0; i < size; i++) {
    edges_vector[i].used = false;
  }
}

void unnull_element_usage(Edge* edges_vector, int size) {
  for(int i = 0; i < size; i++) {
    edges_vector[i].used = false;
  }
}

Edge* concat_arrays(Edge* arr1, Edge* arr2, int size1, int size2) {
  Edge* new_edges = new Edge[size1+size2];
  copy(arr1, arr1 + size1, new_edges);
  copy(arr2, arr2 + size2, new_edges + size1);
  return new_edges;
}

pair<Edge*,int> run_MST(Edge* edges, int N, int size) {
  vector<Edge> edges_vector;
  for(int i = 0; i < size; i++) {
    edges_vector.push_back(edges[i]);
  }
  sort(edges_vector.begin(), edges_vector.end());
  tree_element<int>* elements = new tree_element<int>[N];
  std::vector<Edge>::iterator it = edges_vector.begin();
  int operation_counter = 0;
  int len = 0;
  Edge* edges_recieved = new Edge[2*N];
  for(;it != edges_vector.end();it++) {
    operation_counter++;
    if(!it->used && tree_element<int>::find_set(&elements[it->in]) != tree_element<int>::find_set(&elements[it->out])) {
      it->used = true;
      it->range = len;
      elements[it->in].Union(elements[it->out]);
      edges_recieved[len] = *it;
      len++;
      operation_counter++;
    }
  }
  for(int kk = 0; kk < N; kk++) {
    operation_counter += elements[kk].operation_counter;
  }
  delete [] elements;
  return make_pair(edges_recieved, len);
}

int main(int argc, char **argv) {
  const int tag = 13;
  int rank, wsize;
  MPI_Init (&argc, &argv);      /* starts MPI */
  MPI_Comm_rank (MPI_COMM_WORLD, &rank);        /* get current process id */
  MPI_Comm_size (MPI_COMM_WORLD, &wsize);        /* get number of processes */
  int highest_rank = 0;
  int size = wsize - 1;
  while(size/2 > 0) {
    highest_rank++;
    size = size / 2;
  }
  size = wsize - 1;
  int number_of_processes = wsize;
  //construct Edge type
  const int nitems=6;
  int blocklengths[6] = {1,1,1,1,1,1};
  MPI_Datatype types[6] = {MPI_INT, MPI_INT, MPI_DOUBLE, MPI_INT, MPI_INT, MPI_INT};
  MPI_Datatype mpi_edge;
  MPI_Aint     offsets[6];
  offsets[0] = offsetof(struct Edge, in);
  offsets[1] = offsetof(struct Edge, out);
  offsets[2] = offsetof(struct Edge, cost);
  offsets[3] = offsetof(struct Edge, used);
  offsets[4] = offsetof(struct Edge, range);
  offsets[5] = offsetof(struct Edge, inconsistent);
  MPI_Type_create_struct(nitems, blocklengths, offsets, types, &mpi_edge);
  MPI_Type_commit(&mpi_edge);
  //construct Node type
  const int nitems_=3;
  int blocklengths_[6] = {1,1,1};
  MPI_Datatype types_[3] = {MPI_DOUBLE, MPI_DOUBLE, MPI_INT};
  MPI_Datatype mpi_node;
  MPI_Aint     offsets_[6];
  offsets_[0] = offsetof(struct Node, x);
  offsets_[1] = offsetof(struct Node, y);
  offsets_[2] = offsetof(struct Node, id);
  MPI_Type_create_struct(nitems_, blocklengths_, offsets_, types_, &mpi_node);
  MPI_Type_commit(&mpi_node);
  MPI_Status status;
  //main process
  if(rank == 0) {
    const clock_t begin_time = clock();
    int cases = 1;
    FILE* pf = fopen("./tests1000.txt","r");
    FILE* pfresult = fopen("./result.txt", "a");
    FILE* out_MST = fopen("./MST.txt", "w");
    FILE* out_clasters = fopen("./clasters.txt", "w");
    FILE* out_cluster_points = fopen("./cluster_points.txt", "w");
    int N;
    fscanf(pf,"%d", &N);
    int each = N / (number_of_processes - 1);
    tree_element<int>* elements = new tree_element<int>[N];
    Node* nodes = new Node[N];
    MPI_Bcast(&N, 1 , MPI_INT, 0, MPI_COMM_WORLD);

    for(int j = 0; j < N; j++) {
      double x,y;
      fscanf(pf,"%lf %lf", &x, &y);
      Node n;
      n.x = x;
      n.y = y;
      n.id = j;
      elements[j].elem = j;
      nodes[j] = n;
    }
    MPI_Bcast(nodes, N , mpi_node, 0, MPI_COMM_WORLD);
    vector<Edge> edges_used;
    int full_size = 0;
    int result_size = 0;
    MPI_Recv(&result_size, 1 , MPI_INT, 1, tag, MPI_COMM_WORLD, &status);
    operation_counter++;
    Edge* edges_new_recieved = new Edge[result_size];
    MPI_Recv(edges_new_recieved, result_size , mpi_edge, 1, tag, MPI_COMM_WORLD, &status);
    operation_counter++;
    for(int k = 0; k < result_size; k++) {
      if(edges_new_recieved[k].used) {
        edges_used.push_back(edges_new_recieved[k]);
        full_size++;
      }
    }
    vector<Edge> edges_vector_;
    for(int j = 0; j < full_size; j++) {
      edges_used[j].used = false;
      edges_vector_.push_back(edges_used[j]);
    }
    delete [] elements;
    elements = new tree_element<int>[N];
    sort(edges_vector_.begin(), edges_vector_.end());
    std::vector<Edge>::iterator it = edges_vector_.begin();
    int len = 0;
    for(;it != edges_vector_.end();it++) {
      if(!it->used && tree_element<int>::find_set(&elements[it->in]) != tree_element<int>::find_set(&elements[it->out])) {
        it->used = true;
        len++;
        it->range = len;
        elements[it->in].Union(elements[it->out]);
      }
    }
    for(int kk = 0; kk < N; kk++) {
      operation_counter += elements[kk].operation_counter;
    }
    it = edges_vector_.begin();
    long long int* op_counts = new long long int [number_of_processes];
    long long int all_operation_counter = operation_counter;
    long long int new_counter;
    MPI_Recv(&new_counter, 1 , MPI_LONG_LONG_INT, 1, tag, MPI_COMM_WORLD, &status);
    operation_counter++;
    all_operation_counter += new_counter;
    double clusters = 10;
    fprintf(pfresult, "time: %lf ; number_of_processes: %d ; data vol: %d ; number of operations: %lld ; flops: %lf \n", float( clock () - begin_time ) /  CLOCKS_PER_SEC, number_of_processes, N , all_operation_counter, all_operation_counter / (float( clock () - begin_time ) /  CLOCKS_PER_SEC));
    for(;it != edges_vector_.end();it++) {
      if(it->used) {
        fprintf(out_MST, "%lf %lf %lf %lf %lf\n", nodes[it->in].x, nodes[it->in].y, nodes[it->out].x, nodes[it->out].y, it->cost);
        if(it->range <= len - clusters) {
          fprintf(out_clasters, "%lf %lf %lf %lf %lf\n", nodes[it->in].x, nodes[it->in].y, nodes[it->out].x, nodes[it->out].y, it->cost);
        } else {
          it->inconsistent = true;
        }
      }
    }
    it = edges_vector_.begin();
    tree_element<int>* elements_new = new tree_element<int>[N];
    for(;it != edges_vector_.end();it++) {
      if(it->used && !it->inconsistent) {
        elements_new[it->in].Union(elements_new[it->out]);
      }
    }
    for(int k = 0; k < N; k++) {
      elements_new[k].elem = k;
    }
    for(int k = 0; k < N; k++) {
      int cluster = tree_element<int>::find_set(&elements_new[k])->elem;
      fprintf(out_cluster_points, "%lf %lf %d\n", nodes[k].x, nodes[k].y, cluster);
    }
    delete [] elements_new;
  } else {
    int N;
    MPI_Bcast(&N, 1, MPI_INT, 0, MPI_COMM_WORLD);
    tree_element<int>* elements = new tree_element<int>[N];
    Node* nodes = new Node[N];
    int each = N / (number_of_processes - 1);
    MPI_Bcast(nodes, N , mpi_node, 0, MPI_COMM_WORLD);
    int len = 0;
    int result_start = N - each * (rank - 1);
    vector<Edge> edges_vector;
    for(int i = N - result_start; i < (N - (result_start - each)); i++) {
      operation_counter++;
      for(int j = i + 1; j < N; j++) {
        double dist = getDistance(nodes[i], nodes[j]);
        operation_counter++;
        Edge e;
        e.in = nodes[i].id;
        e.out = nodes[j].id;
        e.cost = dist;
        e.used = 0;
        e.inconsistent = 0;
        edges_vector.push_back(e);
        len++;
      }
    }

    sort(edges_vector.begin(), edges_vector.end());
    std::vector<Edge>::iterator it = edges_vector.begin();
    Edge* edges_recieved = new Edge[2*N];
    int kk = len;
    len = 0;
    for(;it != edges_vector.end();it++) {
      operation_counter++;
      if(!it->used && tree_element<int>::find_set(&elements[it->in]) != tree_element<int>::find_set(&elements[it->out])) {
        it->used = true;
        it->range = len;
        elements[it->in].Union(elements[it->out]);
        edges_recieved[len] = *it;
        len++;
        operation_counter++;
      }
    }
    delete [] elements;
    for(int kk = 0; kk < N; kk++) {
      operation_counter += elements[kk].operation_counter;
    }
    operation_counter++;
    int mult = 1;
    while(highest_rank > 0) {
      mult*=2;
      highest_rank--;
      if((rank -1) % mult == 0) {
        int result_size = 0;
        MPI_Recv(&result_size, 1 , MPI_INT, rank+(mult/2), tag, MPI_COMM_WORLD, &status);
        Edge* edges_new_recieved  = new Edge[result_size];
        MPI_Recv(edges_new_recieved, result_size , mpi_edge, rank+(mult/2), tag, MPI_COMM_WORLD, &status);
        long long int new_counter;
        MPI_Recv(&new_counter, 1 , MPI_LONG_LONG_INT, rank+(mult/2), tag, MPI_COMM_WORLD, &status);
        operation_counter += new_counter;
        Edge* concated_edges = concat_arrays(edges_recieved, edges_new_recieved, len, result_size);
        delete [] edges_recieved;
        delete [] edges_new_recieved;
        null_element_usage(concated_edges,len+result_size);
        pair<Edge*,int> result = run_MST(concated_edges, N, len+result_size);
        len = result.second;
        edges_recieved = result.first;
      } else {
        MPI_Send(&len, 1 , MPI_INT, rank-(mult/2), tag, MPI_COMM_WORLD);
        MPI_Send(edges_recieved, len , mpi_edge, rank-(mult/2), tag, MPI_COMM_WORLD);
        MPI_Send(&operation_counter, 1 , MPI_LONG_LONG_INT, rank-(mult/2), tag, MPI_COMM_WORLD);
        break;
      }
    }
    if(highest_rank ==0 && rank == 1) {
      MPI_Send(&len, 1 , MPI_INT, 0, tag, MPI_COMM_WORLD);
      MPI_Send(edges_recieved, len , mpi_edge, 0, tag, MPI_COMM_WORLD);
      MPI_Send(&operation_counter, 1 , MPI_LONG_LONG_INT, 0, tag, MPI_COMM_WORLD);
    }
  }
  MPI_Finalize();
  return 0;
}