#!/usr/bin/env python

import argparse
from random import randint
import pandas as pd
import numpy as np

ETU_CSV = "etudiants.csv"
SUBJ_CSV = "sujets.tsv"
#RES_CSV = "affectations.tsv"


def main(seed):

    etu_df = pd.read_csv(ETU_CSV)
    res_df = pd.DataFrame()

    for gr in etu_df['groupe'].unique():
        gr_df = etu_df[etu_df['groupe'] == gr]
        subj_df = pd.read_csv(SUBJ_CSV, sep='\t')

        while not gr_df.empty:
            subj = subj_df.sample()
            subj_df = subj_df.drop(index=subj.index)

            n_students = min(gr_df.shape[0], 2)
            etu = gr_df.sample(n=n_students)
            gr_df = gr_df.drop(index=etu.index)
            
            res_df = pd.concat([
                    res_df,
                    etu.merge(subj, how='cross')], 
                ignore_index=True)
        
        # Ajout des sujets non affectes
        subj_df['groupe'] = gr
        res_df = pd.concat([
                    res_df,
                    subj_df], 
                ignore_index=True)

    res_fn = f"affectations_{seed}.tsv"
    res_df.to_csv(res_fn, sep='\t', index=False)
    print("Successfully wrote", res_fn)

def parse_args():
    parser = argparse.ArgumentParser(
        prog='Selection de binomes et de sujets SRA2',
        description='Mange etudiants.csv et sujets.tsv, génère affectations_<seed>.tsv. '+\
        'Vous pouvez lui donner une seed avec l\'option `-s | --seed`.')
    parser.add_argument('-s', '--seed')

    args = parser.parse_args()

    seed = args.seed
    if seed is None:
        seed = randint(0, 2**32)

    np.random.seed(int(seed))
    print("Using following random seed:", seed)
    return seed

if __name__ == '__main__':
    seed = parse_args()
    main(seed)