#!/usr/bin/env python3
import pandas as pd
import numpy as np
import argparse

"""
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

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

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

print("col outliers")
print(selected_df.columns)
print(database_df.columns)

#merge dataframes by bedname, sequence and kmer_number
#outliers_with_position = selected_df.merge(database_df, on=["bedname", "sequence", "kmer_number"])
outliers_with_position = selected_df.merge(database_df, on=["sequence",'bedname'])
print("--outliers with pos------")
print(outliers_with_position)
#outliers_with_position = selected_df.merge(database_df, on=['bedname', "sequence", "kmer_number"], how="left")

#writing "outliers with position" df to csv
outliers_with_position.to_csv(output, index=False, sep='\t')

#df sizes
print(f"selected_df: {len(selected_df)}\ndatabase_df: {len(database_df)}\noutliers_with_position: {len(outliers_with_position)}\n")

####  specific to rice test

# # count how many kmers are into chrX database_df
# database_df['in_chr'] = np.where(database_df['REF'].str.contains("Chr"), True, False)
# print(database_df.groupby(['in_chr']).count())
#
# #count how many kmers are into chrX using df outliers_with_position
# outliers_with_position['in_chr'] = np.where(outliers_with_position['REF'].str.contains("Chr"), True, False)
# print(outliers_with_position.groupby(['in_chr']).count())

