import pandas as pd
import sys

def func(x):
    return ';'.join(x.astype(str))

print('Read data...')
path = sys.argv[1]
if path.split('.')[1] == 'pkl':
    df = pd.read_pickle(path)
else:
    df = pd.read_csv(path, sep='\t', index_col=0)

print('Group by Gene ID and Mirbase ID, get CSs and AIRs strings...')
ls = df.groupby(['Gene ID', 'Mirbase ID']).agg({'AIR': func, 'context++ score': func})

print('Setting Gene ID and mirbase ID back...')
ls['Gene ID'] = ls.index.get_level_values(0)
ls['Mirbase ID'] = ls.index.get_level_values(1)

print(f'Saving to {sys.argv[2]} file...')
ls.to_csv(sys.argv[2], sep='\t', index=None, header=None)

