#ifndef __propulsion_HH__
#define __propulsion_HH__
/********************************* TRICK HEADER *******************************
PURPOSE:
      (Describe the propulsion Module Variables and Algorithm)
LIBRARY DEPENDENCY:
      ((../src/Propulsion.cpp))
PROGRAMMERS:
      (((Chun-Hsu Lai) () () () ))
*******************************************************************************/
#include <armadillo>
#include <functional>
#include "Module.hh"
#include "aux.hh"
#include "datadeck.hh"
#include "global_constants.hh"
#include "sim_services/include/simtime.h"

class Propulsion : public Actuator {
  TRICK_INTERFACE(Propulsion);

 public:
  Propulsion(Data_exchang& input);
  Propulsion(const Propulsion& other);

  Propulsion& operator=(const Propulsion& other);

  virtual void init();
  virtual void algorithm(double int_step);
  void load_proptable(const char* filename);

  enum THRUST_TYPE {
    NO_THRUST = 0,
    INPUT_THRUST = 3,
    LTG_THRUST = 4,
    HOT_STAGE = 5
  };

  enum STAGE { STAGE_1 = 0, STAGE_2 = 1, STAGE_3 = 2, FARING_SEP = 3 };

  void Allocate_stage(unsigned int stage_num);
  void set_stage_var(double isp, double fmass_init, double vmass_init,
                     double aexit_in, double fuel_flow_rate_in, double xcg0,
                     double xcg1, double moi_roll0, double moi_roll1,
                     double moi_pitch0, double moi_pitch1, double moi_yaw0,
                     double moi_yaw1, unsigned int num_stage);
  void set_stage_1();
  void set_stage_2();
  void set_stage_3();
  void set_faring_sep();
  void engine_ignition();
  void set_no_thrust();
  void set_ltg_thrust();
  // void get_input_file_var(double XCG0, double XCG1, double moi_roll0,
  //                         double moi_roll1, double moi_pitch0,
  //                         double moi_pitch1, double moi_yaw0, double
  //                         moi_yaw1, double spi, double fuel_flow_rate);
  void set_faring_mass(double in);
  void set_HOT_STAGE();
  // XXX: get_thrust_state
  enum THRUST_TYPE get_thrust_state();
  std::function<int()> grab_beco_flag;
  /* Input File */
  // void set_vmass0(double);
  // void set_fmass0(double);

  // void set_aexit(double);
  void set_payload_mass(double);

 private:
  /* Rocket Engine state variables */
  struct ENG {
    double ENG_mass;
    double ENG_mass0;
    double ENG_mass1;
    double ENG_isp;

    MATRIX(ENG_I, 3, 3);
    MATRIX(ENG_I_0, 3, 3);
    MATRIX(ENG_I_1, 3, 3);

    VECTOR(ENG_XCG, 3);
    VECTOR(ENG_XCG_0, 3);
    VECTOR(ENG_XCG_1, 3);
  };

  /* Stage variable */
  struct STAGE_VAR {
    double isp;            /* *o (--) Rocket engine's specific impluse */
    double fmass0;         /* ** (kg) Init fuel mass */
    double StageMass0;     /* ** (kg) Stage init mass */
    arma::mat33 IBBB0;     /* *o (--) Initial vehicle's MOI */
    arma::mat33 IBBB1;     /* *o (kg*m2) Final vehicle's MOI */
    arma::vec3 XCG_0;      /* *o (m) Initial vehicle's cg */
    arma::vec3 XCG_1;      /* *o (m) Final vehicle's cg */
    double fmasse;         /* *o (kg) Expended fuel mass */
    double fmassd;         /* ** (kg/s) Fuel mass expended derivative */
    double fuel_flow_rate; /* *o (kg/s) Fuel flow rate of engine */
    double aexit;          /* *o (m2) Engine nozzle area */
    double timer;          /* ** (s) Stage timer */
  };

  Datadeck proptable;
  /* Internal Initializers */
  void default_data();

  /* Internal Propagator / Calculators */
  void propagate_delta_v(double int_step, double spi, double fuel_flow_rate,
                         double vmass);
  void fuel_expend_integrator(double int_step, enum STAGE flag);

  /* Internal Calculators */
  // double calculate_thrust(double press);
  double calculate_fmassr(enum STAGE stage);

  arma::vec3 calculate_XCG(enum STAGE stage);
  arma::mat33 calculate_IBBB(enum STAGE stage);

  /* Variable declaration */
  std::vector<struct ENG*> Eng_list;
  std::vector<struct STAGE_VAR*> Stage_var_list;
  VECTOR(XCG, 3);      /* *o (m) Current vehicle's cg */
  MATRIX(IBBB, 3, 3);  /* *o (kg*m2) Current cheicle's MOI */
  double vmass;        /* *o (kg) Current vehicle mass */
  double fmassr;       /* *o (kg) Remaining fuel mass */
  double mass_ratio;   /* *o (--) fmasse / fmass0 */
  double payload_mass; /* *o (kg) Payload mass */
  double faring_mass;  /* *o (kg) Vehicle faring mass */
  double thrust;       /* *o (N) Thrust generated by vehicle's engine */
  double delta_v;      /* *o (m/s)  Vehicle delta-V generated by engine */

  /* State */
  enum THRUST_TYPE
      thrust_state; /* *o (--)     Propulsion mode, See THRUST TYPE*/
  enum STAGE stage;
};

#endif  // __propulsion_HH__


#include <random>
#include <functional>
#include <vector>
#include <chrono>
#include <boost/format.hpp>
#include <algorithm>
#include "mpi.h"
#include "array3d.h"

using namespace std;

uniform_real_distribution<double> generator { -10, 10 };
default_random_engine re;

int mpi_rank;
int mpi_size;

const double GB = 1024 * 1024 * 1024;

double dt;
double dtdx;
double dtdy;
double dtdz;

double gen() {
    return generator(re);
}

struct vector3d {
    double x = 0, y = 0, z = 0;
};

struct vector6d {
    double ex = 0, ey = 0, ez = 0, bx = 0, by = 0, bz = 0;
};

template <typename T>
void randomize(T & a) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++){
                a(i, j, k) = gen();
            }
        }
    }
}

template <typename T>
void randomize_vector(T & a) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++){
                a(i, j, k).x = gen();
                a(i, j, k).y = gen();
                a(i, j, k).z = gen();
            }
        }
    }
}

template <typename T>
void randomize_vector6d(T & a) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++){
                a(i, j, k).ex = gen();
                a(i, j, k).ey = gen();
                a(i, j, k).ez = gen();
                a(i, j, k).bx = gen();
                a(i, j, k).by = gen();
                a(i, j, k).bz = gen();
            }
        }
    }
}

template <typename T>
void advance_b(const T & ex, const T & ey, const T & ez, T & bx, T & by, T & bz) {
    const int n = ex.get_n();

    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                bx(i,j,k) += 0.5*(-dtdy*(ez(i,j+1,k)-ez(i,j,k)) + dtdz*(ey(i,j,k+1)-ey(i,j,k)));
                by(i,j,k) += 0.5*( dtdx*(ez(i+1,j,k)-ez(i,j,k)) - dtdz*(ex(i,j,k+1)-ex(i,j,k)));
                bz(i,j,k) += 0.5*( dtdy*(ex(i,j+1,k)-ex(i,j,k)) - dtdx*(ey(i+1,j,k)-ey(i,j,k)));
            }
        }
    }
    {int i=0;
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                by(i,j,k) += 0.5*( dtdx*(ez(i+1,j,k)-ez(i,j,k)) - dtdz*(ex(i,j,k+1)-ex(i,j,k)));
                bz(i,j,k) += 0.5*( dtdy*(ex(i,j+1,k)-ex(i,j,k)) - dtdx*(ey(i+1,j,k)-ey(i,j,k)));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        {int j=0;
            for(int k=1;k<n-1;k++) {
                bx(i,j,k) += 0.5*(-dtdy*(ez(i,j+1,k)-ez(i,j,k)) + dtdz*(ey(i,j,k+1)-ey(i,j,k)));
                bz(i,j,k) += 0.5*( dtdy*(ex(i,j+1,k)-ex(i,j,k)) - dtdx*(ey(i+1,j,k)-ey(i,j,k)));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            {int k=0;
                bx(i,j,k) += 0.5*(-dtdy*(ez(i,j+1,k)-ez(i,j,k)) + dtdz*(ey(i,j,k+1)-ey(i,j,k)));
                by(i,j,k) += 0.5*( dtdx*(ez(i+1,j,k)-ez(i,j,k)) - dtdz*(ex(i,j,k+1)-ex(i,j,k)));
            }
        }
    }
}

void advance_b_z(const morton_array<double> & ex, const morton_array<double> & ey, const morton_array<double> & ez,
        morton_array<double> & bx, morton_array<double> & by, morton_array<double> & bz) {
    const unsigned int n = ex.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        if (ex.is_xmax(curr) || ex.is_ymax(curr) || ex.is_zmax(curr)) {
            continue;
        }

        auto x_next = ex.get_x_next(curr);
        auto y_next = ex.get_y_next(curr);
        auto z_next = ex.get_z_next(curr);

        if (ex.is_xmin(curr)) {
            if (ex.is_ymin(curr) || ex.is_zmin(curr)) {
                continue;
            } else {
                by[curr] += 0.5*( dtdx*(ez[x_next]-ez[curr]) - dtdz*(ex[z_next]-ex[curr]));
                bz[curr] += 0.5*( dtdy*(ex[y_next]-ex[curr]) - dtdx*(ey[x_next]-ey[curr]));
            }
        } else if (ex.is_ymin(curr)) {
            if (ex.is_zmin(curr)) {
                continue;
            } else {
                bx[curr] += 0.5*(-dtdy*(ez[y_next]-ez[curr]) + dtdz*(ey[z_next]-ey[curr]));
                bz[curr] += 0.5*( dtdy*(ex[y_next]-ex[curr]) - dtdx*(ey[x_next]-ey[curr]));
            }
        } else if (ex.is_zmin(curr)) {
            bx[curr] += 0.5*(-dtdy*(ez[y_next]-ez[curr]) + dtdz*(ey[z_next]-ey[curr]));
            by[curr] += 0.5*( dtdx*(ez[x_next]-ez[curr]) - dtdz*(ex[z_next]-ex[curr]));
        } else {
            bx[curr] += 0.5*(-dtdy*(ez[y_next]-ez[curr]) + dtdz*(ey[z_next]-ey[curr]));
            by[curr] += 0.5*( dtdx*(ez[x_next]-ez[curr]) - dtdz*(ex[z_next]-ex[curr]));
            bz[curr] += 0.5*( dtdy*(ex[y_next]-ex[curr]) - dtdx*(ey[x_next]-ey[curr]));
        }
    }
}

template <typename T>
void advance_b_vector(const T & e, T & b) {
    const int n = e.get_n();

    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
                b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
                b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
            }
        }
    }
    {int i=0;
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
                b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        {int j=0;
            for(int k=1;k<n-1;k++) {
                b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
                b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            {int k=0;
                b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
                b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
            }
        }
    }
}

void advance_b_vector_z(const morton_array<vector3d> & e, morton_array<vector3d> & b) {
    const unsigned int n = e.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        if (e.is_xmax(curr) || e.is_ymax(curr) || e.is_zmax(curr)) {
            continue;
        }

        auto x_next = e.get_x_next(curr);
        auto y_next = e.get_y_next(curr);
        auto z_next = e.get_z_next(curr);

        if (e.is_xmin(curr)) {
            if (e.is_ymin(curr) || e.is_zmin(curr)) {
                continue;
            } else {
                b[curr].y += 0.5*( dtdx*(e[x_next].z-e[curr].z) - dtdz*(e[z_next].x-e[curr].x));
                b[curr].z += 0.5*( dtdy*(e[y_next].x-e[curr].x) - dtdx*(e[x_next].y-e[curr].y));
            }
        } else if (e.is_ymin(curr)) {
            if (e.is_zmin(curr)) {
                continue;
            } else {
                b[curr].x += 0.5*(-dtdy*(e[y_next].z-e[curr].z) + dtdz*(e[z_next].y-e[curr].y));
                b[curr].z += 0.5*( dtdy*(e[y_next].x-e[curr].x) - dtdx*(e[x_next].y-e[curr].y));
            }
        } else if (e.is_zmin(curr)) {
            b[curr].x += 0.5*(-dtdy*(e[y_next].z-e[curr].z) + dtdz*(e[z_next].y-e[curr].y));
            b[curr].y += 0.5*( dtdx*(e[x_next].z-e[curr].z) - dtdz*(e[z_next].x-e[curr].x));
        } else {
            b[curr].x += 0.5*(-dtdy*(e[y_next].z-e[curr].z) + dtdz*(e[z_next].y-e[curr].y));
            b[curr].y += 0.5*( dtdx*(e[x_next].z-e[curr].z) - dtdz*(e[z_next].x-e[curr].x));
            b[curr].z += 0.5*( dtdy*(e[y_next].x-e[curr].x) - dtdx*(e[x_next].y-e[curr].y));
        }
    }
}

template <typename T>
void advance_b_vector_periodic(const T & e, T & b) {
    const int n = e.get_n();

    for(int i=0;i<n-1;i++) {
        for(int j=0;j<n-1;j++) {
            for(int k=0;k<n-1;k++) {
                b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
                b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
                b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
            }
            const int k = n-1;
            b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,0).y-e(i,j,k).y));
            b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,0).x-e(i,j,k).x));
            b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
        }
        const int j = n-1;
        for(int k=0;k<n-1;k++) {
            b(i,j,k).x += 0.5*(-dtdy*(e(i,0,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
            b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
            b(i,j,k).z += 0.5*( dtdy*(e(i,0,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
        }
        const int k = n-1;
        b(i,j,k).x += 0.5*(-dtdy*(e(i,0,k).z-e(i,j,k).z) + dtdz*(e(i,j,0).y-e(i,j,k).y));
        b(i,j,k).y += 0.5*( dtdx*(e(i+1,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,0).x-e(i,j,k).x));
        b(i,j,k).z += 0.5*( dtdy*(e(i,0,k).x-e(i,j,k).x) - dtdx*(e(i+1,j,k).y-e(i,j,k).y));
    }
    const int i = n-1;
    for(int j=0;j<n-1;j++) {
        for(int k=0;k<n-1;k++) {
            b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
            b(i,j,k).y += 0.5*( dtdx*(e(0,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
            b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(0,j,k).y-e(i,j,k).y));
        }
        const int k = n-1;
        b(i,j,k).x += 0.5*(-dtdy*(e(i,j+1,k).z-e(i,j,k).z) + dtdz*(e(i,j,0).y-e(i,j,k).y));
        b(i,j,k).y += 0.5*( dtdx*(e(0,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,0).x-e(i,j,k).x));
        b(i,j,k).z += 0.5*( dtdy*(e(i,j+1,k).x-e(i,j,k).x) - dtdx*(e(0,j,k).y-e(i,j,k).y));
    }
    const int j = n-1;
    for(int k=0;k<n-1;k++) {
        b(i,j,k).x += 0.5*(-dtdy*(e(i,0,k).z-e(i,j,k).z) + dtdz*(e(i,j,k+1).y-e(i,j,k).y));
        b(i,j,k).y += 0.5*( dtdx*(e(0,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,k+1).x-e(i,j,k).x));
        b(i,j,k).z += 0.5*( dtdy*(e(i,0,k).x-e(i,j,k).x) - dtdx*(e(0,j,k).y-e(i,j,k).y));
    }
    const int k = n-1;
    b(i,j,k).x += 0.5*(-dtdy*(e(i,0,k).z-e(i,j,k).z) + dtdz*(e(i,j,0).y-e(i,j,k).y));
    b(i,j,k).y += 0.5*( dtdx*(e(0,j,k).z-e(i,j,k).z) - dtdz*(e(i,j,0).x-e(i,j,k).x));
    b(i,j,k).z += 0.5*( dtdy*(e(i,0,k).x-e(i,j,k).x) - dtdx*(e(0,j,k).y-e(i,j,k).y));
}

void advance_b_vector_z_periodic(const morton_array<vector3d> & e, morton_array<vector3d> & b) {
    const unsigned int n = e.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        auto x_next = e.get_x_next(curr);
        auto y_next = e.get_y_next(curr);
        auto z_next = e.get_z_next(curr);

        b[curr].x += 0.5*(-dtdy*(e[y_next].z-e[curr].z) + dtdz*(e[z_next].y-e[curr].y));
        b[curr].y += 0.5*( dtdx*(e[x_next].z-e[curr].z) - dtdz*(e[z_next].x-e[curr].x));
        b[curr].z += 0.5*( dtdy*(e[y_next].x-e[curr].x) - dtdx*(e[x_next].y-e[curr].y));
    }
}

template <typename T>
void advance_b_vector6d(const T & v) {
    const int n = v.get_n();

    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                v(i,j,k).bx += 0.5*(-dtdy*(v(i,j+1,k).ez-v(i,j,k).ez) + dtdz*(v(i,j,k+1).ey-v(i,j,k).ey));
                v(i,j,k).by += 0.5*( dtdx*(v(i+1,j,k).ez-v(i,j,k).ez) - dtdz*(v(i,j,k+1).ex-v(i,j,k).ex));
                v(i,j,k).bz += 0.5*( dtdy*(v(i,j+1,k).ex-v(i,j,k).ex) - dtdx*(v(i+1,j,k).ey-v(i,j,k).ey));
            }
        }
    }
    {int i=0;
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                v(i,j,k).by += 0.5*( dtdx*(v(i+1,j,k).ez-v(i,j,k).ez) - dtdz*(v(i,j,k+1).ex-v(i,j,k).ex));
                v(i,j,k).bz += 0.5*( dtdy*(v(i,j+1,k).ex-v(i,j,k).ex) - dtdx*(v(i+1,j,k).ey-v(i,j,k).ey));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        {int j=0;
            for(int k=1;k<n-1;k++) {
                v(i,j,k).bx += 0.5*(-dtdy*(v(i,j+1,k).ez-v(i,j,k).ez) + dtdz*(v(i,j,k+1).ey-v(i,j,k).ey));
                v(i,j,k).bz += 0.5*( dtdy*(v(i,j+1,k).ex-v(i,j,k).ex) - dtdx*(v(i+1,j,k).ey-v(i,j,k).ey));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            {int k=0;
                v(i,j,k).bx += 0.5*(-dtdy*(v(i,j+1,k).ez-v(i,j,k).ez) + dtdz*(v(i,j,k+1).ey-v(i,j,k).ey));
                v(i,j,k).by += 0.5*( dtdx*(v(i+1,j,k).ez-v(i,j,k).ez) - dtdz*(v(i,j,k+1).ex-v(i,j,k).ex));
            }
        }
    }
}

void advance_b_vector6d_z(const morton_array<vector6d> & v) {
    const unsigned int n = v.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        if (v.is_xmax(curr) || v.is_ymax(curr) || v.is_zmax(curr)) {
            continue;
        }

        auto x_next = v.get_x_next(curr);
        auto y_next = v.get_y_next(curr);
        auto z_next = v.get_z_next(curr);

        if (v.is_xmin(curr)) {
            if (v.is_ymin(curr) || v.is_zmin(curr)) {
                continue;
            } else {
                v[curr].by += 0.5*( dtdx*(v[x_next].ez-v[curr].ez) - dtdz*(v[z_next].ex-v[curr].ex));
                v[curr].bz += 0.5*( dtdy*(v[y_next].ex-v[curr].ex) - dtdx*(v[x_next].ey-v[curr].ey));
            }
        } else if (v.is_ymin(curr)) {
            if (v.is_zmin(curr)) {
                continue;
            } else {
                v[curr].bx += 0.5*(-dtdy*(v[y_next].ez-v[curr].ez) + dtdz*(v[z_next].ey-v[curr].ey));
                v[curr].bz += 0.5*( dtdy*(v[y_next].ex-v[curr].ex) - dtdx*(v[x_next].ey-v[curr].ey));
            }
        } else if (v.is_zmin(curr)) {
            v[curr].bx += 0.5*(-dtdy*(v[y_next].ez-v[curr].ez) + dtdz*(v[z_next].ey-v[curr].ey));
            v[curr].by += 0.5*( dtdx*(v[x_next].ez-v[curr].ez) - dtdz*(v[z_next].ex-v[curr].ex));
        } else {
            v[curr].bx += 0.5*(-dtdy*(v[y_next].ez-v[curr].ez) + dtdz*(v[z_next].ey-v[curr].ey));
            v[curr].by += 0.5*( dtdx*(v[x_next].ez-v[curr].ez) - dtdz*(v[z_next].ex-v[curr].ex));
            v[curr].bz += 0.5*( dtdy*(v[y_next].ex-v[curr].ex) - dtdx*(v[x_next].ey-v[curr].ey));
        }
    }
}

template <typename T>
void advance_e(T & ex, T & ey, T & ez, const T & bx, const T & by, const T & bz) {
    const int n = ex.get_n();

    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                ex(i,j,k) +=  dtdy*(bz(i,j,k)-bz(i,j-1,k)) - dtdz*(by(i,j,k)-by(i,j,k-1));
                ey(i,j,k) += -dtdx*(bz(i,j,k)-bz(i-1,j,k)) + dtdz*(bx(i,j,k)-bx(i,j,k-1));
                ez(i,j,k) +=  dtdx*(by(i,j,k)-by(i-1,j,k)) - dtdy*(bx(i,j,k)-bx(i,j-1,k));
            }
        }
    }
    {int i=0;
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                ex(i,j,k) +=  dtdy*(bz(i,j,k)-bz(i,j-1,k)) - dtdz*(by(i,j,k)-by(i,j,k-1));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        {int j=0;
            for(int k=1;k<n-1;k++) {
                ey(i,j,k) += -dtdx*(bz(i,j,k)-bz(i-1,j,k)) + dtdz*(bx(i,j,k)-bx(i,j,k-1));
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            {int k=0;
                ez(i,j,k) +=  dtdx*(by(i,j,k)-by(i-1,j,k)) - dtdy*(bx(i,j,k)-bx(i,j-1,k));
            }
        }
    }
}

void advance_e_z(morton_array<double> & ex, morton_array<double> & ey, morton_array<double> & ez,
        const morton_array<double> & bx, const morton_array<double> & by, const morton_array<double> & bz) {
    const unsigned int n = ex.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        if (ex.is_xmax(curr) || ex.is_ymax(curr) || ex.is_zmax(curr)) {
            continue;
        }

        auto x_prev = ex.get_x_prev(curr);
        auto y_prev = ex.get_y_prev(curr);
        auto z_prev = ex.get_z_prev(curr);

        if (ex.is_xmin(curr)) {
            if (ex.is_ymin(curr) || ex.is_zmin(curr)) {
                continue;
            } else {
                ex[curr] +=  dtdy*(bz[curr]-bz[y_prev]) - dtdz*(by[curr]-by[z_prev]);
            }
        } else if (ex.is_ymin(curr)) {
            if (ex.is_zmin(curr)) {
                continue;
            } else {
                ey[curr] += -dtdx*(bz[curr]-bz[x_prev]) + dtdz*(bx[curr]-bx[z_prev]);
            }
        } else if (ex.is_zmin(curr)) {
            ez[curr] +=  dtdx*(by[curr]-by[x_prev]) - dtdy*(bx[curr]-bx[y_prev]);
        } else {
            ex[curr] +=  dtdy*(bz[curr]-bz[y_prev]) - dtdz*(by[curr]-by[z_prev]);
            ey[curr] += -dtdx*(bz[curr]-bz[x_prev]) + dtdz*(bx[curr]-bx[z_prev]);
            ez[curr] +=  dtdx*(by[curr]-by[x_prev]) - dtdy*(bx[curr]-bx[y_prev]);
        }
    }
}

template <typename T>
void advance_e_vector(T & e, const T & b) {
    const int n = e.get_n();

    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,j-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,k-1).y);
                e(i,j,k).y += -dtdx*(b(i,j,k).z-b(i-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,k-1).x);
                e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(i-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,j-1,k).x);
            }
        }
    }
    {int i=0;
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,j-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,k-1).y);
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        {int j=0;
            for(int k=1;k<n-1;k++) {
                e(i,j,k).y += -dtdx*(b(i,j,k).z-b(i-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,k-1).x);
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            {int k=0;
                e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(i-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,j-1,k).x);
            }
        }
    }
}

void advance_e_vector_z(morton_array<vector3d> & e, const morton_array<vector3d> & b) {
    const unsigned int n = e.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        if (e.is_xmax(curr) || e.is_ymax(curr) || e.is_zmax(curr)) {
            continue;
        }

        auto x_prev = e.get_x_prev(curr);
        auto y_prev = e.get_y_prev(curr);
        auto z_prev = e.get_z_prev(curr);

        if (e.is_xmin(curr)) {
            if (e.is_ymin(curr) || e.is_zmin(curr)) {
                continue;
            } else {
                e[curr].x +=  dtdy*(b[curr].z-b[y_prev].z) - dtdz*(b[curr].y-b[z_prev].y);
            }
        } else if (e.is_ymin(curr)) {
            if (e.is_zmin(curr)) {
                continue;
            } else {
                e[curr].y += -dtdx*(b[curr].z-b[x_prev].z) + dtdz*(b[curr].x-b[z_prev].x);
            }
        } else if (e.is_zmin(curr)) {
            e[curr].z +=  dtdx*(b[curr].y-b[x_prev].y) - dtdy*(b[curr].x-b[y_prev].x);
        } else {
            e[curr].x +=  dtdy*(b[curr].z-b[y_prev].z) - dtdz*(b[curr].y-b[z_prev].y);
            e[curr].y += -dtdx*(b[curr].z-b[x_prev].z) + dtdz*(b[curr].x-b[z_prev].x);
            e[curr].z +=  dtdx*(b[curr].y-b[x_prev].y) - dtdy*(b[curr].x-b[y_prev].x);
        }
    }
}

template <typename T>
void advance_e_vector_periodic(T & e, const T & b) {
    const int n = e.get_n();

    for(int i=1;i<n;i++) {
        for(int j=1;j<n;j++) {
            for(int k=1;k<n;k++) {
                e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,j-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,k-1).y);
                e(i,j,k).y += -dtdx*(b(i,j,k).z-b(i-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,k-1).x);
                e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(i-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,j-1,k).x);
            }
            const int k = 0;
            e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,j-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,n-1).y);
            e(i,j,k).y += -dtdx*(b(i,j,k).z-b(i-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,n-1).x);
            e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(i-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,j-1,k).x);
        }
        const int j = 0;
        for(int k=1;k<n;k++) {
            e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,n-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,k-1).y);
            e(i,j,k).y += -dtdx*(b(i,j,k).z-b(i-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,k-1).x);
            e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(i-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,n-1,k).x);
        }
        const int k = 0;
        e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,n-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,n-1).y);
        e(i,j,k).y += -dtdx*(b(i,j,k).z-b(i-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,n-1).x);
        e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(i-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,n-1,k).x);
    }
    const int i = 0;
    for(int j=1;j<n;j++) {
        for(int k=1;k<n;k++) {
            e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,j-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,k-1).y);
            e(i,j,k).y += -dtdx*(b(i,j,k).z-b(n-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,k-1).x);
            e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(n-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,j-1,k).x);
        }
        const int k = 0;
        e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,j-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,n-1).y);
        e(i,j,k).y += -dtdx*(b(i,j,k).z-b(n-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,n-1).x);
        e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(n-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,j-1,k).x);
    }
    const int j = 0;
    for(int k=1;k<n;k++) {
        e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,n-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,k-1).y);
        e(i,j,k).y += -dtdx*(b(i,j,k).z-b(n-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,k-1).x);
        e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(n-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,n-1,k).x);
    }
    const int k = 0;
    e(i,j,k).x +=  dtdy*(b(i,j,k).z-b(i,n-1,k).z) - dtdz*(b(i,j,k).y-b(i,j,n-1).y);
    e(i,j,k).y += -dtdx*(b(i,j,k).z-b(n-1,j,k).z) + dtdz*(b(i,j,k).x-b(i,j,n-1).x);
    e(i,j,k).z +=  dtdx*(b(i,j,k).y-b(n-1,j,k).y) - dtdy*(b(i,j,k).x-b(i,n-1,k).x);
}

void advance_e_vector_z_periodic(morton_array<vector3d> & e, const morton_array<vector3d> & b) {
    const unsigned int n = e.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        auto x_prev = e.get_x_prev(curr);
        auto y_prev = e.get_y_prev(curr);
        auto z_prev = e.get_z_prev(curr);

        e[curr].x +=  dtdy*(b[curr].z-b[y_prev].z) - dtdz*(b[curr].y-b[z_prev].y);
        e[curr].y += -dtdx*(b[curr].z-b[x_prev].z) + dtdz*(b[curr].x-b[z_prev].x);
        e[curr].z +=  dtdx*(b[curr].y-b[x_prev].y) - dtdy*(b[curr].x-b[y_prev].x);
    }
}

template <typename T>
void advance_e_vector6d(T & v) {
    const int n = v.get_n();

    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                v(i,j,k).ex +=  dtdy*(v(i,j,k).bz-v(i,j-1,k).bz) - dtdz*(v(i,j,k).by-v(i,j,k-1).by);
                v(i,j,k).ey += -dtdx*(v(i,j,k).bz-v(i-1,j,k).bz) + dtdz*(v(i,j,k).bx-v(i,j,k-1).bx);
                v(i,j,k).ez +=  dtdx*(v(i,j,k).by-v(i-1,j,k).by) - dtdy*(v(i,j,k).bx-v(i,j-1,k).bx);
            }
        }
    }
    {int i=0;
        for(int j=1;j<n-1;j++) {
            for(int k=1;k<n-1;k++) {
                v(i,j,k).ex +=  dtdy*(v(i,j,k).bz-v(i,j-1,k).bz) - dtdz*(v(i,j,k).by-v(i,j,k-1).by);
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        {int j=0;
            for(int k=1;k<n-1;k++) {
                v(i,j,k).ey += -dtdx*(v(i,j,k).bz-v(i-1,j,k).bz) + dtdz*(v(i,j,k).bx-v(i,j,k-1).bx);
            }
        }
    }
    for(int i=1;i<n-1;i++) {
        for(int j=1;j<n-1;j++) {
            int k=0;
            v(i,j,k).ez +=  dtdx*(v(i,j,k).by-v(i-1,j,k).by) - dtdy*(v(i,j,k).bx-v(i,j-1,k).bx);
        }
    }
}

void advance_e_vector6d_z(morton_array<vector6d> & v) {
    const unsigned int n = v.get_size();

    for (unsigned int curr = 0; curr < n; curr++) {
        if (v.is_xmax(curr) || v.is_ymax(curr) || v.is_zmax(curr)) {
            continue;
        }

        auto x_prev = v.get_x_prev(curr);
        auto y_prev = v.get_y_prev(curr);
        auto z_prev = v.get_z_prev(curr);

        if (v.is_xmin(curr)) {
            if (v.is_ymin(curr) || v.is_zmin(curr)) {
                continue;
            } else {
                v[curr].ex +=  dtdy*(v[curr].bz-v[y_prev].bz) - dtdz*(v[curr].by-v[z_prev].by);
            }
        } else if (v.is_ymin(curr)) {
            if (v.is_zmin(curr)) {
                continue;
            } else {
                v[curr].ey += -dtdx*(v[curr].bz-v[x_prev].bz) + dtdz*(v[curr].bx-v[z_prev].bx);
            }
        } else if (v.is_zmin(curr)) {
            v[curr].ez +=  dtdx*(v[curr].by-v[x_prev].by) - dtdy*(v[curr].bx-v[y_prev].bx);
        } else {
            v[curr].ex +=  dtdy*(v[curr].bz-v[y_prev].bz) - dtdz*(v[curr].by-v[z_prev].by);
            v[curr].ey += -dtdx*(v[curr].bz-v[x_prev].bz) + dtdz*(v[curr].bx-v[z_prev].bx);
            v[curr].ez +=  dtdx*(v[curr].by-v[x_prev].by) - dtdy*(v[curr].bx-v[y_prev].bx);
        }
    }
}

template <typename T>
void advance(T & ex, T & ey, T & ez, T & bx, T & by, T & bz) {
    advance_b(ex, ey, ez, bx, by, bz);
    advance_e(ex, ey, ez, bx, by, bz);
    advance_b(ex, ey, ez, bx, by, bz);
}

void advance_z(morton_array<double> & ex, morton_array<double> & ey, morton_array<double> & ez,
        morton_array<double> & bx, morton_array<double> & by, morton_array<double> & bz) {
    advance_b_z(ex, ey, ez, bx, by, bz);
    advance_e_z(ex, ey, ez, bx, by, bz);
    advance_b_z(ex, ey, ez, bx, by, bz);
}

template <typename T>
void advance_vector(T & e, T & b) {
    advance_b_vector(e, b);
    advance_e_vector(e, b);
    advance_b_vector(e, b);
}

void advance_vector_z(morton_array<vector3d> & e, morton_array<vector3d> & b) {
    advance_b_vector_z(e, b);
    advance_e_vector_z(e, b);
    advance_b_vector_z(e, b);
}

template <typename T>
void advance_vector_periodic(T & e, T & b) {
    advance_b_vector_periodic(e, b);
    advance_e_vector_periodic(e, b);
    advance_b_vector_periodic(e, b);
}

void advance_vector_z_periodic(morton_array<vector3d> & e, morton_array<vector3d> & b) {
    advance_b_vector_z_periodic(e, b);
    advance_e_vector_z_periodic(e, b);
    advance_b_vector_z_periodic(e, b);
}

template <typename T>
void advance_vector6d(T & v) {
    advance_b_vector6d(v);
    advance_e_vector6d(v);
    advance_b_vector6d(v);
}

void advance_vector6d_z(morton_array<vector6d> & v) {
    advance_b_vector6d_z(v);
    advance_e_vector6d_z(v);
    advance_b_vector6d_z(v);
}

template <typename T>
void run_test(std::function<void(T&, T&, T&, T&, T&, T&)> func, const std::string & testname, int size, int iterations=10) {
    T ex(size);
    T ey(size);
    T ez(size);
    T bx(size);
    T by(size);
    T bz(size);
    vector<double> times(iterations);
    for (int i = 0; i < iterations; i++) {
        randomize(ex);
        randomize(ey);
        randomize(ez);
        randomize(bx);
        randomize(by);
        randomize(bz);
        MPI_Barrier(MPI_COMM_WORLD);
        chrono::steady_clock::time_point begin = chrono::steady_clock::now();
        func(ex, ey, ez, bx, by, bz);
        MPI_Barrier(MPI_COMM_WORLD);
        chrono::steady_clock::time_point end = chrono::steady_clock::now();
        times[i] = chrono::duration_cast<chrono::nanoseconds>(end - begin).count() / 1e6;
    }
    if (mpi_rank == 0) {
        sort(times.begin(), times.end());
        times.pop_back();
        double avg = accumulate(times.begin(), times.end(), 0.0) / times.size();
        double avg_sq = inner_product(times.begin(), times.end(), times.begin(), 0.0) / times.size();
        double standard_deviation = sqrt(avg_sq - avg * avg);
        double throughput = (6 * static_cast<double>(mpi_size) * size * size * size * sizeof(double)) / GB * 1e3 / avg; // Gb / s
        cout << boost::format("%15s ") % testname << boost::format("%5d: ") % size << boost::format("%9.3f ms") % avg
                << boost::format(" (+-%7.3f ms)") % standard_deviation << boost::format(" %5.2f Gb/s") % throughput
                << endl;
    }
}

template <typename T>
void run_test_vector(std::function<void(T&, T&)> func, const std::string & testname, int size, int iterations=10) {
    T a(size);
    T b(size);
    vector<double> times(iterations);
    for (int i = 0; i < iterations; i++) {
        randomize_vector(a);
        randomize_vector(b);
        MPI_Barrier(MPI_COMM_WORLD);
        chrono::steady_clock::time_point begin = chrono::steady_clock::now();
        func(a, b);
        MPI_Barrier(MPI_COMM_WORLD);
        chrono::steady_clock::time_point end = chrono::steady_clock::now();
        times[i] = chrono::duration_cast<chrono::nanoseconds>(end - begin).count() / 1e6;
    }
    if (mpi_rank == 0) {
        sort(times.begin(), times.end());
        times.pop_back();
        double avg = accumulate(times.begin(), times.end(), 0.0) / times.size();
        double avg_sq = inner_product(times.begin(), times.end(), times.begin(), 0.0) / times.size();
        double standard_deviation = sqrt(avg_sq - avg * avg);
        double throughput = (2 * static_cast<double>(mpi_size) * size * size * size * sizeof(vector3d)) / GB * 1e3 / avg; // Gb / s
        cout << boost::format("%15s ") % testname << boost::format("%9.3f ms") % avg
                << boost::format(" (+-%7.3f ms)") % standard_deviation << boost::format(" %5.2f Gb/s") % throughput
                << endl;
    }
}

template <typename T>
void run_test_vector6d(std::function<void(T&)> func, const std::string & testname, int size, int iterations=10) {
    T a(size);
    vector<double> times(iterations);
    for (int i = 0; i < iterations; i++) {
        randomize_vector6d(a);
        MPI_Barrier(MPI_COMM_WORLD);
        chrono::steady_clock::time_point begin = chrono::steady_clock::now();
        func(a);
        MPI_Barrier(MPI_COMM_WORLD);
        chrono::steady_clock::time_point end = chrono::steady_clock::now();
        times[i] = chrono::duration_cast<chrono::nanoseconds>(end - begin).count() / 1e6;
    }
    if (mpi_rank == 0) {
        sort(times.begin(), times.end());
        times.pop_back();
        double avg = accumulate(times.begin(), times.end(), 0.0) / times.size();
        double avg_sq = inner_product(times.begin(), times.end(), times.begin(), 0.0) / times.size();
        double standard_deviation = sqrt(avg_sq - avg * avg);
        double throughput = (static_cast<double>(mpi_size) * size * size * size * sizeof(vector6d)) / GB * 1e3 / avg; // Gb / s
        cout << boost::format("%15s ") % testname << boost::format("%5d: ") % size << boost::format("%9.3f ms") % avg
                << boost::format(" (+-%7.3f ms)") % standard_deviation << boost::format(" %5.2f Gb/s") % throughput
                << endl;
    }
}

template <typename T, typename V>
bool is_equal(T & a, V & b) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++) {
                if (std::fabs(b(i, j, k) - a(i, j, k)) > 1e-6) return false;
            }
        }
    }
    return true;
}

template <typename T, typename V>
bool is_equal_vector(T & a, V & b) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++) {
                if (std::fabs(b(i, j, k).x - a(i, j, k).x) > 1e-6) return false;
                if (std::fabs(b(i, j, k).y - a(i, j, k).y) > 1e-6) return false;
                if (std::fabs(b(i, j, k).z - a(i, j, k).z) > 1e-6) return false;
            }
        }
    }
    return true;
}

template <typename T, typename V>
bool is_equal_vector6d(T & a, V & b) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++) {
                if (std::fabs(b(i, j, k).ex - a(i, j, k).ex) > 1e-6) return false;
                if (std::fabs(b(i, j, k).ey - a(i, j, k).ey) > 1e-6) return false;
                if (std::fabs(b(i, j, k).ez - a(i, j, k).ez) > 1e-6) return false;
                if (std::fabs(b(i, j, k).bx - a(i, j, k).bx) > 1e-6) return false;
                if (std::fabs(b(i, j, k).by - a(i, j, k).by) > 1e-6) return false;
                if (std::fabs(b(i, j, k).bz - a(i, j, k).bz) > 1e-6) return false;
            }
        }
    }
    return true;
}

void run_checks(int size) {
    cout << "Checking size " << size << "..." << endl;

    {
        // ordinary checks

        morton_array<double> ex1(size);
        morton_array<double> ey1(size);
        morton_array<double> ez1(size);
        morton_array<double> bx1(size);
        morton_array<double> by1(size);
        morton_array<double> bz1(size);

        morton_array<double> ex2(size);
        morton_array<double> ey2(size);
        morton_array<double> ez2(size);
        morton_array<double> bx2(size);
        morton_array<double> by2(size);
        morton_array<double> bz2(size);

        randomize(ex1);
        randomize(ey1);
        randomize(ez1);

        ex2 = ex1;
        ey2 = ey1;
        ez2 = ez1;

        randomize(bx1);
        randomize(by1);
        randomize(bz1);

        bx2 = bx1;
        by2 = by1;
        bz2 = bz1;

        advance_b(ex1, ey1, ez1, bx1, by1, bz1);
        advance_b_z(ex2, ey2, ez2, bx2, by2, bz2);

        if (!is_equal(bx1, bx2) || !is_equal(by1, by2) || !is_equal(bz1, bz2))
            cout << "advance_b_z is different" << endl;

        randomize(ex1);
        randomize(ey1);
        randomize(ez1);

        ex2 = ex1;
        ey2 = ey1;
        ez2 = ez1;

        randomize(bx1);
        randomize(by1);
        randomize(bz1);

        bx2 = bx1;
        by2 = by1;
        bz2 = bz1;

        advance_e(ex1, ey1, ez1, bx1, by1, bz1);
        advance_e_z(ex2, ey2, ez2, bx2, by2, bz2);

        if (!is_equal(ex1, ex2) || !is_equal(ey1, ey2) || !is_equal(ez1, ez2))
            cout << "advance_e_z is different" << endl;

        randomize(ex1);
        randomize(ey1);
        randomize(ez1);

        ex2 = ex1;
        ey2 = ey1;
        ez2 = ez1;

        randomize(bx1);
        randomize(by1);
        randomize(bz1);

        bx2 = bx1;
        by2 = by1;
        bz2 = bz1;

        advance(ex1, ey1, ez1, bx1, by1, bz1);
        advance_z(ex2, ey2, ez2, bx2, by2, bz2);

        if (!is_equal(ex1, ex2) || !is_equal(ey1, ey2) || !is_equal(ez1, ez2)
                || !is_equal(bx1, bx2) || !is_equal(by1, by2) || !is_equal(bz1, bz2))
            cout << "advance_z is different" << endl;
    }

    {
        // vector3d checks

        morton_array<vector3d> e1(size);
        morton_array<vector3d> e2(size);
        morton_array<vector3d> b1(size);
        morton_array<vector3d> b2(size);

        randomize_vector(e1);
        e2 = e1;
        randomize_vector(b1);
        b2 = b1;

        advance_b_vector(e1, b1);
        advance_b_vector_z(e2, b2);

        if (!is_equal_vector(b1, b2)) cout << "advance_b_vector_z is different" << endl;

        randomize_vector(e1);
        e2 = e1;
        randomize_vector(b1);
        b2 = b1;

        advance_e_vector(e1, b1);
        advance_e_vector_z(e2, b2);

        if (!is_equal_vector(e1, e2)) cout << "advance_e_vector_z is different" << endl;

        randomize_vector(e1);
        e2 = e1;
        randomize_vector(b1);
        b2 = b1;

        advance_vector(e1, b1);
        advance_vector_z(e2, b2);

        if (!is_equal_vector(e1, e2) || !is_equal_vector(b1, b2)) cout << "advance_vector_z is different" << endl;

        randomize_vector(e1);
        e2 = e1;
        randomize_vector(b1);
        b2 = b1;

        advance_b_vector_periodic(e1, b1);
        advance_b_vector_z_periodic(e2, b2);

        if (!is_equal_vector(b1, b2)) cout << "advance_b_vector_z_periodic is different" << endl;

        randomize_vector(e1);
        e2 = e1;
        randomize_vector(b1);
        b2 = b1;

        advance_e_vector_periodic(e1, b1);
        advance_e_vector_z_periodic(e2, b2);

        if (!is_equal_vector(e1, e2)) cout << "advance_e_vector_z_periodic is different" << endl;

        randomize_vector(e1);
        e2 = e1;
        randomize_vector(b1);
        b2 = b1;

        advance_vector_periodic(e1, b1);
        advance_vector_z_periodic(e2, b2);

        if (!is_equal_vector(e1, e2) || !is_equal_vector(b1, b2)) cout << "advance_vector_z_periodic is different" << endl;
    }

    {
        // vector6d checks

        morton_array<vector6d> v1(size);
        morton_array<vector6d> v2(size);

        randomize_vector6d(v1);
        v2 = v1;

        advance_b_vector6d(v1);
        advance_b_vector6d_z(v2);

        if (!is_equal_vector6d(v1, v2)) cout << "advance_b_vector6d_z is different" << endl;

        randomize_vector6d(v1);
        v2 = v1;

        advance_e_vector6d(v1);
        advance_e_vector6d_z(v2);

        if (!is_equal_vector6d(v1, v2)) cout << "advance_e_vector6d_z is different" << endl;

        randomize_vector6d(v1);
        v2 = v1;

        advance_vector6d(v1);
        advance_vector6d_z(v2);

        if (!is_equal_vector6d(v1, v2)) cout << "advance_vector6d_z is different" << endl;
    }

}

template <typename T>
void print_array_vector(T& a) {
    int n = a.get_n();
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++) {
                cout << boost::format("[%6.2f,%6.2f,%6.2f]") % a(i, j, k).x % a(i, j, k).y % a(i, j, k).z;
            }
            cout << endl;
        }
        cout << endl;
    }
}

int main(int argc, char **argv) {
    dt = 0.01;
    double dx = 0.05;
    double dy = 0.05;
    double dz = 0.05;

    dtdx = dt / dx;
    dtdy = dt / dy;
    dtdz = dt / dz;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);

    re.seed(chrono::system_clock::now().time_since_epoch().count());

    const int iterations = 6;

    vector<int> sizes = {32, 64, 128, 256, 512};

    if (argc > 1) {
        sizes.erase(sizes.begin(), sizes.end());
        for (int i = 1; i < argc; i++) {
            sizes.push_back(atoi(argv[i]));
        }
    }

    if (mpi_rank == 0) {
        cout << "Sizes:";
        for (auto size : sizes) {
            cout << " " << size;
        }
        cout << endl;
    }

    if (mpi_rank == 0) {
        for (auto size : sizes) {
            run_checks(size);
        }
    }

    for (auto size : sizes) {
        if (mpi_rank == 0) {
            cout << "Size " << size << endl;
            cout << endl;
        }

        if (mpi_rank == 0) {
            cout << "scalar" << endl;
        }
        run_test<simple_array<double>>(advance_b<simple_array<double>>, "simple(b)", size, iterations);
        run_test<cached_array<double>>(advance_b<cached_array<double>>, "cached(b)", size, iterations);
        run_test<morton_array<double>>(advance_b<morton_array<double>>, "morton_bad(b)", size, iterations);
        run_test<morton_array<double>>(advance_b_z, "morton(b)", size, iterations);

        run_test<simple_array<double>>(advance_e<simple_array<double>>, "simple(e)", size, iterations);
        run_test<cached_array<double>>(advance_e<cached_array<double>>, "cached(e)", size, iterations);
        run_test<morton_array<double>>(advance_e<morton_array<double>>, "morton_bad(e)", size, iterations);
        run_test<morton_array<double>>(advance_e_z, "morton(e)", size, iterations);

        run_test<simple_array<double>>(advance<simple_array<double>>, "simple", size, iterations);
        run_test<cached_array<double>>(advance<cached_array<double>>, "cached", size, iterations);
        run_test<morton_array<double>>(advance<morton_array<double>>, "morton_bad", size, iterations);
        run_test<morton_array<double>>(advance_z, "morton", size, iterations);

        if (mpi_rank == 0) {
            cout << "vector 3d" << endl;
        }
        run_test_vector<simple_array<vector3d>>(advance_b_vector<simple_array<vector3d>>, "simple(b)", size, iterations);
        run_test_vector<cached_array<vector3d>>(advance_b_vector<cached_array<vector3d>>, "cached(b)", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_b_vector<morton_array<vector3d>>, "morton_bad(b)", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_b_vector_z, "morton(b)", size, iterations);

        run_test_vector<simple_array<vector3d>>(advance_e_vector<simple_array<vector3d>>, "simple(e)", size, iterations);
        run_test_vector<cached_array<vector3d>>(advance_e_vector<cached_array<vector3d>>, "cached(e)", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_e_vector<morton_array<vector3d>>, "morton_bad(e)", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_e_vector_z, "morton(e)", size, iterations);

        run_test_vector<simple_array<vector3d>>(advance_vector<simple_array<vector3d>>, "simple", size, iterations);
        run_test_vector<cached_array<vector3d>>(advance_vector<cached_array<vector3d>>, "cached", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_vector<morton_array<vector3d>>, "morton_bad", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_vector_z, "morton", size, iterations);

        run_test_vector<simple_array<vector3d>>(advance_vector_periodic<simple_array<vector3d>>, "simple_p", size, iterations);
        run_test_vector<cached_array<vector3d>>(advance_vector_periodic<cached_array<vector3d>>, "cached_p", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_vector_periodic<morton_array<vector3d>>, "morton_bad_p", size, iterations);
        run_test_vector<morton_array<vector3d>>(advance_vector_z_periodic, "morton_p", size, iterations);

        if (mpi_rank == 0) {
            cout << "vector 6d" << endl;
        }

        run_test_vector6d<simple_array<vector6d>>(advance_b_vector6d<simple_array<vector6d>>, "simple(b)", size, iterations);
        run_test_vector6d<cached_array<vector6d>>(advance_b_vector6d<cached_array<vector6d>>, "cached(b)", size, iterations);
        run_test_vector6d<morton_array<vector6d>>(advance_b_vector6d<morton_array<vector6d>>, "morton_bad(b)", size, iterations);
        run_test_vector6d<morton_array<vector6d>>(advance_b_vector6d_z, "morton(b)", size, iterations);

        run_test_vector6d<simple_array<vector6d>>(advance_e_vector6d<simple_array<vector6d>>, "simple(e)", size, iterations);
        run_test_vector6d<cached_array<vector6d>>(advance_e_vector6d<cached_array<vector6d>>, "cached(e)", size, iterations);
        run_test_vector6d<morton_array<vector6d>>(advance_e_vector6d<morton_array<vector6d>>, "morton_bad(e)", size, iterations);
        run_test_vector6d<morton_array<vector6d>>(advance_e_vector6d_z, "morton(e)", size, iterations);

        run_test_vector6d<simple_array<vector6d>>(advance_vector6d<simple_array<vector6d>>, "simple", size, iterations);
        run_test_vector6d<cached_array<vector6d>>(advance_vector6d<cached_array<vector6d>>, "cached", size, iterations);
        run_test_vector6d<morton_array<vector6d>>(advance_vector6d<morton_array<vector6d>>, "morton_bad", size, iterations);
        run_test_vector6d<morton_array<vector6d>>(advance_vector6d_z, "morton", size, iterations);
    }

    MPI_Finalize();
    return 0;
}


#include <Eigen/Eigen>
#include <catch2/catch.hpp>
#include <entt/entity/registry.hpp>
#include <lmlib/math_constants.h>
#include <lmng/hierarchy.h>
#include <lmng/logging.h>
#include <lmng/meta/reflect_component.h>
#include <lmng/name.h>
#include <lmng/transform.h>

TEST_CASE("Hierarchy system")
{
    entt::registry registry;
    lmng::hierarchy_system hierarchy_system{registry};

    auto parent = registry.create();
    auto child = registry.create();

    registry.emplace<lmng::parent>(child, parent);

    unsigned counter{0};

    for (auto child : lmng::child_range{registry, parent})
    {
        counter++;
    }

    CHECK(counter == 1);

    lmng::transform transform{
      Eigen::Vector3f{1.f, 1.f, 1.f},
      Eigen::Quaternionf{
        Eigen::AngleAxisf{lm::pi / 2, Eigen::Vector3f::UnitY()}}};

    registry.emplace<lmng::transform>(parent, transform);
    registry.emplace<lmng::transform>(child, transform);

    auto composed_transform = lmng::compose_transforms(transform, transform);
    auto child_resolved = lmng::resolve_transform(registry, child);

    CHECK(child_resolved.position == composed_transform.position);
    CHECK(child_resolved.rotation.isApprox(composed_transform.rotation));

    SECTION("Orphan child")
    {
        lmng::orphan_children(registry, parent);
        auto child_transform = registry.get<lmng::transform>(child);
        CHECK(child_transform.position == composed_transform.position);
        CHECK(child_transform.rotation.isApprox(composed_transform.rotation));
    }

    SECTION("Reparent")
    {
        auto new_parent = registry.create();
        lmng::transform new_parent_transform{
          Eigen::Vector3f{10.f, 10.f, 10.f},
          Eigen::Quaternionf{
            Eigen::AngleAxisf{-lm::pi / 4, Eigen::Vector3f::UnitX()}}};

        registry.emplace<lmng::transform>(new_parent, new_parent_transform);

        auto before_reparent_transform =
          lmng::resolve_transform(registry, child);

        lmng::reparent(registry, child, new_parent);

        auto after_reparent_transform =
          lmng::resolve_transform(registry, child);

        CHECK(after_reparent_transform.position.isApprox(
          before_reparent_transform.position));

        auto dist = after_reparent_transform.rotation.angularDistance(
          before_reparent_transform.rotation);
        CAPTURE(dist);
        CHECK(after_reparent_transform.rotation.isApprox(
          before_reparent_transform.rotation));
    }
}

TEST_CASE("Recursive hierarchy range")
{
    entt::registry registry;
    lmng::hierarchy_system hierarchy_system{registry};

    auto parent = registry.create();
    registry.emplace<lmng::name>(parent, "Parent");

    auto child_1 = registry.create();
    registry.emplace<lmng::name>(child_1, "Child 1");
    registry.emplace<lmng::parent>(child_1, parent);

    auto child_1_child_1 = registry.create();
    registry.emplace<lmng::name>(child_1_child_1, "Child 1 child 1");
    registry.emplace<lmng::parent>(child_1_child_1, child_1);

    auto child_1_child_2 = registry.create();
    registry.emplace<lmng::name>(child_1_child_2, "Child 1 child 2");
    registry.emplace<lmng::parent>(child_1_child_2, child_1);

    auto child_2 = registry.create();
    registry.emplace<lmng::name>(child_2, "Child 2");
    registry.emplace<lmng::parent>(child_2, parent);

    auto child_2_child = registry.create();
    registry.emplace<lmng::name>(child_2_child, "Child 2 child");
    registry.emplace<lmng::parent>(child_2_child, child_2);

    unsigned counter{0};

    for (auto curr_child : lmng::recursive_child_range{registry, parent})
    {
        if (counter >= registry.size() - 1)
            FAIL("Too many entities iterated");

        registry.emplace<unsigned>(curr_child, counter++);
    }

    REQUIRE(registry.has<unsigned>(child_1));
    CHECK(registry.get<unsigned>(child_1) == 0);

    REQUIRE(registry.has<unsigned>(child_1_child_1));
    CHECK(registry.get<unsigned>(child_1_child_1) == 1);

    REQUIRE(registry.has<unsigned>(child_1_child_2));
    CHECK(registry.get<unsigned>(child_1_child_2) == 2);

    REQUIRE(registry.has<unsigned>(child_2));
    CHECK(registry.get<unsigned>(child_2) == 3);

    REQUIRE(registry.has<unsigned>(child_2_child));
    CHECK(registry.get<unsigned>(child_2_child) == 4);

    auto childless = registry.create();

    unsigned childless_child_visit_count{0};
    for (auto unused : lmng::recursive_child_range{registry, childless})
    {
        childless_child_visit_count++;
    }

    REQUIRE(childless_child_visit_count == 0);
}

TEST_CASE("Hierarchy manipulation")
{
    entt::registry registry;
    lmng::hierarchy_system hierarchy_system{registry};
    lmng::connect_component_logging(registry);

    auto parent = registry.create();

    std::vector<entt::entity> children{8, entt::entity{entt::null}};

    registry.create(children.begin(), children.end());

    registry.insert(children.begin(), children.end(), lmng::parent{parent});

    registry.destroy(parent);
}


/*
 *  Copyright 2011-2015 Maxim Milakov
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#include "intensity_2d_data_transformer.h"

#include "neural_network_exception.h"
#include "data_transformer_util.h"

#include <opencv2/core/core.hpp>
#include <boost/format.hpp>

namespace nnforge
{
	intensity_2d_data_transformer::intensity_2d_data_transformer(
		float max_contrast_factor,
		float max_absolute_brightness_shift)
	{
		generator = rnd::get_random_generator();

		contrast_distribution = nnforge_uniform_real_distribution<float>(1.0F / max_contrast_factor, max_contrast_factor);
		brightness_shift_distribution = nnforge_uniform_real_distribution<float>(-max_absolute_brightness_shift, max_absolute_brightness_shift);
	}

	intensity_2d_data_transformer::~intensity_2d_data_transformer()
	{
	}

	void intensity_2d_data_transformer::transform(
		const void * data,
		void * data_transformed,
		neuron_data_type::input_type type,
		const layer_configuration_specific& original_config,
		unsigned int sample_id)
	{
		if (type != neuron_data_type::type_byte)
			throw neural_network_exception("intensity_2d_data_transformer is implemented for data stored as bytes only");

		if (original_config.dimension_sizes.size() < 2)
			throw neural_network_exception((boost::format("intensity_2d_data_transformer is processing at least 2d data, data is passed with number of dimensions %1%") % original_config.dimension_sizes.size()).str());

		float contrast = contrast_distribution.min();
		if (contrast_distribution.max() > contrast_distribution.min())
			contrast = contrast_distribution(generator);
		float brightness_shift = brightness_shift_distribution.min() * 255.0F;
		if (brightness_shift_distribution.max() > brightness_shift_distribution.min())
			brightness_shift = brightness_shift_distribution(generator) * 255.0F;

		unsigned int neuron_count_per_image = original_config.dimension_sizes[0] * original_config.dimension_sizes[1];
		unsigned int image_count = original_config.get_neuron_count() / neuron_count_per_image;
		for(unsigned int image_id = 0; image_id < image_count; ++image_id)
		{
			cv::Mat1b image(static_cast<int>(original_config.dimension_sizes[1]), static_cast<int>(original_config.dimension_sizes[0]), static_cast<unsigned char *>(data_transformed) + (image_id * neuron_count_per_image));

			data_transformer_util::change_brightness_and_contrast(
				image,
				contrast,
				brightness_shift);
		}
	}

 	bool intensity_2d_data_transformer::is_deterministic() const
	{
		return false;
	}
}


/**
 *
 * Copyright (c) 2010 Matthias Walter (xammy@xammy.homelinux.net)
 *
 * Authors: Matthias Walter
 *
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 */

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/lookup_edge.hpp>
#include <boost/core/lightweight_test.hpp>
#include <boost/graph/bipartite.hpp>

/// Verifies a 2-coloring

template < typename Graph, typename ColorMap >
void check_two_coloring(const Graph& g, const ColorMap color_map)
{
    typedef boost::graph_traits< Graph > traits;
    typename traits::edge_iterator edge_iter, edge_end;

    for (boost::tie(edge_iter, edge_end) = boost::edges(g);
         edge_iter != edge_end; ++edge_iter)
    {
        typename traits::vertex_descriptor source, target;
        source = boost::source(*edge_iter, g);
        target = boost::target(*edge_iter, g);
        BOOST_TEST(
            boost::get(color_map, source) != boost::get(color_map, target));
    }
}

/// Tests for a vertex sequence to define an odd cycle

template < typename Graph, typename RandomAccessIterator >
void check_odd_cycle(
    const Graph& g, RandomAccessIterator first, RandomAccessIterator beyond)
{
    typedef boost::graph_traits< Graph > traits;

    typename traits::vertex_descriptor first_vertex, current_vertex,
        last_vertex;

    BOOST_TEST((beyond - first) % 2 == 1);

    //  std::cout << "odd_cycle: " << int(*first) << std::endl;

    for (first_vertex = current_vertex = *first++; first != beyond; ++first)
    {
        //    std::cout << "odd_cycle: " << int(*first) << std::endl;

        last_vertex = current_vertex;
        current_vertex = *first;

        BOOST_TEST(
            boost::lookup_edge(current_vertex, last_vertex, g).second);
    }

    BOOST_TEST(boost::lookup_edge(first_vertex, current_vertex, g).second);
}

/// Call the is_bipartite and find_odd_cycle functions and verify their results.

template < typename Graph, typename IndexMap >
void check_bipartite(const Graph& g, IndexMap index_map, bool is_bipartite)
{
    typedef boost::graph_traits< Graph > traits;
    typedef std::vector< boost::default_color_type > partition_t;
    typedef std::vector< typename traits::vertex_descriptor > vertex_vector_t;
    typedef boost::iterator_property_map< partition_t::iterator, IndexMap >
        partition_map_t;

    partition_t partition(boost::num_vertices(g));
    partition_map_t partition_map(partition.begin(), index_map);

    vertex_vector_t odd_cycle(boost::num_vertices(g));

    bool first_result = boost::is_bipartite(g, index_map, partition_map);

    BOOST_TEST(first_result == boost::is_bipartite(g, index_map));

    if (first_result)
        check_two_coloring(g, partition_map);

    BOOST_TEST(first_result == is_bipartite);

    typename vertex_vector_t::iterator second_first = odd_cycle.begin();
    typename vertex_vector_t::iterator second_beyond
        = boost::find_odd_cycle(g, index_map, partition_map, second_first);

    if (is_bipartite)
    {
        BOOST_TEST(second_beyond == second_first);
        check_two_coloring(g, partition_map);
    }
    else
    {
        check_odd_cycle(g, second_first, second_beyond);
    }

    second_beyond = boost::find_odd_cycle(g, index_map, second_first);
    if (is_bipartite)
    {
        BOOST_TEST(second_beyond == second_first);
    }
    else
    {
        check_odd_cycle(g, second_first, second_beyond);
    }
}

int main(int argc, char** argv)
{
    typedef boost::adjacency_list< boost::vecS, boost::vecS,
        boost::undirectedS >
        vector_graph_t;
    typedef boost::adjacency_list< boost::listS, boost::listS,
        boost::undirectedS >
        list_graph_t;
    typedef std::pair< int, int > E;

    typedef std::map< boost::graph_traits< list_graph_t >::vertex_descriptor,
        size_t >
        index_map_t;
    typedef boost::associative_property_map< index_map_t > index_property_map_t;

    /**
     * Create the graph drawn below.
     *
     *       0 - 1 - 2
     *       |       |
     *   3 - 4 - 5 - 6
     *  /      \   /
     *  |        7
     *  |        |
     *  8 - 9 - 10
     **/

    E bipartite_edges[]
        = { E(0, 1), E(0, 4), E(1, 2), E(2, 6), E(3, 4), E(3, 8), E(4, 5),
              E(4, 7), E(5, 6), E(6, 7), E(7, 10), E(8, 9), E(9, 10) };
    vector_graph_t bipartite_vector_graph(&bipartite_edges[0],
        &bipartite_edges[0] + sizeof(bipartite_edges) / sizeof(E), 11);
    list_graph_t bipartite_list_graph(&bipartite_edges[0],
        &bipartite_edges[0] + sizeof(bipartite_edges) / sizeof(E), 11);

    /**
     * Create the graph drawn below.
     *
     *       2 - 1 - 0
     *       |       |
     *   3 - 6 - 5 - 4
     *  /      \   /
     *  |        7
     *  |       /
     *  8 ---- 9
     *
     **/

    E non_bipartite_edges[] = { E(0, 1), E(0, 4), E(1, 2), E(2, 6), E(3, 4),
        E(3, 8), E(4, 5), E(4, 7), E(5, 6), E(6, 7), E(7, 9), E(8, 9) };
    vector_graph_t non_bipartite_vector_graph(&non_bipartite_edges[0],
        &non_bipartite_edges[0] + sizeof(non_bipartite_edges) / sizeof(E), 10);
    list_graph_t non_bipartite_list_graph(&non_bipartite_edges[0],
        &non_bipartite_edges[0] + sizeof(non_bipartite_edges) / sizeof(E), 10);

    /// Create index maps

    index_map_t bipartite_index_map, non_bipartite_index_map;
    boost::graph_traits< list_graph_t >::vertex_iterator vertex_iter,
        vertex_end;
    size_t i = 0;
    for (boost::tie(vertex_iter, vertex_end)
         = boost::vertices(bipartite_list_graph);
         vertex_iter != vertex_end; ++vertex_iter)
    {
        bipartite_index_map[*vertex_iter] = i++;
    }
    index_property_map_t bipartite_index_property_map
        = index_property_map_t(bipartite_index_map);

    i = 0;
    for (boost::tie(vertex_iter, vertex_end)
         = boost::vertices(non_bipartite_list_graph);
         vertex_iter != vertex_end; ++vertex_iter)
    {
        non_bipartite_index_map[*vertex_iter] = i++;
    }
    index_property_map_t non_bipartite_index_property_map
        = index_property_map_t(non_bipartite_index_map);

    /// Call real checks

    check_bipartite(bipartite_vector_graph,
        boost::get(boost::vertex_index, bipartite_vector_graph), true);
    check_bipartite(bipartite_list_graph, bipartite_index_property_map, true);

    check_bipartite(non_bipartite_vector_graph,
        boost::get(boost::vertex_index, non_bipartite_vector_graph), false);
    check_bipartite(
        non_bipartite_list_graph, non_bipartite_index_property_map, false);

    /// Test some more interfaces

    BOOST_TEST(is_bipartite(bipartite_vector_graph));
    BOOST_TEST(!is_bipartite(non_bipartite_vector_graph));

    return boost::report_errors();
}


//Copyright (c) 2015 Zachary Kann
//
//Permission is hereby granted, free of charge, to any person obtaining a copy
//of this software and associated documentation files (the "Software"), to deal
//in the Software without restriction, including without limitation the rights
//to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//copies of the Software, and to permit persons to whom the Software is
//furnished to do so, subject to the following conditions:
//
//The above copyright notice and this permission notice shall be included in all
//copies or substantial portions of the Software.
//
//THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//SOFTWARE.

// ---
// Author: Zachary Kann

// This class is for dealing with time-correlation functions (TCFs) between
// scalars or vectors.

#include <armadillo>

#ifndef _Z_TCF_HPP_
#define _Z_TCF_HPP_

template <class T>
class TCF {
 public:
  TCF(int length, int x_spacing, int interval = 1, int zeros = 0)
      : length_(length), x_spacing_(x_spacing), interval_(interval),
        zeros_(zeros) {
    // TODO(Zak): I think this should be -1 not -2... test that.
    correlation_function_.zeros(length_+zeros_);
  }

  // Calculations correlation between two vector time series assuming
  // time-symmetry around t=0.
  void Correlate(const arma::Mat<T>& mat1, const arma::Mat<T>& mat2);

  // Calculations correlation between two vector time series assuming
  // time-symmetry around t=0. Allows time series to start at arbitrary index
  // and wrap around end of vector.
  void Correlate(const arma::Mat<T>& mat1, const arma::Mat<T>& mat2, const int mod);

  // Calculations correlation between two vector time series without assuming
  // time-symmetry around t=0.
  void CorrelateOneDirection(const arma::Mat<T>& mat1,
                             const arma::Mat<T>& mat2);

  // Calculations correlation between two vector time series without assuming
  // time-symmetry around t=0. Allows time series to start at arbitrary index
  // and wrap around end of vector.
  void CorrelateOneDirection(const arma::Mat<T>& mat1, const arma::Mat<T>& mat2,
                               const int mod);

  // Wrapper for Correlate to avoid ambiguity
  inline void Correlate(const arma::subview_cube<T>& vec1,
                        const arma::subview_cube<T>& vec2) {
    Correlate(arma::Row<T>(vec1), arma::Row<T>(vec2));
  }

  // Wrapper for Correlate to avoid ambiguity
  inline void Correlate(const arma::subview_cube<T>& vec1,
                        const arma::subview_cube<T>& vec2, const int mod) {
    Correlate(arma::Row<T>(vec1), arma::Row<T>(vec2), mod);
  }

  // Wrapper for CorrelateOneDirection to avoid ambiguity
  inline void CorrelateOneDirection(const arma::subview_cube<T>& vec1,
                                    const arma::subview_cube<T>& vec2) {
    CorrelateOneDirection(arma::Row<T>(vec1), arma::Row<T>(vec2));
  }

  // Wrapper for CorrelateOneDirection to avoid ambiguity
  inline void CorrelateOneDirection(const arma::subview_cube<T>& vec1,
                                    const arma::subview_cube<T>& vec2,
                                    const int mod) {
    CorrelateOneDirection(arma::Row<T>(vec1), arma::Row<T>(vec2), mod);
  }

  // Calculations correlation between two scalar time series assuming
  // time-symmetry around t=0.
  void Correlate(const arma::Row<T>& vec1, const arma::Row<T>& vec2);

  // Calculations correlation between two scalar time series assuming
  // time-symmetry around t=0. Allows time series to start at arbitrary index
  // and wrap around end of vector.
  void Correlate(const arma::Row<T>& vec1, const arma::Row<T>& vec2,
                    const int mod);

  // Calculations correlation between two scalar time series without assuming
  // time-symmetry around t=0.
  void CorrelateOneDirection(const arma::Row<T>& vec1,
                             const arma::Row<T>& vec2);

  // Calculations correlation between two scalar time series without assuming
  // time-symmetry around t=0. Allows time series to start at arbitrary index
  // and wrap around end of vector.
  void CorrelateOneDirection(const arma::Row<T>& vec1,
                                  const arma::Row<T>& vec2, const int mod);

  // Wrapper for CorrelateOneDirection that allows vector self-correlation.
  inline void CorrelateOneDirection(const arma::Mat<T>& mat) {
    CorrelateOneDirection(mat, mat);
  }

  // Wrapper for CorrelateOneDirection that allows vector self-correlation and
  // allows time series to start at arbitrary index and wrap around end of
  // vector.
  inline void CorrelateOneDirection(const arma::Mat<T>& vec, const int mod) {
    CorrelateOneDirection(vec, vec, mod);
  }

  // Wrapper for CorrelateOneDirection that allows scalar self-correlation.
  inline void CorrelateOneDirection(const arma::Row<T>& vec) {
    CorrelateOneDirection(vec, vec);
  }

  // Wrapper for CorrelateOneDirection that allows scalar self-correlation and
  // allows time series to start at arbitrary index and wrap around end of
  // vector.
  inline void CorrelateOneDirection(const arma::Row<T>& vec,
                                         const int mod) {
    CorrelateOneDirection(vec, vec, mod);
  }

  inline int CalculateNumCorr(const arma::Row<T>& vec) {
    return (vec.n_elem - length_)/interval_ + 1;
  }

  inline int CalculateNumCorr(const int steps) {
    return (steps - length_)/interval_ + 1;
  }

  inline void MultiplyTCF(const double multiplier) {
    correlation_function_ *= multiplier;
  }

  inline void MultiplyTCF(const int i, const T multiplier) {
    correlation_function_(i) *= multiplier;
  }

  inline void AddTCF(const int i, const T addend) {
    correlation_function_(i) += addend;
  }

  // Mutators
  inline void set_tcf(const int i, const T data) {
    correlation_function_(i) = data;
  }

  // Accessors
  int length() const { return length_; }
  int x_spacing() const { return x_spacing_; }
  int zeros() const { return zeros_; }
  int total_length() const { return length_+zeros_; }
  int interval() const { return interval_; }
  inline T tcf(int i) const { return correlation_function_(i); }
  inline arma::Row<T> tcf() const { return correlation_function_; }

 protected:
  arma::Row<T> correlation_function_;

 private:
  int length_, x_spacing_, interval_, zeros_, number_;

  inline int CalculateNumCorr(const arma::Row<T>& vec1,
                              const arma::Row<T>& vec2) {
    int n1 = vec1.n_elem;
    int n2 = vec2.n_elem;
    int shorter = (n1 < n2) ? n1 : n2;
    return (shorter - length_)/interval_ + 1;
  }

  inline int CalculateNumCorr(const arma::Mat<T>& mat) {
    return (mat.n_rows - length_)/interval_ + 1;
  }

  inline int CalculateNumCorr(const arma::Mat<T>& mat1,
                              const arma::Mat<T>& mat2) {
    int n1 = mat1.n_rows;
    int n2 = mat2.n_rows;
    int shorter = (n1 < n2) ? n1 : n2;
    return (shorter - length_)/interval_ + 1;
  }
};
#endif


// Copyright 2004 The Trustees of Indiana University.

// Use, modification and distribution is subject to the Boost Software
// License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

//  Authors: Douglas Gregor
//           Andrew Lumsdaine
#include <boost/graph/fruchterman_reingold.hpp>
#include <boost/graph/random_layout.hpp>
#include <boost/graph/kamada_kawai_spring_layout.hpp>
#include <boost/graph/circle_layout.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/point_traits.hpp>
#include <boost/random/linear_congruential.hpp>
#include <boost/test/minimal.hpp>
#include <iostream>
#include <boost/limits.hpp>
#include <fstream>
#include <string>
using namespace boost;

enum vertex_position_t { vertex_position };
namespace boost { BOOST_INSTALL_PROPERTY(vertex, position); }

typedef square_topology<>::point_type point;

template<typename Graph, typename PositionMap, typename Topology>
void print_graph_layout(const Graph& g, PositionMap position, const Topology& topology)
{
  typedef typename Topology::point_type Point;
  // Find min/max ranges
  Point min_point = position[*vertices(g).first], max_point = min_point;
  BGL_FORALL_VERTICES_T(v, g, Graph) {
    min_point = topology.pointwise_min(min_point, position[v]);
    max_point = topology.pointwise_max(max_point, position[v]);
  }

  for (int y = (int)min_point[1]; y <= (int)max_point[1]; ++y) {
    for (int x = (int)min_point[0]; x <= (int)max_point[0]; ++x) {
      typename graph_traits<Graph>::vertex_iterator vi, vi_end;
      // Find vertex at this position
      typename graph_traits<Graph>::vertices_size_type index = 0;
      for (tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi, ++index) {
        if ((int)position[*vi][0] == x && (int)position[*vi][1] == y)
          break;
      }

      if (vi == vi_end) std::cout << ' ';
      else std::cout << (char)(index + 'A');
    }
    std::cout << std::endl;
  }
}

template<typename Graph, typename PositionMap>
void dump_graph_layout(std::string name, const Graph& g, PositionMap position)
{
  std::ofstream out((name + ".dot").c_str());
  out << "graph " << name << " {" << std::endl;

  typename graph_traits<Graph>::vertex_iterator vi, vi_end;
  for (tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi) {
    out << "  n" << get(vertex_index, g, *vi) << "[ pos=\""
        << (int)position[*vi][0] + 25 << ", " << (int)position[*vi][1] + 25
        << "\" ];\n";
  }

  typename graph_traits<Graph>::edge_iterator ei, ei_end;
  for (tie(ei, ei_end) = edges(g); ei != ei_end; ++ei) {
    out << "  n" << get(vertex_index, g, source(*ei, g)) << " -- n"
        << get(vertex_index, g, target(*ei, g)) << ";\n";
  }
  out << "}\n";
}

template<typename Graph>
void
test_circle_layout(Graph*, typename graph_traits<Graph>::vertices_size_type n)
{
  typedef typename graph_traits<Graph>::vertex_descriptor vertex;
  typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator;
  typedef typename graph_traits<Graph>::vertices_size_type vertices_size_type;
  typedef typename graph_traits<Graph>::edges_size_type edges_size_type;

  Graph g(n);

  // Initialize vertex indices
  vertex_iterator vi = vertices(g).first;
  for (vertices_size_type i = 0; i < n; ++i, ++vi)
    put(vertex_index, g, *vi, i);

  circle_graph_layout(g, get(vertex_position, g), 10.0);

  std::cout << "Regular polygon layout with " << n << " points.\n";
  square_topology<> topology;
  print_graph_layout(g, get(vertex_position, g), topology);
}

struct simple_edge
{
  int first, second;
};

struct kamada_kawai_done
{
  kamada_kawai_done() : last_delta() {}

  template<typename Graph>
  bool operator()(double delta_p,
                  typename boost::graph_traits<Graph>::vertex_descriptor p,
                  const Graph& g,
                  bool global)
  {
    if (global) {
      double diff = last_delta - delta_p;
      if (diff < 0) diff = -diff;
      last_delta = delta_p;
      return diff < 0.01;
    } else {
      return delta_p < 0.01;
    }
  }

  double last_delta;
};

template<typename Graph>
void
test_triangle(Graph*)
{
  typedef typename graph_traits<Graph>::vertex_descriptor vertex_descriptor;
  typedef typename graph_traits<Graph>::edge_descriptor edge_descriptor;

  Graph g;

  vertex_descriptor u = add_vertex(g); put(vertex_index, g, u, 0);
  vertex_descriptor v = add_vertex(g); put(vertex_index, g, v, 1);
  vertex_descriptor w = add_vertex(g); put(vertex_index, g, w, 2);

  edge_descriptor e1 = add_edge(u, v, g).first; put(edge_weight, g, e1, 1.0);
  edge_descriptor e2 = add_edge(v, w, g).first; put(edge_weight, g, e2, 1.0);
  edge_descriptor e3 = add_edge(w, u, g).first; put(edge_weight, g, e3, 1.0);

  circle_graph_layout(g, get(vertex_position, g), 25.0);

  bool ok = kamada_kawai_spring_layout(g,
                                       get(vertex_position, g),
                                       get(edge_weight, g),
                                       square_topology<>(50.0),
                                       side_length(50.0));
  BOOST_CHECK(ok);

  std::cout << "Triangle layout (Kamada-Kawai).\n";
  print_graph_layout(g, get(vertex_position, g));
}

template<typename Graph>
void
test_cube(Graph*)
{
  enum {A, B, C, D, E, F, G, H};
  simple_edge cube_edges[12] = {
    {A, E}, {A, B}, {A, D}, {B, F}, {B, C}, {C, D}, {C, G}, {D, H},
    {E, H}, {E, F}, {F, G}, {G, H}
  };

  Graph g(&cube_edges[0], &cube_edges[12], 8);

  typedef typename graph_traits<Graph>::edge_iterator edge_iterator;
  typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator;

  vertex_iterator vi, vi_end;
  int i = 0;
  for (tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi)
    put(vertex_index, g, *vi, i++);

  edge_iterator ei, ei_end;
  for (tie(ei, ei_end) = edges(g); ei != ei_end; ++ei) {
    put(edge_weight, g, *ei, 1.0);
    std::cerr << "(" << (char)(get(vertex_index, g, source(*ei, g)) + 'A')
              << ", " << (char)(get(vertex_index, g, target(*ei, g)) + 'A')
              << ") ";
  }
  std::cerr << std::endl;

  circle_graph_layout(g, get(vertex_position, g), 25.0);

  bool ok = kamada_kawai_spring_layout(g,
                                       get(vertex_position, g),
                                       get(edge_weight, g),
                                       square_topology<>(50.0),
                                       side_length(50.0),
                                       kamada_kawai_done());
  BOOST_CHECK(ok);

  std::cout << "Cube layout (Kamada-Kawai).\n";
  print_graph_layout(g, get(vertex_position, g), square_topology<>(50.));

  dump_graph_layout("cube", g, get(vertex_position, g));

  minstd_rand gen;
  typedef square_topology<> Topology;
  Topology topology(gen, 50.0);
  std::vector<Topology::point_difference_type> displacements(num_vertices(g));
  Topology::point_type origin;
  origin[0] = origin[1] = 50.0;
  Topology::point_difference_type extent;
  extent[0] = extent[1] = 50.0;
  rectangle_topology<> rect_top(gen, 0, 0, 50, 50);
  random_graph_layout(g, get(vertex_position, g), rect_top);

  fruchterman_reingold_force_directed_layout
    (g,
     get(vertex_position, g),
     topology,
     origin,
     extent,
     square_distance_attractive_force(),
     square_distance_repulsive_force(),
     all_force_pairs(),
     linear_cooling<double>(100),
     make_iterator_property_map(displacements.begin(),
                                get(vertex_index, g),
                                Topology::point_difference_type()));

  std::cout << "Cube layout (Fruchterman-Reingold).\n";
  print_graph_layout(g, get(vertex_position, g), square_topology<>(50.));

  dump_graph_layout("cube-fr", g, get(vertex_position, g));
}

template<typename Graph>
void
test_triangular(Graph*)
{
  enum {A, B, C, D, E, F, G, H, I, J};
  simple_edge triangular_edges[18] = {
    {A, B}, {A, C}, {B, C}, {B, D}, {B, E}, {C, E}, {C, F}, {D, E}, {D, G},
    {D, H}, {E, F}, {E, H}, {E, I}, {F, I}, {F, J}, {G, H}, {H, I}, {I, J}
  };

  Graph g(&triangular_edges[0], &triangular_edges[18], 10);

  typedef typename graph_traits<Graph>::edge_iterator edge_iterator;
  typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator;

  vertex_iterator vi, vi_end;
  int i = 0;
  for (tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi)
    put(vertex_index, g, *vi, i++);

  edge_iterator ei, ei_end;
  for (tie(ei, ei_end) = edges(g); ei != ei_end; ++ei) {
    put(edge_weight, g, *ei, 1.0);
    std::cerr << "(" << (char)(get(vertex_index, g, source(*ei, g)) + 'A')
              << ", " << (char)(get(vertex_index, g, target(*ei, g)) + 'A')
              << ") ";
  }
  std::cerr << std::endl;

  typedef square_topology<> Topology;
  minstd_rand gen;
  Topology topology(gen, 50.0);
  Topology::point_type origin;
  origin[0] = origin[1] = 50.0;
  Topology::point_difference_type extent;
  extent[0] = extent[1] = 50.0;

  circle_graph_layout(g, get(vertex_position, g), 25.0);

  bool ok = kamada_kawai_spring_layout(g,
                                       get(vertex_position, g),
                                       get(edge_weight, g),
                                       topology,
                                       side_length(50.0),
                                       kamada_kawai_done());
  BOOST_CHECK(ok);

  std::cout << "Triangular layout (Kamada-Kawai).\n";
  print_graph_layout(g, get(vertex_position, g), square_topology<>(50.));

  dump_graph_layout("triangular-kk", g, get(vertex_position, g));

  rectangle_topology<> rect_top(gen, -25, -25, 25, 25);
  random_graph_layout(g, get(vertex_position, g), rect_top);

  dump_graph_layout("random", g, get(vertex_position, g));

  std::vector<Topology::point_difference_type> displacements(num_vertices(g));
  fruchterman_reingold_force_directed_layout
    (g,
     get(vertex_position, g),
     topology,
     origin,
     extent,
     attractive_force(square_distance_attractive_force()).
     cooling(linear_cooling<double>(100)));

  std::cout << "Triangular layout (Fruchterman-Reingold).\n";
  print_graph_layout(g, get(vertex_position, g), square_topology<>(50.));

  dump_graph_layout("triangular-fr", g, get(vertex_position, g));
}

template<typename Graph>
void
test_disconnected(Graph*)
{
  enum {A, B, C, D, E, F, G, H};
  simple_edge triangular_edges[13] = {
    {A, B}, {B, C}, {C, A},
    {D, E}, {E, F}, {F, G}, {G, H}, {H, D},
    {D, F}, {F, H}, {H, E}, {E, G}, {G, D}
  };

  Graph g(&triangular_edges[0], &triangular_edges[13], 8);

  typedef typename graph_traits<Graph>::edge_iterator edge_iterator;
  typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator;

  vertex_iterator vi, vi_end;
  int i = 0;
  for (tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi)
    put(vertex_index, g, *vi, i++);

  edge_iterator ei, ei_end;
  for (tie(ei, ei_end) = edges(g); ei != ei_end; ++ei) {
    put(edge_weight, g, *ei, 1.0);
    std::cerr << "(" << (char)(get(vertex_index, g, source(*ei, g)) + 'A')
              << ", " << (char)(get(vertex_index, g, target(*ei, g)) + 'A')
              << ") ";
  }
  std::cerr << std::endl;

  circle_graph_layout(g, get(vertex_position, g), 25.0);

  bool ok = kamada_kawai_spring_layout(g,
                                       get(vertex_position, g),
                                       get(edge_weight, g),
                                       square_topology<>(50.0),
                                       side_length(50.0),
                                       kamada_kawai_done());
  BOOST_CHECK(!ok);

  minstd_rand gen;
  rectangle_topology<> rect_top(gen, -25, -25, 25, 25);
  random_graph_layout(g, get(vertex_position, g), rect_top);

  typedef square_topology<> Topology;
  Topology topology(gen, 50.0);
  std::vector<Topology::point_difference_type> displacements(num_vertices(g));
  Topology::point_type origin;
  origin[0] = origin[1] = 50.0;
  Topology::point_difference_type extent;
  extent[0] = extent[1] = 50.0;
  fruchterman_reingold_force_directed_layout
    (g,
     get(vertex_position, g),
     topology,
     origin,
     extent,
     attractive_force(square_distance_attractive_force()).
     cooling(linear_cooling<double>(50)));

  std::cout << "Disconnected layout (Fruchterman-Reingold).\n";
  print_graph_layout(g, get(vertex_position, g), square_topology<>(50.));

  dump_graph_layout("disconnected-fr", g, get(vertex_position, g));
}

int test_main(int, char*[])
{
  typedef adjacency_list<listS, listS, undirectedS,
                         // Vertex properties
                         property<vertex_index_t, int,
                         property<vertex_position_t, point> >,
                         // Edge properties
                         property<edge_weight_t, double> > Graph;

  test_circle_layout((Graph*)0, 5);
  test_cube((Graph*)0);
  test_triangular((Graph*)0);
  test_disconnected((Graph*)0);
  return 0;
}


/* Copyright (c) 2021, the adamantine authors.
 *
 * This file is subject to the Modified BSD License and may not be distributed
 * without copyright and license information. Please refer to the file LICENSE
 * for the text and further information on this license.
 */

#ifndef DATA_ASSIMILATOR_HH
#define DATA_ASSIMILATOR_HH

#include <experimental_data.hh>
#include <types.hh>

#include <deal.II/fe/mapping_q1_eulerian.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/la_vector.h>
#include <deal.II/lac/solver_gmres.h>
#include <deal.II/lac/sparse_matrix.h>

#include <map>
#include <random>

namespace adamantine
{
/**
 * Forward declaration of the tester friend class to DataAssimilator.
 */
class DataAssimilatorTester;

/**
 * This class implements an ensemble Kalman filter (EnKF) for data assimilation.
 * It updates an ensemble of solution vectors to be consistent with experimental
 * observations.
 *
 * The EnKF implementation here is largely based on Chapter 6 of Data
 * Assimilation: Method and Applications by Asch, Bocquet, and Nodet.
 */
class DataAssimilator
{
  friend class DataAssimilatorTester;

public:
  /**
   * Constructor.
   */
  DataAssimilator(boost::property_tree::ptree const &database);

  /**
   * This is the main public interface for the class and is called to perform
   * one assimilation process. It takes in an ensemble of simulation data
   * (sim_data), experimental observation data (expt_data), and the covariance
   * matrix for the experimental observations (R). This function updates
   * sim_data using the EnKF method.
   *
   * Overall it performs x = x + K (y + u - H x) for each ensemble member,
   * where:
   * x is the simulation solution
   * K is the Kalman gain
   * y is the observed solution
   * u is a perturbation to the observed solution
   * H is the observation matrix
   */
  void update_ensemble(
      std::vector<dealii::LA::distributed::Vector<double>> &sim_data,
      std::vector<double> const &expt_data, dealii::SparseMatrix<double> &R);

  /**
   * This updates the internal mapping between the indices of the entries in
   * expt_data and the indices of the entries in the sim_data ensemble members
   * in updateEnsemble. This must be called before updateEnsemble whenever there
   * are changes to the simulation mesh or the observation locations. The
   * indices_and_offsets variable is the output from
   * adamantine::set_with_experimental_data.
   */
  template <int dim>
  void update_dof_mapping(
      dealii::DoFHandler<dim> const &dof_handler,
      std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets);

private:
  /**
   * This calculates the Kalman gain and applies it to the perturbed innovation.
   */
  std::vector<dealii::LA::distributed::Vector<double>> apply_kalman_gain(
      std::vector<dealii::LA::distributed::Vector<double>> &vec_ensemble,
      dealii::SparseMatrix<double> &R,
      std::vector<dealii::Vector<double>> &perturbed_innovation);

  /**
   * This calculates the observation matrix.
   */
  dealii::SparseMatrix<double> calc_H(dealii::SparsityPattern &pattern) const;

  /**
   * This calculates the application of the observation matrix on an ensemble
   * member from the simulation (sim_ensemble_member), avoiding explicit
   * construction of a sparse matrix for H.
   */
  dealii::Vector<double> calc_Hx(
      const dealii::LA::distributed::Vector<double> &sim_ensemble_member) const;

  /**
   * This fills a vector (vec) with noise from a multivariate normal
   * distribution defined by a covariance matrix (R).
   */
  void fill_noise_vector(dealii::Vector<double> &vec,
                         dealii::SparseMatrix<double> &R);

  /**
   * This calculates the sample covariance for an input ensemble of vectors
   * (vec_ensemble). Currently this is for a full (i.e. not sparse) matrix. For
   * improved computational performance and reduced spurious correlations a
   * sparse matrix version of this with localization should also be implemented.
   * This is templated so that it can be called on both simulated and
   * experimental data.
   */
  template <typename VectorType>
  dealii::FullMatrix<double>
  calc_sample_covariance_dense(std::vector<VectorType> vec_ensemble) const;

  /**
   * The number of ensemble members in the simulation.
   */
  unsigned int _num_ensemble_members;

  /**
   * The length of the data vector for each simulation ensemble member.
   */
  unsigned int _sim_size;

  /**
   * The length of the data vector the experimental observations.
   */
  unsigned int _expt_size;

  /**
   * The pseudo-random number generator, used for the perturbations to the
   * innovation vectors.
   */
  std::mt19937 _prng;

  /**
   * Random variate generator for a normal distribution, used for the
   * perturbations to the innovation vectors.
   */
  std::normal_distribution<> _normal_dist_generator;

  /**
   * The mapping between the index in the experimental observation data vector
   * to the DoF in the simulation data vectors. This is simpler to use than the
   * indices and offsets that are passed into updateEnsemble, and means that the
   * simulation DoFHandler only needs to be used by updateDofMapping where this
   * variable is set.
   */
  std::pair<std::vector<int>, std::vector<int>> _expt_to_dof_mapping;

  /**
   * Standardized settings for the GMRES solver needed for the matrix inversion
   * in the Kalman gain calculation.
   */
  dealii::SolverControl _solver_control;

  /**
   * Additional settings for the GMRES solver needed for the matrix inversion in
   * the Kalman gain calculation.
   */
  dealii::SolverGMRES<dealii::Vector<double>>::AdditionalData _additional_data;
};
} // namespace adamantine

#endif