/* Copyright (C) 2010 Ion Torrent Systems, Inc. All Rights Reserved */
#include <malloc.h>
#include "Utils.h"
#include "TraceStoreCol.h"
#include "IonH5Eigen.h"
//#define EIGEN_USE_MKL_ALL 1
#include <Eigen/Dense>
#include <Eigen/LU>

#define MIN_SAMPLE_WELL 100
#define SMOOTH_REDUCE_STEP 10
#define SMOOTH_REDUCE_REGION 100
#define INTEGRATION_START 6
#define INTEGRATION_END 12
void TraceStoreCol::WellProj(TraceStoreCol &store,
                             std::vector<KeySeq> & key_vectors,
                             vector<char> &filter,
                             vector<float> &mad) {
  int useable_flows = 0;
  for (size_t i = 0; i < key_vectors.size(); i++) {
    useable_flows = std::max((int)key_vectors[i].usableKeyFlows, useable_flows);
    //    useable_flows = std::max(store.GetNumFlows(), (size_t)useable_flows);
  }
  Eigen::VectorXf norm(store.mFrameStride * useable_flows);
  Eigen::VectorXf sum(store.mFrameStride * useable_flows);

  norm.setZero();
  sum.setZero();
  int start_frame = store.GetNumFrames() - 2;
  int end_frame = store.GetNumFrames();
  int count = 0;
  for (int frame_ix = start_frame; frame_ix < end_frame; frame_ix++) {
    count++;
    int16_t *__restrict data_start = store.GetMemPtr() + frame_ix * store.mFlowFrameStride;
    int16_t *__restrict data_end = data_start + store.mFrameStride * useable_flows;
    float *__restrict norm_start = &norm[0];
    while(data_start != data_end) {
      *norm_start++ += *data_start++;
    }
  }

  // std::vector<ChipReduction> smoothed_avg(useable_flows);
  // int x_clip = mCols;
  // int y_clip = mRows;
  // if (mUseMeshNeighbors == 0) {
  //   x_clip = THUMBNAIL_SIZE;
  //   y_clip = THUMBNAIL_SIZE;
  // }
  // int last_frame = store.GetNumFrames() - 1;
  // for (int flow_ix = 0; flow_ix < useable_flows; flow_ix++) {
  //   smoothed_avg[flow_ix].Init(mRows, mCols, 1,
  //                              SMOOTH_REDUCE_STEP, SMOOTH_REDUCE_STEP,
  //                              y_clip, x_clip,
  //                              SMOOTH_REDUCE_STEP * SMOOTH_REDUCE_STEP * .4);
  //   smoothed_avg[flow_ix].ReduceFrame(&mData[0] + flow_ix * mFrameStride + last_frame  * mFlowFrameStride, &filter[0], 0);
  //   smoothed_avg[flow_ix].SmoothBlocks(SMOOTH_REDUCE_REGION, SMOOTH_REDUCE_REGION);
  // }
  // int x_count = 0;
  // for (int flow_ix = 0; flow_ix < useable_flows; flow_ix++) {
  //   for (size_t row_ix = 0; row_ix < mRows; row_ix++) {
  //     for (size_t col_ix = 0; col_ix < mCols; col_ix++) {
  //       float avg = smoothed_avg[flow_ix].GetSmoothEst(row_ix, col_ix, 0);
  //       norm[x_count] = avg;
  //       x_count++;
  //     }
  //   }
  // }

  norm = norm / count;
  //  start_frame = INTEGRATION_START;
  //  end_frame = INTEGRATION_END;
  start_frame = 5;
  end_frame = store.GetNumFrames() - 6;

  for (int frame_ix = start_frame; frame_ix < end_frame; frame_ix++) {
    int16_t *__restrict data_start = store.GetMemPtr() + frame_ix * store.mFlowFrameStride;
    int16_t *__restrict data_end = data_start + store.mFrameStride * useable_flows;
    float *__restrict norm_start = &norm[0];
    float *__restrict sum_start = &sum[0];
    int well_offset = 0;
    while(data_start != data_end) {
      if (*norm_start == 0.0f) {
        *norm_start = 1.0f;
        // avoid divide by zero but mark that something is wrong with this well..
        if (filter[well_offset] == 0) {
          filter[well_offset] = 6;
        }
      }
      *sum_start += *data_start / *norm_start;
      well_offset++;
      // reset each flow
      if (well_offset == (int) store.mFrameStride) {
        well_offset = 0;
      }
      sum_start++;
      data_start++;
      norm_start++;
    }
  }

  Eigen::MatrixXf flow_mat(store.mFrameStride, useable_flows);
  for (int flow_ix = 0; flow_ix < flow_mat.cols(); flow_ix++) {
    for (size_t well_ix = 0; well_ix < store.mFrameStride; well_ix++) {
      flow_mat(well_ix, flow_ix) = sum(flow_ix * store.mFrameStride + well_ix);
    }
  }

  Eigen::VectorXf n2;
  n2 = flow_mat.rowwise().squaredNorm();
  n2 = n2.array().sqrt();
  float *n2_start = n2.data();
  float *n2_end = n2_start + n2.rows();
  while (n2_start != n2_end) {
    if (*n2_start == 0.0f) { *n2_start = 1.0f; }
    n2_start++;
  }
  for (int flow_ix = 0; flow_ix < flow_mat.cols(); flow_ix++) {
    flow_mat.col(flow_ix).array() = flow_mat.col(flow_ix).array() / n2.array();
  }

  Eigen::VectorXf proj;
  Eigen::VectorXf max_abs_proj(flow_mat.rows());
  max_abs_proj.setZero();
  for (size_t key_ix = 0; key_ix < key_vectors.size(); key_ix++) {
    Eigen::VectorXf key(useable_flows);
    key.setZero();
    for (int f_ix = 0; f_ix < useable_flows; f_ix++) {
      key[f_ix] = key_vectors[key_ix].flows[f_ix];
    }
    proj = flow_mat * key;
    proj = proj.array().abs();
    for (int i = 0; i < proj.rows(); i++) {
      max_abs_proj(i) = max(max_abs_proj(i), proj(i));
    }
  }

  for (int i = 0; i < max_abs_proj.rows(); i++) {
    mad[i] = max_abs_proj(i);
  }

}

void TraceStoreCol::Init(Mask &mask, size_t frames, const char *flowOrder,
       int numFlowsBuff, int maxFlow, int rowStep, int colStep) {

    pthread_mutex_init (&mLock, NULL);
    mUseMeshNeighbors = 1;
    mRowRefStep = rowStep;
    mColRefStep = colStep;
    mMinRefProbes = floor (mRowRefStep * mColRefStep * .1);
    mRows = mask.H();
    mCols = mask.W();
    mFrames = frames;
    mFrameStride = mRows * mCols;
    mFlows = mFlowsBuf = maxFlow;
    mFlowFrameStride = mFrameStride * maxFlow;
    mMaxDist = 2 * sqrt(rowStep*rowStep + colStep+colStep);
    mFlowOrder = flowOrder;

    mWells = mRows * mCols;
    mUseAsReference.resize (mWells, false);
    int keep = 0;
    int empties = 0;
    mRefGridsValid.resize (mFlowsBuf, 0);
    mRefGrids.resize (mFlowsBuf);
    mData.resize (mWells * mFrames * mFlowsBuf);
    std::fill (mData.begin(), mData.end(), 0);
  }

/**
 * Create the basis splines for order requested at a particular value of x.
 * From "A Practical Guide To Splines - Revised Edition" by Carl de Boor p 111
 * the algorithm for BSPLVB. The value of the spline coefficients can be calculated
 * based on a cool recursion. This is a simplified and numerically stable algorithm
 * that seems to be the basis of most bspline implementations including gsl, matlab,etc.
 * @param all_knots - Original knots vector augmented with order number of endpoints
 * @param n_knots - Total number of knots.
 * @param order - Order or order of polynomials, 4 for cubic splines.
 * @param i - index of value x in knots such that all_knots[i] <= x && all_knots[i+1] > x
 * @param x - value that we want to evaluate spline at.
 * @param b - memory to write our coefficients into, must be at least order long.
 */
void basis_spline_xi_v2(float *all_knots, int n_knots, int order, int i, float x, float *b) {
  b[0] = 1;
  double kr[order];
  double kl[order];
  double term;
  double saved;
  for (int j = 0; j < order - 1; j++) {
    kr[j] = all_knots[i + j + 1] - x; // k right values
    kl[j] = x - all_knots[i - j];     // k left values
    saved = 0;
    for (int r = 0; r <= j; r++) {
      term = b[r] / (kr[r] + kl[j-r]);
      b[r] = saved + kr[r] * term;
      saved = kl[j-r] * term;
    }
    b[j+1] = saved;
  }
}

/**
 * Create the augmented knots vector and fill in matrix Xt with the spline basis vectors
 */
void basis_splines_endreps_local_v2(float *knots, int n_knots, int order, int *boundaries, int n_boundaries,  Eigen::MatrixXf &Xt) {
  assert(n_boundaries == 2 && boundaries[0] < boundaries[1]);
  int n_rows = boundaries[1] - boundaries[0]; // this is frames so evaluate at each frame we'll need
  int n_cols = n_knots + order;  // number of basis vectors we'll have at end
  Xt.resize(n_cols, n_rows); // swapped to transpose later. This order lets us use it as a scratch space
  Xt.fill(0);
  int n_std_knots = n_knots + 2 * order;
  float std_knots[n_std_knots];

  // Repeat the boundary knots on there to ensure linear out side of knots
  for (int i = 0; i < order; i++) {
    std_knots[i] = boundaries[0];
  }
  // Our original boundary knots here
  for (int i = 0; i < n_knots; i++) {
    std_knots[i+order] = knots[i];
  }
  // Repeat the boundary knots on there to ensure linear out side of knots
  for (int i = 0; i < order; i++) {
    std_knots[i+order+n_knots] = boundaries[1];
  }
  // Evaluate our basis splines at each frame.
  for (int i = boundaries[0]; i < boundaries[1]; i++) {
    int idx = -1;
    // find index such that i >= knots[idx] && i < knots[idx+1]
    float *val = std::upper_bound(std_knots, std_knots + n_std_knots - 1, 1.0f * i);
    idx = val - std_knots - 1;
    assert(idx >= 0);
    float *f = Xt.data() + i * n_cols + idx - (order - 1); //column offset
    basis_spline_xi_v2(std_knots, n_std_knots, order, idx, i, f);
  }
  // Put in our conventional format where each column is a basis vector
  Xt.transposeInPlace();
}

void FillInKnots(const std::string &strategy, int n_frame, std::vector<float> &knots) {
  if (strategy != "no-knots") {
    if (strategy == "even4") {
      float stride = n_frame / 3.0f;
      knots.push_back(stride);
      knots.push_back(stride * 2.0f);
    }
    else if (strategy == "every4") {
      int current = 4;
      while (current < (n_frame -1)) {
        knots.push_back(current);
        current+=4;
      }
    }
    else if (strategy == "every3") {
      int current = 3;
      while (current < (n_frame -1)) {
        knots.push_back(current);
        current+=3;
      }
    }
    else if (strategy == "middle") {
      float stride = n_frame / 2.0f;
      knots.push_back(stride);
    }
    else if (strategy.find("explicit:") == 0) {
      string spec = strategy.substr(9, strategy.length() - 9);
      knots = char2Vec<float>(spec.c_str(), ',');
    }
    else {
      assert(false);
    }
  }
}

void TraceStoreCol::SplineLossyCompress(const std::string &strategy, int order, char *bad_wells, float *mad) {
  Eigen::MatrixXf Basis;
  vector<float> knots;
  FillInKnots(strategy, mFrames, knots);
  //  Eigen::Map<Eigen::MatrixXf, Eigen::Aligned> Basis(, compressed.n_frames, compressed.n_basis);
  if (!knots.empty()) {
    int boundaries[2];
    boundaries[0] = 0;
    boundaries[1] = mFrames;
    basis_splines_endreps_local_v2(&knots[0], knots.size(), order, boundaries, sizeof(boundaries)/sizeof(boundaries[0]), Basis);
  }
  Eigen::MatrixXf SX = (Basis.transpose() * Basis).inverse() * Basis.transpose();
  Eigen::MatrixXf Y(mFlowFrameStride, mFrames);
  //  Eigen::MatrixXf FlowMeans(mFlows, mFrames);

  //  FlowMeans.setZero();
  int good_wells = 0;
  char *bad_start = bad_wells;
  char *bad_end = bad_start + mFrameStride;
  while(bad_start != bad_end) {
    if (*bad_start++ == 0) {
      good_wells++;
    }
  }

  // if nothing good then skip it
  if (good_wells < MIN_SAMPLE_WELL) {
    return;
  }

  std::vector<ChipReduction> smoothed_avg(mFlows);
  int x_clip = mCols;
  int y_clip = mRows;
  if (mUseMeshNeighbors == 0) {
    x_clip = THUMBNAIL_SIZE;
    y_clip = THUMBNAIL_SIZE;
  }
  for (size_t flow_ix = 0; flow_ix < mFlows; flow_ix++) {
    smoothed_avg[flow_ix].Init(mRows, mCols, mFrames,
                               SMOOTH_REDUCE_STEP, SMOOTH_REDUCE_STEP,
                               y_clip, x_clip,
                               SMOOTH_REDUCE_STEP * SMOOTH_REDUCE_STEP * .4);
    for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
      smoothed_avg[flow_ix].ReduceFrame(&mData[0] + flow_ix * mFrameStride + frame_ix * mFlowFrameStride, bad_wells, frame_ix);
    }
    smoothed_avg[flow_ix].SmoothBlocks(SMOOTH_REDUCE_REGION, SMOOTH_REDUCE_REGION);
  }

  float *y_start = Y.data();
  float *y_end = Y.data() + mData.size();
  int16_t *trace_start = &mData[0];
  while(y_start != y_end) {
    *y_start++ = *trace_start++;
  }

  // // get the flow means per frame.
  // for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
  //   for(size_t flow_ix = 0; flow_ix < mFlows; flow_ix++) {
  //     float *start = Y.data() + mFlowFrameStride * frame_ix + flow_ix * mFrameStride;
  //     float *end = start + mFrameStride;
  //     float *sum = &FlowMeans(flow_ix, frame_ix);
  //     char *bad = bad_wells;
  //     while (start != end) {
  //       if (*bad == 0) {
  //         *sum += *start;
  //       }
  //       start++;
  //       bad++;
  //     }
  //     *sum /= good_wells;
  //   }
  // }

  // subtract off flow,frame avg
  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    for(size_t flow_ix = 0; flow_ix < mFlows; flow_ix++) {
      float *start = Y.data() + mFlowFrameStride * frame_ix + flow_ix * mFrameStride;
      float *end = start + mFrameStride;
      for (size_t row = 0; row < mRows; row++) {
        for (size_t col = 0; col < mCols; col++) {
          float avg = smoothed_avg[flow_ix].GetSmoothEst(row, col, frame_ix);
          *start++ -= avg;
        }
      }
    }
  }

  // // subtract them off
  // Eigen::VectorXf col_mean = Y.colwise().sum();
  // col_mean /= Y.rows();

  // for (int i = 0; i < Y.cols(); i++) {
  //   Y.col(i).array() -= col_mean.coeff(i);
  // }

  // Get coefficients to solve
  Eigen::MatrixXf B = Y * SX.transpose();
  // Uncompress data into yhat matrix
  Eigen::MatrixXf Yhat = B * Basis.transpose();

  // add the flow/frame averages back
  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    for(size_t flow_ix = 0; flow_ix < mFlows; flow_ix++) {
      float *start = Y.data() + mFlowFrameStride * frame_ix + flow_ix * mFrameStride;
      float *end = start + mFrameStride;
      float *hstart = Yhat.data() + mFlowFrameStride * frame_ix + flow_ix * mFrameStride;
      for (size_t row = 0; row < mRows; row++) {
        for (size_t col = 0; col < mCols; col++) {
          float avg = smoothed_avg[flow_ix].GetSmoothEst(row, col, frame_ix);
          *start++ += avg;
          *hstart++ += avg;
        }
      }
    }
  }

  // for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
  //   for(size_t flow_ix = 0; flow_ix < mFlows; flow_ix++) {
  //     float *start = Y.data() + mFlowFrameStride * frame_ix + flow_ix * mFrameStride;
  //     float *hstart = Yhat.data() + mFlowFrameStride * frame_ix + flow_ix * mFrameStride;
  //     float *end = start + mFrameStride;
  //     float avg = FlowMeans(flow_ix, frame_ix);
  //     while (start != end) {
  //       *start++ += avg;
  //       *hstart++ += avg;
  //     }
  //   }
  // }

  // for (int i = 0; i < Yhat.cols(); i++) {
  //   Yhat.col(i).array() += col_mean.coeff(i);
  //   Y.col(i).array() += col_mean.coeff(i);
  // }

  float *yhat_start = Yhat.data();
  float *yhat_end = Yhat.data() + mData.size();
  trace_start = &mData[0];
  while(yhat_start != yhat_end) {
    *trace_start++ = (int)(*yhat_start + .5);
    yhat_start++;
  }

  Y = Y - Yhat;
  Eigen::VectorXf M = Y.rowwise().squaredNorm();

  for (size_t flow_ix = 0; flow_ix < mFlows; flow_ix++) {
    float *mad_start = mad;
    float *mad_end = mad_start + mFrameStride;
    float *m_start = M.data() + flow_ix * mFrameStride;
    while (mad_start != mad_end) {
      *mad_start += *m_start;
      mad_start++;
      m_start++;
    }
  }

  float *mad_start = mad;
  float *mad_end = mad + mFrameStride;
  int norm_factor = mFlows * mFrames;
  while (mad_start != mad_end) {
    *mad_start /= norm_factor;
    mad_start++;
  }
}

void TraceStoreCol::SplineLossyCompress(const std::string &strategy, int order, int flow_ix, char *bad_wells, float *mad) {
  Eigen::MatrixXf Basis;
  vector<float> knots;
  FillInKnots(strategy, mFrames, knots);
  //  Eigen::Map<Eigen::MatrixXf, Eigen::Aligned> Basis(, compressed.n_frames, compressed.n_basis);
  if (!knots.empty()) {
    int boundaries[2];
    boundaries[0] = 0;
    boundaries[1] = mFrames;
    basis_splines_endreps_local_v2(&knots[0], knots.size(), order, boundaries, sizeof(boundaries)/sizeof(boundaries[0]), Basis);
  }
  Eigen::MatrixXf SX = (Basis.transpose() * Basis).inverse() * Basis.transpose();
  Eigen::MatrixXf Y(mFrameStride, mFrames);
  //  Eigen::MatrixXf FlowMeans(mFlows, mFrames);

  //  FlowMeans.setZero();
  int good_wells = 0;
  char *bad_start = bad_wells;
  char *bad_end = bad_start + mFrameStride;
  while(bad_start != bad_end) {
    if (*bad_start++ == 0) {
      good_wells++;
    }
  }

  // if nothing good then skip it
  if (good_wells < MIN_SAMPLE_WELL) {
    return;
  }

  ChipReduction smoothed_avg;
  int x_clip = mCols;
  int y_clip = mRows;
  if (mUseMeshNeighbors == 0) {
    x_clip = THUMBNAIL_SIZE;
    y_clip = THUMBNAIL_SIZE;
  }

  smoothed_avg.Init(mRows, mCols, mFrames,
                    SMOOTH_REDUCE_STEP, SMOOTH_REDUCE_STEP,
                    y_clip, x_clip,
                    SMOOTH_REDUCE_STEP * SMOOTH_REDUCE_STEP * .4);
  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    smoothed_avg.ReduceFrame(&mData[0] + flow_ix * mFrameStride + frame_ix * mFlowFrameStride, bad_wells, frame_ix);
  }
  smoothed_avg.SmoothBlocks(SMOOTH_REDUCE_REGION, SMOOTH_REDUCE_REGION);

  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    float *y_start = Y.data() + frame_ix * mFrameStride;
    float *y_end = y_start + mFrameStride;
    int16_t *trace_start = &mData[0] + flow_ix * mFrameStride + frame_ix * mFlowFrameStride;
    while(y_start != y_end) {
      *y_start++ = *trace_start++;
    }
  }

  // subtract off flow,frame avg
  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    float *start = Y.data() + mFrameStride * frame_ix;
    float *end = start + mFrameStride;
    for (size_t row = 0; row < mRows; row++) {
      for (size_t col = 0; col < mCols; col++) {
        float avg = smoothed_avg.GetSmoothEst(row, col, frame_ix);
        *start++ -= avg;
      }
    }
  }

  // Get coefficients to solve
  Eigen::MatrixXf B = Y * SX.transpose();
  // Uncompress data into yhat matrix
  Eigen::MatrixXf Yhat = B * Basis.transpose();

  // add the flow/frame averages back
  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    float *start = Y.data() + mFrameStride * frame_ix;
    float *end = start + mFrameStride;
    float *hstart = Yhat.data() + mFrameStride * frame_ix;
    for (size_t row = 0; row < mRows; row++) {
      for (size_t col = 0; col < mCols; col++) {
        float avg = smoothed_avg.GetSmoothEst(row, col, frame_ix);
        *start++ += avg;
        *hstart++ += avg;
      }
    }
  }

  for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
    float *yhat_start = Yhat.data() + frame_ix * mFrameStride;
    float *yhat_end = yhat_start + mFrameStride;
    int16_t *trace_start = &mData[0] + flow_ix * mFrameStride + frame_ix * mFlowFrameStride;
    while(yhat_start != yhat_end) {
      *trace_start++ = (int)(*yhat_start + .5f);
      yhat_start++;
    }
  }

  Y = Y - Yhat;
  Eigen::VectorXf M = Y.rowwise().squaredNorm();

  float *mad_start = mad;
  float *mad_end = mad_start + mFrameStride;
  float *m_start = M.data();
  while (mad_start != mad_end) {
    *mad_start += *m_start;
    mad_start++;
    m_start++;
  }

  mad_start = mad;
  mad_end = mad + mFrameStride;
  int norm_factor = mFrames;// mFlows * mFrames;
  while (mad_start != mad_end) {
    *mad_start /= norm_factor;
    mad_start++;
  }
}

// Compress a block of a data using pca
void TraceStoreCol::PcaLossyCompress(int row_start, int row_end,
                                     int col_start, int col_end,
                                     int flow_ix,
                                     float *ssq, char *filters,
                                     int row_step, int col_step,
                                     int num_pca) {

  Eigen::MatrixXf Ysub, Y, S, Basis;

  int loc_num_wells = (col_end - col_start) * (row_end - row_start);
  int loc_num_cols = col_end - col_start;

  // Take a sample of the data at rate step, avoiding flagged wells
  // Count the good rows
  int sample_wells = 0;
  for (int row_ix = row_start; row_ix < row_end; row_ix+= row_step) {
    char *filt_start = filters + row_ix * mCols + col_start;
    char *filt_end = filt_start + loc_num_cols;
    while (filt_start < filt_end) {
      if (*filt_start == 0) {
        sample_wells++;
      }
      filt_start += col_step;
    }
  }
  // try backing off to every well rather than just sampled if we didn't get enough
  if (sample_wells < MIN_SAMPLE_WELL) {
    row_step = 1;
    col_step = 1;
    int sample_wells = 0;
    for (int row_ix = row_start; row_ix < row_end; row_ix+= row_step) {
      char *filt_start = filters + row_ix * mCols + col_start;
      char *filt_end = filt_start + loc_num_cols;
      while (filt_start < filt_end) {
        if (*filt_start == 0) {
          sample_wells++;
        }
        filt_start += col_step;
      }
    }
  }

  if (sample_wells < MIN_SAMPLE_WELL) {
    return; // just give up
  }

  // Copy the sampled data in Matrix, frame major
  Ysub.resize(sample_wells, mFrames);
  for (int frame_ix = 0; frame_ix < (int)mFrames; frame_ix++) {
    int sample_offset = 0;
    for (int row_ix = row_start; row_ix < row_end; row_ix+=row_step) {
      size_t store_offset = row_ix * mCols + col_start;
      char *filt_start = filters + store_offset;
      char *filt_end = filt_start + loc_num_cols;
      int16_t *trace_start = &mData[0] + (mFlowFrameStride * frame_ix) + (flow_ix * mFrameStride) + store_offset;
      float *ysub_start = Ysub.data() + sample_wells * frame_ix + sample_offset;
      while (filt_start < filt_end) {
        if (*filt_start == 0) {
          *ysub_start = *trace_start;
          ysub_start++;
          sample_offset++;
        }
        trace_start += col_step;
        filt_start += col_step;
      }
    }
  }

  // Copy in all the data into working matrix
  Y.resize(loc_num_wells, (int)mFrames);
  for (int frame_ix = 0; frame_ix < (int)mFrames; frame_ix++) {
    for (int row_ix = row_start; row_ix < row_end; row_ix++) {
      size_t store_offset = row_ix * mCols + col_start;
      int16_t *trace_start = &mData[0] + (mFlowFrameStride * frame_ix) + (flow_ix * mFrameStride) + store_offset;
      int16_t *trace_end = trace_start + loc_num_cols;
      float * y_start = Y.data() + loc_num_wells * frame_ix + (row_ix - row_start) * loc_num_cols;
      while( trace_start != trace_end ) {
        *y_start++ = *trace_start++;
      }
    }
  }
  Eigen::VectorXf col_mean = Y.colwise().sum();
  col_mean /= Y.rows();

  for (int i = 0; i < Y.cols(); i++) {
    Y.col(i).array() -= col_mean.coeff(i);
  }
  // Create scatter matrix
  S = Ysub.transpose() * Ysub;
  // Compute the eigenvectors
  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXf> es;
  es.compute(S);
  Eigen::MatrixXf Pca_Basis = es.eigenvectors();
  Eigen::VectorXf Pca_Values = es.eigenvalues();
  // Copy top eigen vectors into basis for projection
  Basis.resize(mFrames, num_pca);
  for (int i = 0; i < Basis.cols(); i++) {
    //    Basis.col(i) = es.eigenvectors().col(es.eigenvectors().cols() - i -1);
    Basis.col(i) = Pca_Basis.col(Pca_Basis.cols() - i - 1);
  }
  // Create solver matrix, often not a good way of solving things but eigen vectors should be stable and fast
  Eigen::MatrixXf SX = (Basis.transpose() * Basis).inverse() * Basis.transpose();
  // Get coefficients to solve
  Eigen::MatrixXf B = Y * SX.transpose();
  // Uncompress data into yhat matrix
  Eigen::MatrixXf Yhat = B * Basis.transpose();

  for (int i = 0; i < Yhat.cols(); i++) {
    Yhat.col(i).array() += col_mean.coeff(i);
    Y.col(i).array() += col_mean.coeff(i);
  }

  // H5File h5("pca_lossy.h5");
  // h5.Open();
  // char buff[256];
  // snprintf(buff, sizeof(buff), "/Y_%d_%d_%d", flow_ix, row_start, col_start);
  // H5Eigen::WriteMatrix(h5, buff, Y);
  // snprintf(buff, sizeof(buff), "/Yhat_%d_%d_%d", flow_ix, row_start, col_start);
  // H5Eigen::WriteMatrix(h5, buff, Yhat);
  // snprintf(buff, sizeof(buff), "/Basis_%d_%d_%d", flow_ix, row_start, col_start);
  // H5Eigen::WriteMatrix(h5, buff, Basis);
  // h5.Close();
  // Copy data out of yhat matrix into original data structure, keeping track of residuals
  for (int frame_ix = 0; frame_ix < (int)mFrames; frame_ix++) {
    for (int row_ix = row_start; row_ix < row_end; row_ix++) {
      size_t store_offset = row_ix * mCols + col_start;
      int16_t *trace_start = &mData[0] + mFlowFrameStride * frame_ix + flow_ix * mFrameStride + store_offset;
      int16_t *trace_end = trace_start + loc_num_cols;
      float * ssq_start = ssq + store_offset;
      size_t loc_offset = (row_ix - row_start) * loc_num_cols;
      float * y_start = Y.data() + loc_num_wells * frame_ix + loc_offset;
      float * yhat_start = Yhat.data() + loc_num_wells * frame_ix + loc_offset;
      while( trace_start != trace_end ) {
        *trace_start = (int16_t)(*yhat_start + .5);
        float val = *y_start - *yhat_start;
        *ssq_start += val * val;
        y_start++;
        yhat_start++;
        trace_start++;
        ssq_start++;
      }
    }
  }

  // divide ssq data out for per frame avg
  for (int row_ix = row_start; row_ix < row_end; row_ix++) {
    size_t store_offset = row_ix * mCols + col_start;
    float * ssq_start = ssq + store_offset;
    float * ssq_end = ssq_start + loc_num_cols;
    while (ssq_start != ssq_end) {
      *ssq_start /= mFrames;
      ssq_start++;
    }
  }
}

bool TraceStoreCol::PcaLossyCompressChunk(int row_start, int row_end,
                                          int col_start, int col_end,
                                          int num_rows, int num_cols, int num_frames,
                                          int frame_stride,
                                          int flow_ix, int flow_frame_stride,
                                          short *data, bool replace,
                                          float *ssq, char *filters,
                                          int row_step, int col_step,
                                          int num_pca) {

  Eigen::MatrixXf Ysub, Y, S, Basis;

  int loc_num_wells = (col_end - col_start) * (row_end - row_start);
  int loc_num_cols = col_end - col_start;

  // Take a sample of the data at rate step, avoiding flagged wells
  // Count the good rows
  int sample_wells = 0;
  for (int row_ix = row_start; row_ix < row_end; row_ix+= row_step) {
    char *filt_start = filters + row_ix * num_cols + col_start;
    char *filt_end = filt_start + loc_num_cols;
    while (filt_start < filt_end) {
      if (*filt_start == 0) {
        sample_wells++;
      }
      filt_start += col_step;
    }
  }
  // try backing off to every well rather than just sampled if we didn't get enough
  if (sample_wells < MIN_SAMPLE_WELL) {
    row_step = 1;
    col_step = 1;
    int sample_wells = 0;
    for (int row_ix = row_start; row_ix < row_end; row_ix+= row_step) {
      char *filt_start = filters + row_ix * num_cols + col_start;
      char *filt_end = filt_start + loc_num_cols;
      while (filt_start < filt_end) {
        if (*filt_start == 0) {
          sample_wells++;
        }
        filt_start += col_step;
      }
    }
  }

  if (sample_wells < MIN_SAMPLE_WELL) {
    return false; // just give up
  }

  // Got enough data to work with, zero out the ssq array for accumulation
  for (int row_ix = row_start; row_ix < row_end; row_ix++) {
    float *ssq_start = ssq + row_ix * num_cols + col_start;
    float *ssq_end = ssq_start + loc_num_cols;
    while (ssq_start != ssq_end) {
      *ssq_start++ = 0;
    }
  }
  // Copy the sampled data in Matrix, frame major
  Ysub.resize(sample_wells, num_frames);
  for (int frame_ix = 0; frame_ix < (int)num_frames; frame_ix++) {
    int sample_offset = 0;
    for (int row_ix = row_start; row_ix < row_end; row_ix+=row_step) {
      size_t store_offset = row_ix * num_cols + col_start;
      char *filt_start = filters + store_offset;
      char *filt_end = filt_start + loc_num_cols;
      int16_t *trace_start = data + (flow_frame_stride * frame_ix) + (flow_ix * frame_stride) + store_offset;
      float *ysub_start = Ysub.data() + sample_wells * frame_ix + sample_offset;
      while (filt_start < filt_end) {
        if (*filt_start == 0) {
          *ysub_start = *trace_start;
          ysub_start++;
          sample_offset++;
        }
        trace_start += col_step;
        filt_start += col_step;
      }
    }
  }

  // Copy in all the data into working matrix
  Y.resize(loc_num_wells, (int)num_frames);
  for (int frame_ix = 0; frame_ix < (int)num_frames; frame_ix++) {
    for (int row_ix = row_start; row_ix < row_end; row_ix++) {
      size_t store_offset = row_ix * num_cols + col_start;
      int16_t *trace_start = data + (flow_frame_stride * frame_ix) + (flow_ix * frame_stride) + store_offset;
      int16_t *trace_end = trace_start + loc_num_cols;
      float * y_start = Y.data() + loc_num_wells * frame_ix + (row_ix - row_start) * loc_num_cols;
      while( trace_start != trace_end ) {
        *y_start++ = *trace_start++;
      }
    }
  }
  Eigen::VectorXf col_mean = Y.colwise().sum();
  col_mean /= Y.rows();

  for (int i = 0; i < Y.cols(); i++) {
    Y.col(i).array() -= col_mean.coeff(i);
  }
  // Create scatter matrix
  S = Ysub.transpose() * Ysub;
  // Compute the eigenvectors
  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXf> es;
  es.compute(S);
  Eigen::MatrixXf Pca_Basis = es.eigenvectors();
  Eigen::VectorXf Pca_Values = es.eigenvalues();
  // Copy top eigen vectors into basis for projection
  Basis.resize(num_frames, num_pca);
  for (int i = 0; i < Basis.cols(); i++) {
    //    Basis.col(i) = es.eigenvectors().col(es.eigenvectors().cols() - i -1);
    Basis.col(i) = Pca_Basis.col(Pca_Basis.cols() - i - 1);
  }
  // Create solver matrix, often not a good way of solving things but eigen vectors should be stable and fast
  Eigen::MatrixXf SX = (Basis.transpose() * Basis).inverse() * Basis.transpose();
  // Get coefficients to solve
  Eigen::MatrixXf B = Y * SX.transpose();
  // Uncompress data into yhat matrix
  Eigen::MatrixXf Yhat = B * Basis.transpose();

  for (int i = 0; i < Yhat.cols(); i++) {
    Yhat.col(i).array() += col_mean.coeff(i);
    Y.col(i).array() += col_mean.coeff(i);
  }

  // H5File h5("pca_lossy.h5");
  // h5.Open();
  // char buff[256];
  // snprintf(buff, sizeof(buff), "/Y_%d_%d_%d", flow_ix, row_start, col_start);
  // H5Eigen::WriteMatrix(h5, buff, Y);
  // snprintf(buff, sizeof(buff), "/Yhat_%d_%d_%d", flow_ix, row_start, col_start);
  // H5Eigen::WriteMatrix(h5, buff, Yhat);
  // snprintf(buff, sizeof(buff), "/Basis_%d_%d_%d", flow_ix, row_start, col_start);
  // H5Eigen::WriteMatrix(h5, buff, Basis);
  // h5.Close();
  // Copy data out of yhat matrix into original data structure, keeping track of residuals
  for (int frame_ix = 0; frame_ix < (int)num_frames; frame_ix++) {
    for (int row_ix = row_start; row_ix < row_end; row_ix++) {
      size_t store_offset = row_ix * num_cols + col_start;
      int16_t *trace_start = data + flow_frame_stride * frame_ix + flow_ix * frame_stride + store_offset;
      int16_t *trace_end = trace_start + loc_num_cols;
      float * ssq_start = ssq + store_offset;
      size_t loc_offset = (row_ix - row_start) * loc_num_cols;
      float * y_start = Y.data() + loc_num_wells * frame_ix + loc_offset;
      float * yhat_start = Yhat.data() + loc_num_wells * frame_ix + loc_offset;
      while( trace_start != trace_end ) {
        if (replace) {
          *trace_start = (int16_t)(*yhat_start + .5);
        }
        float val = *y_start - *yhat_start;
        *ssq_start += val * val;
        y_start++;
        yhat_start++;
        trace_start++;
        ssq_start++;
      }
    }
  }

  // divide ssq data out for per frame avg
  for (int row_ix = row_start; row_ix < row_end; row_ix++) {
    size_t store_offset = row_ix * num_cols + col_start;
    float * ssq_start = ssq + store_offset;
    float * ssq_end = ssq_start + loc_num_cols;
    while (ssq_start != ssq_end) {
      *ssq_start /= num_frames;
      ssq_start++;
    }
  }
  return true;
}

void TraceStoreCol::SplineLossyCompress(const std::string &strategy, int order, int flow_ix, char *bad_wells,
                                        float *mad, size_t num_rows, size_t num_cols, size_t num_frames, size_t num_flows,
                                        int use_mesh_neighbors, size_t frame_stride, size_t flow_frame_stride, int16_t *data) {
  Eigen::MatrixXf Basis;
  vector<float> knots;
  FillInKnots(strategy, num_frames, knots);
  if (!knots.empty()) {
    int boundaries[2];
    boundaries[0] = 0;
    boundaries[1] = num_frames;
    basis_splines_endreps_local_v2(&knots[0], knots.size(), order, boundaries, sizeof(boundaries)/sizeof(boundaries[0]), Basis);
  }
  Eigen::MatrixXf SX = (Basis.transpose() * Basis).inverse() * Basis.transpose();
  Eigen::MatrixXf Y(frame_stride, num_frames);
  //  Eigen::MatrixXf FlowMeans(num_flows, num_frames);

  //  FlowMeans.setZero();
  int good_wells = 0;
  char *bad_start = bad_wells;
  char *bad_end = bad_start + frame_stride;
  while(bad_start != bad_end) {
    if (*bad_start++ == 0) {
      good_wells++;
    }
  }

  // if nothing good then skip it
  if (good_wells < MIN_SAMPLE_WELL) {
    return;
  }

  ChipReduction smoothed_avg;
  int x_clip = num_cols;
  int y_clip = num_rows;
  if (use_mesh_neighbors == 0) {
    x_clip = THUMBNAIL_SIZE;
    y_clip = THUMBNAIL_SIZE;
  }

  smoothed_avg.Init(num_rows, num_cols, num_frames,
                    SMOOTH_REDUCE_STEP, SMOOTH_REDUCE_STEP,
                    y_clip, x_clip,
                    SMOOTH_REDUCE_STEP * SMOOTH_REDUCE_STEP * .4);
  for (size_t frame_ix = 0; frame_ix < num_frames; frame_ix++) {
    smoothed_avg.ReduceFrame(data + flow_ix * frame_stride + frame_ix * flow_frame_stride, bad_wells, frame_ix);
  }
  smoothed_avg.SmoothBlocks(SMOOTH_REDUCE_REGION, SMOOTH_REDUCE_REGION);

  for (size_t frame_ix = 0; frame_ix < num_frames; frame_ix++) {
    float *y_start = Y.data() + frame_ix * frame_stride;
    float *y_end = y_start + frame_stride;
    int16_t *trace_start = data + flow_ix * frame_stride + frame_ix * flow_frame_stride;
    while(y_start != y_end) {
      *y_start++ = *trace_start++;
    }
  }

  // subtract off flow,frame avg
  for (size_t frame_ix = 0; frame_ix < num_frames; frame_ix++) {
    float *start = Y.data() + frame_stride * frame_ix;
    float *end = start + frame_stride;
    for (size_t row = 0; row < num_rows; row++) {
      for (size_t col = 0; col < num_cols; col++) {
        float avg = smoothed_avg.GetSmoothEst(row, col, frame_ix);
        *start++ -= avg;
      }
    }
  }

  // Get coefficients to solve
  Eigen::MatrixXf B = Y * SX.transpose();
  // Uncompress data into yhat matrix
  Eigen::MatrixXf Yhat = B * Basis.transpose();

  // add the flow/frame averages back
  for (size_t frame_ix = 0; frame_ix < num_frames; frame_ix++) {
    float *start = Y.data() + frame_stride * frame_ix;
    float *end = start + frame_stride;
    float *hstart = Yhat.data() + frame_stride * frame_ix;
    for (size_t row = 0; row < num_rows; row++) {
      for (size_t col = 0; col < num_cols; col++) {
        float avg = smoothed_avg.GetSmoothEst(row, col, frame_ix);
        *start++ += avg;
        *hstart++ += avg;
      }
    }
  }

  for (size_t frame_ix = 0; frame_ix < num_frames; frame_ix++) {
    float *yhat_start = Yhat.data() + frame_ix * frame_stride;
    float *yhat_end = yhat_start + frame_stride;
    int16_t *trace_start = data + flow_ix * frame_stride + frame_ix * flow_frame_stride;
    while(yhat_start != yhat_end) {
      *trace_start++ = (int)(*yhat_start + .5f);
      yhat_start++;
    }
  }

  Y = Y - Yhat;
  Eigen::VectorXf M = Y.rowwise().squaredNorm();

  float *mad_start = mad;
  float *mad_end = mad_start + frame_stride;
  float *m_start = M.data();
  while (mad_start != mad_end) {
    *mad_start += *m_start;
    mad_start++;
    m_start++;
  }

  mad_start = mad;
  mad_end = mad + frame_stride;
  int norm_factor = num_frames;// num_flows * num_frames;
  while (mad_start != mad_end) {
    *mad_start /= norm_factor;
    mad_start++;
  }
}

int TraceStoreCol::PrepareReference(size_t flowIx, std::vector<char> &filteredWells) {
    mRefWells.resize(mUseAsReference.size());
    for (size_t i = 0; i < mRefWells.size(); i++) {
      if (mUseAsReference[i]) {
        mRefWells[i] = 0;
      }
      else {
        mRefWells[i] = 1;
      }
    }
    assert(flowIx < mRefReduction.size());

    int x_clip = mCols;
    int y_clip = mRows;
    if (mUseMeshNeighbors == 0) {
      x_clip = THUMBNAIL_SIZE;
      y_clip = THUMBNAIL_SIZE;
    }
    mRefReduction[flowIx].Init(mRows, mCols, mFrames,
                               REF_REDUCTION_SIZE, REF_REDUCTION_SIZE,
                               y_clip, x_clip, 1);
    for (size_t frame_ix = 0; frame_ix < mFrames; frame_ix++) {
      mRefReduction[flowIx].ReduceFrame(&mData[0] + flowIx * mFrameStride + frame_ix * mFlowFrameStride, &mRefWells[0], frame_ix);
    }
    mRefReduction[flowIx].SmoothBlocks(REF_SMOOTH_SIZE, REF_SMOOTH_SIZE);
    return TSM_OK;
  }

int TraceStoreCol::PrepareReferenceOld (size_t flowIx, std::vector<char> &filteredWells) {
    int fIdx = flowIx;
    CalcReference (mRowRefStep, mColRefStep, flowIx, mRefGrids[fIdx], filteredWells);
    mRefGridsValid[fIdx] = 1;
    mFineRefGrids.resize (mRefGrids.size());
    mFineRefGrids[fIdx].Init (mRows, mCols, mRowRefStep/2, mColRefStep/2);
    int numBin = mFineRefGrids[fIdx].GetNumBin();
    int rowStart = -1, rowEnd = -1, colStart = -1, colEnd = -1;
    for (int binIx = 0; binIx < numBin; binIx++) {
      mFineRefGrids[fIdx].GetBinCoords (binIx, rowStart, rowEnd, colStart, colEnd);
      vector<float> &trace = mFineRefGrids[fIdx].GetItem (binIx);
      CalcMedianReference ( (rowEnd + rowStart) /2, (colEnd + colStart) /2, mRefGrids[fIdx],
                            mDist, mValues, trace);
    }
    return TSM_OK;
  }

int TraceStoreCol::CalcMedianReference (size_t row, size_t col,
                           GridMesh<std::vector<float> > &regionMed,
                           std::vector<double> &dist,
                           std::vector<std::vector<float> *> &values,
                           std::vector<float> &reference) {
    int retVal = TraceStore::TS_OK;
    reference.resize(mFrames);
    std::fill(reference.begin(), reference.end(), 0.0);
    regionMed.GetClosestNeighbors (row, col, mUseMeshNeighbors, dist, values);
    int num_good = 0;
    size_t valSize = values.size();
    for (size_t i =0; i < valSize; i++) {
      if (values[i]->size() > 0) {
        num_good++;
      }
    }
    // try reaching a little farther if no good reference close by
    if (num_good == 0) {
      regionMed.GetClosestNeighbors (row, col, 2*mUseMeshNeighbors, dist, values);
    }
    size_t size = 0;
    double maxDist = 0;
    for (size_t i = 0; i < values.size(); i++) {
        size = max (values[i]->size(), size);
        maxDist = max(dist[i], maxDist);
    }
    reference.resize (size);
    std::fill (reference.begin(), reference.end(), 0.0);
    double distWeight = 0;
         valSize = values.size();

    for (size_t i = 0; i < valSize; i++) {
      if (values[i]->size()  == 0) {
        continue;
      }
      double w = TraceStore::WeightDist (dist[i], mRowRefStep); //1/sqrt(dist[i]+1);
      distWeight += w;
      size_t vSize = values[i]->size();
      for (size_t j = 0; j < vSize; j++) {
        reference[j] += w * values[i]->at (j);
      }
    }
    // Divide by our total weight to get weighted mean
    if (distWeight > 0)  {
      for (size_t i = 0; i < reference.size(); i++) {
        reference[i] /= distWeight;
      }
      retVal = TraceStore::TS_OK;
    }
    else {
      retVal = TraceStore::TS_BAD_DATA;
    }
    return retVal;
  }


// graph.cpp -- See graph.h and graph_impl.h.
//
// This file is distributed under the MIT license, see LICENSE.txt.

#include <boost/graph/one_bit_color_map.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/graph/stoer_wagner_min_cut.hpp>
#include <boost/graph/edmonds_karp_max_flow.hpp>

#include "graph.h"

// Use Boost's Stoer Wagner Min-Cut algorithm to find the global minimum cut on an undirected graph.
template <>
edge_t UndirectedGraph::partition()
{
  // Boost types can be quite a pain, so let's just use the magic of type
  // inference rather than typing out typenames that don't even fit on a line.

  // Boost's implementation of this algorithm stores the cut (i.e. the
  // partitioning) in a parity map, so create one for it to use.  One bit per
  // node, to represent which side of the cut that node is on.
  auto num_vertices = boost::num_vertices(g);
  auto vertices = boost::get(boost::vertex_index, g);
  BOOST_AUTO(parities, boost::make_one_bit_color_map(num_vertices, vertices));

  // Perform the main calculation.
  auto weight_map = boost::get(boost::edge_weight, g);
  auto parity_map = boost::parity_map(parities);
  edge_t w = boost::stoer_wagner_min_cut(g, weight_map, parity_map);

  // Mark the found partition on the graph's vertex properties.
  for (auto v: Types::Vertices(g)) {
    g[v].partition = boost::get(parities, v);
  }

  return w;
}

// Perform a simple DFS and set the partition on all reachable nodes.
// NOTE: Here we could create a filtered_graph of all edges in the residual
// network with remaining capacity, and then do a reachability search on that
// using boost, but this approach is simpler.
static void partitionReachable(FlowNetwork::Types::G &g,
                               FlowNetwork::Types::V &start,
                               int partition)
{
  // Avoid infinite loops on cycles, stop at already processed nodes.
  if (g[start].partition == partition) return;

  // We've reached this vertex, so mark it.
  g[start].partition = partition;

  // Now visit all neighbours that we can reach via residual flow capacity.
  for (auto e : FlowNetwork::Types::OutEdges(start, g)) {
    if (boost::get(boost::edge_residual_capacity, g, e) > 0) {
      FlowNetwork::Types::V target = boost::target(e, g);
      partitionReachable(g, target, partition);
    }
  }
}

// This function finds an s-t min-cut, i.e. the min-cut that also satisfies the
// requirement of having the nodes 's' and 't' on opposite sides of the cut.
//
// This is not implemented directly, but via the max-flow/min-cut duality.
// What this means is that we map this graph to a flow network, solve for
// max-flow, and then map that flow back to a min-cut.
template <>
edge_t UndirectedGraph::partition(node_t source, node_t sink)
{
  Types::V sourceV = node(source);
  Types::V sinkV = node(sink);
  assert(sourceV != Types::Null());
  assert(sinkV != Types::Null());

  using namespace boost;

  FlowNetwork::Types::G fg;

  // The max-flow algorithm needs to know which pairs of edges are the reverse
  // of each other, so we have a map to store that property.
  auto reverse = get(edge_reverse, fg);

  // Now we convert this undirected graph to a flow network.  Firstly, the flow
  // network is directed, so for every edge a--b in the original graph  we
  // create an edge a->b, and one b->a, each with the same weight as a--b.
  // Creating edges as pairs like this also makes it trivial to build the
  // required reverse-edge map (each edge in a pair is the reverse of the
  // other).
  for (auto e : Types::Edges(g)) {
    // This probably shouldn't assume that a vertex_descriptor from an edge in
    // 'g' has a one-to-one map with a 'fg' vertex_descriptor, but it works.

    Types::V s = boost::source(e, g);
    Types::V t = boost::target(e, g);
    edge_t weight = get(edge_weight, g, e);

    FlowNetwork::Types::E e1, e2;
    bool inserted;

    tie(e1, inserted) = add_edge(s, t, weight, fg);
    tie(e2, inserted) = add_edge(t, s, weight, fg);
    reverse[e1] = e2;
    reverse[e2] = e1;
  }

  // Run the max-flow algorithm.
  edge_t flow = edmonds_karp_max_flow(fg, sourceV, sinkV);

  // We now have a max-flow of the flow network.  To map that back to a min-cut
  // of the original graph we take the residual network (i.e. every edge which
  // is not at maximum capacity in the max-flow) and perform a reachability
  // search from the original source node.  Every node that we can reach goes
  // in one partition, every node that we can't goes in another.  If this seems
  // like voodoo search the internet for "max-flow/min-cut duality", there is
  // plenty written on the topic.

  // We set all functions to Xeon-Phi, but then we do a reachability search
  // from the source node and mark all reachable nodes as Xeon again.
  for (auto v: FlowNetwork::Types::Vertices(fg)) {
    fg[v].partition = 1;
  }

  partitionReachable(fg, sourceV, 0);

  // Now we copy the flow graph partitioning back into the original graph.
  for (auto v: FlowNetwork::Types::Vertices(fg)) {
    g[v].partition = fg[v].partition;
  }

  return flow;
}

template <>
void DirectedGraph::draw(const char *file) const
{
  draw_implementation(g, file, "digraph", "-->");
}

template <>
void UndirectedGraph::draw(const char *file) const
{
  draw_implementation(g, file, "graph", "--");
}


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 */

#define BOOST_TEST_MAIN

#include <boost/make_shared.hpp>
#include <boost/lambda/lambda.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/unit_test.hpp>

#include <Eigen/Core>

#include "Tudat/Astrodynamics/ElectroMagnetism/radiationPressureInterface.h"

namespace tudat
{
namespace unit_tests
{

BOOST_AUTO_TEST_SUITE( test_radiation_pressure_interface )

//! Test implementation of radiation pressure calculation
BOOST_AUTO_TEST_CASE( testRadiationPressureCalculation )
{
    // Calculate total solar power (Montenbruck & Gill, 2000, p.77)
    double totalSolarPower = 1367.0 * 4.0 * mathematical_constants::PI *
            physical_constants::ASTRONOMICAL_UNIT * physical_constants::ASTRONOMICAL_UNIT;

    // Calculate radiation pressure at 1 AU.
    double calculatedRadiationPressure = electro_magnetism::calculateRadiationPressure(
                totalSolarPower, physical_constants::ASTRONOMICAL_UNIT );

    // Set literature radiation pressure (Montenbruck & Gill, 2000, p.77)
    double expectedRadiationPressure = 4.56E-6;

    BOOST_CHECK_CLOSE_FRACTION( calculatedRadiationPressure, expectedRadiationPressure, 1.0E-4 );

    // Test calculation of radiation pressure from class interface.
    std::shared_ptr< electro_magnetism::RadiationPressureInterface > radiationPressureInterface =
            std::make_shared< electro_magnetism::RadiationPressureInterface >(
                [ & ]( ){ return totalSolarPower; },
                [ & ]( ){ return
                    ( Eigen::Vector3d( ) <<
                      0.0, physical_constants::ASTRONOMICAL_UNIT / std::sqrt( 2.0 ), 0.0 ).finished( ); },
                [ & ]( ){ return
                    ( Eigen::Vector3d( ) <<
                      physical_constants::ASTRONOMICAL_UNIT / std::sqrt( 2.0 ), 0.0, 0.0 ).finished( ); },
                1.0, 1.0 );
    radiationPressureInterface->updateInterface( );
    double classCalculatedRadiationPressure =
            radiationPressureInterface->getCurrentRadiationPressure( );

    BOOST_CHECK_CLOSE_FRACTION( classCalculatedRadiationPressure, expectedRadiationPressure,
                                1.0E-4 );
    BOOST_CHECK_CLOSE_FRACTION( classCalculatedRadiationPressure, calculatedRadiationPressure,
                                2.0 * std::numeric_limits< double >::epsilon( ) );

}

//! Test application of shadow function.
BOOST_AUTO_TEST_CASE( testShadowFunctionLink )
{
    // Set test geometry (see testShadowFunctionForPartialShadow for details)
    double totalSolarPower = 1367.0 * 4.0 * mathematical_constants::PI *
            physical_constants::ASTRONOMICAL_UNIT * physical_constants::ASTRONOMICAL_UNIT;

    const Eigen::Vector3d occultingBodyPosition = Eigen::Vector3d::Zero( );
    const double occultedBodyRadius = 6.96e8; // Siedelmann 1992.
    const double occultingBodyRadius = 6378.137e3; // WGS-84.

    std::vector< std::function< Eigen::Vector3d( ) > > occultingBodyPositionFunctions;
    occultingBodyPositionFunctions.push_back( [ & ]( ){ return occultingBodyPosition; } );
    std::vector< double > occultingBodyRadii;
    occultingBodyRadii.push_back( occultingBodyRadius );

    Eigen::Vector3d satelliteDirection( 0.018, 1.0, 0.0 );
    satelliteDirection.normalize( );
    const Eigen::Vector3d satellitePosition = ( occultingBodyRadius + 1.0e3 ) * satelliteDirection;
    const Eigen::Vector3d occultedBodyPosition = -149598000.0e3 * Eigen::Vector3d( 1.0, 0.0, 0.0 );

    // Create radiation pressure interface with occultation.
    std::shared_ptr< electro_magnetism::RadiationPressureInterface >
            occultedRadiationPressureInterface =
            std::make_shared< electro_magnetism::RadiationPressureInterface >(
                [ & ]( ){ return totalSolarPower; },
                [ & ]( ){ return occultedBodyPosition; },
                [ & ]( ){ return satellitePosition; },
                1.0, 1.0, occultingBodyPositionFunctions, occultingBodyRadii,
                occultedBodyRadius );
    occultedRadiationPressureInterface->updateInterface( );

    // Create radiation pressure interface without occultation.
    std::shared_ptr< electro_magnetism::RadiationPressureInterface >
            unoccultedRadiationPressureInterface =
            std::make_shared< electro_magnetism::RadiationPressureInterface >(
                [ & ]( ){ return totalSolarPower; },
                [ & ]( ){ return occultedBodyPosition; },
                [ & ]( ){ return satellitePosition; },
                1.0, 1.0 );
    unoccultedRadiationPressureInterface->updateInterface( );

    // Test application of shadow function (see testShadowFunctionForPartialShadow)
    // Satellite partially visible. Satellite is located between the locations of the previous
    // tests in penumbra. According to analytical derivations in Matlab the shadow function should
    // be around 0.4547. Altitude of satellite = 1000 km (from unitTestMissionGeometry.cpp).
    BOOST_CHECK_CLOSE_FRACTION(
                occultedRadiationPressureInterface->getCurrentRadiationPressure( ) /
                unoccultedRadiationPressureInterface->getCurrentRadiationPressure( ) , 0.4547,
                                1.0E-4 );

}

BOOST_AUTO_TEST_SUITE_END( )

} // namespace unit_tests
} // namespace tudat