import sys

import pandas as pd

from imblearn.under_sampling import RandomUnderSampler

def resample_data(t):

    t = t[['text', 'readm_30d']]

    label = t.pop('readm_30d')

    rus = RandomUnderSampler(random_state=42)

    X, y = rus.fit_resample(t, label.astype('category'))

    df = pd.DataFrame(X[:, 0])

    df.columns = ['text']

    df['readm_30d'] = pd.Series(y)

    return df

if __name__ == '__main__':

    input_file = args = sys.argv[1]

    output_file = args = sys.argv[2]

    df = pd.read_csv(input_file)

    df = resample_data(df)

    print("new shape: {}".format(df.shape))

    df.to_csv(output_file)