#!/usr/bin/env python
import sys, re
import argparse
import pandas as pd

def read_representatives_file(fi):
    df = pd.read_csv(fi)
    df = df[["genome","cluster"]]
    df = df.set_index('cluster')
    return df

def read_clusters_file(fi, df_repr, output_fi):
    df_clust = pd.read_csv(fi)
    df_clust = df_clust[["genome","secondary_cluster"]]
    composition = df_clust.groupby(['secondary_cluster'])['genome'].apply(list)
    df_clust = composition.to_frame()['genome'].apply(lambda x: ",".join(x) ).to_frame()
    df = pd.merge(df_clust, df_repr,  left_index=True, right_index=True)
    df = df.rename(columns={"genome_y": "representative", "genome_x": "composition"})
    df = df[["representative","composition"]]

    df.to_csv(output_fi ,sep='\t', index=False)

def parse_arguments():
    # Manage parameters.
    parser = argparse.ArgumentParser()

    parser.add_argument('-c', '--clusters', required = True, help = \
    'data_tables/Cdb.csv dRep output file of bins cluster compositions.')

    parser.add_argument('-r', '--representatives', required = True, help = \
    'data_tables/Wdb.csv dRep output file of representative bins clusters.')

    parser.add_argument('-o', '--output_file', required = True, \
    help = 'representative-cluster_to_bins output file.')

    args = parser.parse_args()
    return args

###################################################################################################################################################
#
# MAIN
#
##################################################################################################################################################


def main():

    args = parse_arguments()

    df_repr = read_representatives_file(args.representatives)

    read_clusters_file(args.clusters, df_repr, args.output_file)

if __name__ == '__main__':
    main()