import pandas as pd

import numpy as np

import glob

import math

import re

import sys

import multiprocessing

def downsampleRow(args):

    row, targetSum = args

    currentCount = row.sum()

    downsampledRow = row.copy()

    while currentCount > targetSum and currentCount != 0:

        possible = downsampledRow[(downsampledRow > 0)]

        desiredTossCount = int(currentCount - targetSum)

        probabilities = [p / currentCount for p in possible]

        for indexToLower in np.random.choice(

                possible.index, max(0, desiredTossCount),

                replace=True, p=probabilities):

            if downsampledRow[indexToLower] > 0:

                downsampledRow[indexToLower] -= 1

        currentCount = downsampledRow.sum()

    return downsampledRow

# downsample_to = sample to this amount of counts per column

# min_feature_abundance = remove all rows which have less than these counts

def downsampleDataFrame(df, downsample_to, min_feature_abundance=50):

    pool = multiprocessing.Pool(8)

    try:

        df = df.loc[:, df.sum() > downsample_to]

        df = df.loc[df.sum(1) > min_feature_abundance, :]

        subset = df.transpose()

        dfDownsampled = subset.copy()

        for idx, drow in enumerate(

            pool.map(

                downsampleRow, [

                (row, downsample_to) for i, row in subset.iterrows()])):

            dfDownsampled.iloc[idx, :] = drow

    except Exception as e:

        print(e)

    pool.close()

    return dfDownsampled.transpose()