#!/usr/bin/env python3
import pandas as pd
import argparse
from pathlib import Path, PurePosixPath

"""
outliers_and_positions.py cross outliers and mapping generated by ikiss (pca and bwa) (lfmm and bwa)  are crossed by pandas.
Used for iKISS pipeline
"""

parser = argparse.ArgumentParser(description='crossing outliers and mapping dataframes')
parser.add_argument('--outliers', help="outliers file generated by pcadapt or lfmm")
parser.add_argument('--kmers_position', help="kmer_position file")
parser.add_argument('--output', help="output path")
args = parser.parse_args()

outliers = args.outliers
positions = args.kmers_position
output = args.output
output = Path(output).resolve()
name_output = PurePosixPath(output).stem
dir_output = PurePosixPath(output).parent
output_stats = f"{dir_output}/stats.txt"

# import files
## outliers
selected_df = pd.read_csv(outliers, delimiter='\t', header="infer")
## kmers with position in ref
database_df = pd.read_csv(positions, delimiter='\t', header="infer")

# over each sequence we remove "_0" characters
def remove_str(seq):
    if "_0" in seq:
        return str(seq)[:-2]
    else:
        return seq
selected_df['sequence'] = selected_df['sequence'].apply(remove_str)

#merge dataframes by bedname, sequence and kmer_number
outliers_with_position = selected_df.merge(database_df, on=["sequence",'bedname'])

#df stats
dico_stats = {}
dico_stats['nb_total_kmers'] = int(len(database_df))
dico_stats['nb_kmers_with_position'] = int(len(selected_df))
dico_stats['nb_outliers_with_position'] = int(len(outliers_with_position))
dico_stats['percentage_outliers_with_position'] = round(dico_stats['nb_outliers_with_position']*100/dico_stats['nb_kmers_with_position'], 3)
stats_df = pd.DataFrame.from_dict(dico_stats, orient="index")
stats_df.to_json(output_stats)

#df_kmer_module.columns = ['nb_kmers_mapped_in_ref', 'nb_outliers_lfmm', 'nb_outliers_with_position','percentage_kmers_with_position']
#writing "outliers with position" df to csv
outliers_with_position.to_csv(output, index=False, sep='\t')


                

