#include "counts.hpp"

#include <algorithm>

#include <boost/tokenizer.hpp>

#include <exception>

#include <fstream>

#include <unordered_map>

#include "aux.hpp"

#include "compression.hpp"

#include "io.hpp"

#include "log.hpp"

#include "parallel.hpp"

using namespace std;

namespace PF = PoissonFactorization;

using Int = PF::Int;

using IMatrix = PF::IMatrix;

using Matrix = PF::Matrix;

using Vector = PF::Vector;

Counts::Counts(const string &path_, const string &separator)

: path(path_),

row_names(),

col_names(),

counts(parse_file<IMatrix>(path, read_counts, separator, row_names,

col_names)) {}

void select_top(vector<Counts> &counts_v, size_t top) {

if (top == 0 or counts_v.empty() or counts_v[0].row_names.size() <= top)

return;

LOG(verbose) << "Selecting top " << top;

Vector gene_sums = rowSums<Vector>(counts_v[0].counts);

for (size_t i = 1; i < counts_v.size(); ++i)

gene_sums += rowSums<Vector>(counts_v[i].counts);

const size_t G = gene_sums.n_elem;

vector<size_t> order(G);

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);

});

order.resize(top);

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.counts.n_cols;

IMatrix m(top, T, arma::fill::zeros);

for (size_t i = 0; i < top; ++i)

m.row(i) = counts.counts.row(order[i]);

counts.counts = m;

counts.row_names = names;

}

}

void discard_empty_spots(Counts &c) {

auto cs = colSums<Vector>(c.counts);

const size_t N = cs.n_elem;

for (size_t n = 0; n < N; ++n)

if (cs(N - n - 1) == 0) {

c.counts.shed_col(N - n - 1);

c.col_names.erase(begin(c.col_names) + N - n - 1);

}

}

vector<Counts> load_data(const vector<string> &paths, bool intersect,

size_t top, bool discard_empty) {

vector<Counts> counts_v;

for (auto &path : paths) {

LOG(verbose) << "Loading " << path;

counts_v.push_back(Counts(path));

}

if (intersect)

gene_intersection(counts_v);

else

gene_union(counts_v);

select_top(counts_v, top);

if (discard_empty)

for (auto &counts : counts_v)

discard_empty_spots(counts);

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;

for (auto &counts : counts_v) {

const size_t H = counts.counts.n_rows;

const size_t S = counts.counts.n_cols;

IMatrix new_counts(G, S, arma::fill::zeros);

for (size_t h = 0; h < H; ++h) {

auto iter = gene_map.find(counts.row_names[h]);

if (iter != end(gene_map))

new_counts.row(iter->second) = counts.counts.row(h);

}

counts.counts = new_counts;

counts.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 = counts.n_cols;

Matrix d(n, n, arma::fill::zeros);

#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 (counts.n_rows == 0)

return Matrix(0, 0);

const size_t n = split_on_x<double>(col_names[0]).size();

Matrix coords(counts.n_cols, n);

for (size_t i = 0; i < counts.n_cols; ++i) {

auto coord = split_on_x<double>(col_names[i]);

for (size_t 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.n_cols == b.n_cols);

Matrix m(a.n_rows, b.n_rows, arma::fill::zeros);

for (size_t i = 0; i < a.n_rows; ++i)

for (size_t j = 0; j < b.n_rows; ++j)

for (size_t k = 0; k < a.n_cols; ++k) {

const double x = a(i, k) - b(j, k);

m(i, j) += x * x;

}

return m;

}

Matrix row_normalize(Matrix m) {

m.each_row(do_normalize<arma::Row<double>>);

return m;

}

size_t sum_rows(const vector<Counts> &c) {

size_t n = 0;

for (auto &x : c)

n += x.counts.n_rows;

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.counts.n_rows);

return x;

}