#!/usr/bin/env python3

import fire
import numpy as np
import pandas as pd
from pathlib import PurePath


def alt_allele_num(reads_inf):
    reads_inf_str_list = str(reads_inf).split(',')
    reads_inf_int_list = [int(read_i) for read_i in reads_inf_str_list]
    if len(reads_inf_int_list) == 1:
        return 0
    else:
        return reads_inf_int_list[1]


def allele_depth(reads_inf, allele_num=2):
    reads_inf_str_list = str(reads_inf).split(',')
    if len(reads_inf_str_list) <= allele_num:
        reads_inf_int_list = [int(read_i) for read_i in reads_inf_str_list]
        return sum(reads_inf_int_list)
    else:
        return np.nan


def table2hd5(table_file, sep='\t'):
    table_df = pd.read_csv(table_file, sep=sep)
    sample_name = table_df.columns[-1]

    out_table_df = table_df.copy()
    out_table_df.loc[:, 'alt_count'] = table_df.loc[:, sample_name].map(
        alt_allele_num)
    out_table_df.loc[:, 'dep_count'] = table_df.loc[:, sample_name].map(
        allele_depth)
    out_table_df.loc[:, 'sample_id'] = sample_name
    out_table_df = out_table_df.set_index(['Chr', 'Pos', 'Alt',
                                           'sample_id']).unstack('sample_id')

    table_file = PurePath(table_file)
    table_file_name = table_file.stem
    table_file_h5 = table_file.with_suffix('.h5')
    out_table_df.to_hdf(table_file_h5, table_file_name)


if __name__ == '__main__':
    fire.Fire(table2hd5)
