#include "counts.hpp"
#include <algorithm>
#include <boost/tokenizer.hpp>
#include <exception>
#include <fstream>
#include <numeric>
#include <unordered_map>
#include "aux.hpp"
#include "compression.hpp"
#include "io.hpp"
#include "log.hpp"
#include "parallel.hpp"

using namespace std;
using Int = STD::Int;
using Matrix = STD::Matrix;
using Vector = STD::Vector;

Counts::Counts() : path(""), row_names(0), col_names(0), matrix(nullptr) {}

Counts::Counts(const string &path_, bool transpose, const string &separator)
    : path(path_),
      row_names(),
      col_names(),
      matrix(make_shared<Matrix>(parse_file<Matrix>(
          path, read_floats, separator, row_names, col_names))) {
  if (transpose) {
    std::swap(row_names, col_names);
    matrix->transposeInPlace();
  }
}

size_t Counts::operator()(size_t g, size_t t) const { return (*matrix)(g, t); }
// size_t &Counts::operator()(size_t g, size_t t) { return (*matrix)(g, t); }

size_t Counts::num_genes() const { return matrix->rows(); }

size_t Counts::num_samples() const { return matrix->cols(); }

void select_top_or_bottom(vector<Counts> &counts_v, size_t n, bool select_top) {
  if (n == 0 or counts_v.empty() or counts_v[0].row_names.size() <= n)
    return;
  if (select_top)
    LOG(verbose) << "Selecting top " << n;
  else
    LOG(verbose) << "Selecting bottom " << n;

  Vector gene_sums = rowSums<Vector>(*counts_v[0].matrix);
  for (size_t i = 1; i < counts_v.size(); ++i)
    gene_sums += rowSums<Vector>(*counts_v[i].matrix);

  const size_t G = gene_sums.size();

  vector<size_t> order(G);
  std::iota(begin(order), end(order), 0);
  sort(begin(order), end(order), [&gene_sums](size_t a, size_t b) {
    return gene_sums(a) > gene_sums(b);
  });

  if (not select_top)
    std::reverse(begin(order), end(order));

  order.resize(n);

  vector<string> names;
  for (auto &o : order)
    names.push_back(counts_v[0].row_names[o]);

  for (auto &counts : counts_v) {
    const size_t T = counts.matrix->cols();
    Matrix m = Matrix::Zero(n, T);
    for (size_t i = 0; i < n; ++i)
      m.row(i) = counts.matrix->row(order[i]);
    *counts.matrix = m;
    counts.row_names = names;
  }
}

void discard_spots(Counts &cnt, size_t min_reads) {
  const size_t R = cnt.matrix->rows();
  const size_t C = cnt.matrix->cols();
  auto cs = colSums<Vector>(*cnt.matrix);
  auto num_discarded = 0;
  for (size_t c = 0; c < C; ++c)
    if (cs(C - c - 1) < min_reads) {
      num_discarded++;
      cnt.col_names.erase(begin(cnt.col_names) + C - c - 1);
    }

  auto old = *cnt.matrix;
  *cnt.matrix = Matrix(R, C - num_discarded);
  size_t c_ = 0;
  for (size_t c = 0; c < C; ++c)
    if (cs(c) >= min_reads)
      cnt.matrix->col(c_++) = old.col(c);
  LOG(verbose) << "Discarded " << num_discarded << " spots with fewer than "
               << min_reads << " counts.";
}

void discard_genes(vector<Counts> &cnts, size_t min_reads) {
  if (cnts.empty())
    return;
  const size_t R = cnts[0].matrix->rows();
  Vector rs = Vector::Zero(R);
  for (auto &cnt : cnts)
    rs += rowSums<Vector>(*cnt.matrix);
  auto num_discarded = 0;
  for (size_t r = 0; r < R; ++r)
    if (rs(R - r - 1) < min_reads) {
      num_discarded++;
      for (auto &cnt : cnts)
        cnt.row_names.erase(begin(cnt.row_names) + R - r - 1);
    }

  for (auto &cnt : cnts) {
    const size_t C = cnt.matrix->cols();
    auto old = *cnt.matrix;
    *cnt.matrix = Matrix(R - num_discarded, C);
    size_t r_ = 0;
    for (size_t r = 0; r < R; ++r)
      if (rs(r) >= min_reads)
        cnt.matrix->row(r_++) = old.row(r);
  }
  LOG(verbose) << "Discarded " << num_discarded << " genes with fewer than "
               << min_reads << " counts.";
}

vector<Counts> load_data(const vector<string> &paths, bool intersect,
                         size_t top, size_t bottom, size_t min_reads_spot,
                         size_t min_reads_gene, bool transpose) {
  vector<Counts> counts_v(paths.size());
#pragma omp parallel for if (DO_PARALLEL)
  for (size_t i = 0; i < paths.size(); ++i) {
    LOG(verbose) << "Loading " << paths[i];
    counts_v[i] = Counts(paths[i], transpose);
  }

  if (intersect)
    gene_intersection(counts_v);
  else
    gene_union(counts_v);

  select_top_or_bottom(counts_v, top, true);
  select_top_or_bottom(counts_v, bottom, false);

  for (auto &counts : counts_v)
    discard_spots(counts, min_reads_spot);
  discard_genes(counts_v, min_reads_gene);

  LOG(verbose) << "Done loading";
  return counts_v;
}

template <typename Fnc>
void match_genes(vector<Counts> &counts_v, Fnc fnc) {
  LOG(verbose) << "Matching genes";
  unordered_map<string, size_t> present;
  for (auto &counts : counts_v)
    for (auto &name : counts.row_names)
      present[name]++;

  vector<string> selected;
  for (auto &entry : present)
    if (fnc(entry.second))
      selected.push_back(entry.first);

  const size_t G = selected.size();

  sort(begin(selected), end(selected));

  unordered_map<string, size_t> gene_map;
  for (size_t g = 0; g < G; ++g)
    gene_map[selected[g]] = g;

#pragma omp parallel for if (DO_PARALLEL)
  for (size_t i = 0; i < counts_v.size(); ++i) {
    const size_t H = counts_v[i].matrix->rows();
    const size_t S = counts_v[i].matrix->cols();
    Matrix new_counts = Matrix::Zero(G, S);
    for (size_t h = 0; h < H; ++h) {
      auto iter = gene_map.find(counts_v[i].row_names[h]);
      if (iter != end(gene_map))
        new_counts.row(iter->second) = counts_v[i].matrix->row(h);
    }
    *counts_v[i].matrix = new_counts;
    counts_v[i].row_names = selected;
  }
}

void gene_union(vector<Counts> &counts_v) {
  match_genes(counts_v, [](size_t x) { return x > 0; });
}

void gene_intersection(vector<Counts> &counts_v) {
  const size_t n = counts_v.size();
  match_genes(counts_v, [n](size_t x) { return x == n; });
}

template <typename T>
vector<T> split_on_x(const string &s, const string &token = "x") {
  vector<T> v;
  size_t last_pos = 0;
  while (true) {
    auto pos = s.find(token, last_pos);
    if (pos != string::npos) {
      v.push_back(atof(s.substr(last_pos, pos - last_pos).c_str()));
    } else {
      v.push_back(atof(s.substr(last_pos).c_str()));
      break;
    }
    last_pos = pos + 1;
  }
  return v;
}

double sq_distance(const string &a, const string &b) {
  auto x = split_on_x<double>(a);
  auto y = split_on_x<double>(b);
  const size_t n = x.size();
  double d = 0;
  for (size_t i = 0; i < n; ++i) {
    double z = x[i] - y[i];
    d += z * z;
  }
  return d;
}

Matrix Counts::compute_distances() const {
  size_t n = matrix->cols();
  Matrix d = Matrix::Zero(n, n);
#pragma omp parallel for if (DO_PARALLEL)
  for (size_t i = 0; i < n; ++i)
    for (size_t j = i + 1; j < n; ++j)
      d(i, j) = d(j, i) = sq_distance(col_names[i], col_names[j]);
  return d;
}

Matrix Counts::parse_coords() const {
  if (matrix->cols() == 0)
    return Matrix(0, 0);
  const STD::Index n = split_on_x<double>(col_names[0]).size();
  Matrix coords(matrix->cols(), n);
  for (STD::Index i = 0; i < matrix->cols(); ++i) {
    auto coord = split_on_x<double>(col_names[i]);
    for (STD::Index j = 0; j < n; ++j)
      coords(i, j) = coord[j];
  }
  return coords;
}

template <typename T>
void do_normalize(T &v) {
  double z = 0;
  for (auto x : v)
    z += x;
  if (z > 0)
    for (auto &x : v)
      x /= z;
}

Matrix compute_sq_distances(const Matrix &a, const Matrix &b) {
  assert(a.cols() == b.cols());
  Matrix m = Matrix::Zero(a.rows(), b.rows());
  for (STD::Index i = 0; i < a.rows(); ++i)
    for (STD::Index j = 0; j < b.rows(); ++j)
      for (STD::Index k = 0; k < a.cols(); ++k) {
        const double x = a(i, k) - b(j, k);
        m(i, j) += x * x;
      }
  return m;
}

Matrix row_normalize(Matrix m) {
  for (STD::Index r = 0; r < m.rows(); ++r) {
    double z = m.row(r).sum();
    if (z > 0)
      m.row(r) /= z;
  }
  return m;
}

size_t sum_rows(const vector<Counts> &c) {
  size_t n = 0;
  for (auto &x : c)
    n += x.matrix->rows();
  return n;
}

size_t sum_cols(const vector<Counts> &c) {
  size_t n = 0;
  for (auto &x : c)
    n += x.matrix->cols();
  return n;
}

size_t max_row_number(const vector<Counts> &c) {
  size_t x = 0;
  for (auto &m : c)
    x = max<size_t>(x, m.matrix->rows());
  return x;
}