import argparse
import pandas as pd
import numpy as np
import subprocess
from itertools import chain

seq_exist = pd.read_csv('/home/omic/ConPLex/MANE_all_transcipts.csv')

def get_round_2(threshold, workdir, round):
    #get output names
    #get patient fasta name
    fasta_name = subprocess.check_output([f'find {workdir} -name \*_variants_transcript_id_mutations.fasta' ], shell = True)
    fasta_name = str(fasta_name).split('/')[-1].split('_variants_transcript_id_mutations.fasta')[0]
    #get test drug 
    drug_name = subprocess.check_output([f'find {workdir} -name \*.csv' ], shell = True)
    drug_name = [x for x in str(drug_name).split('/') if "round_1.csv" not in x][-1].split('.csv')[0]
    name_out = drug_name + '_' + fasta_name
    #get all work dir conplex files
    #res_list = subprocess.check_output([f'ls {workdir}/*_results.tsv'], shell = True)
    res_list = subprocess.check_output([f'find {workdir} -name \*_results.tsv' ], shell = True)
    round_1_score_list = str(res_list).split('\'')[1].split('\\n')[:-1]
    #read 1. round
    transcipts_1 = [pd.read_csv(i, sep = '\t', header=None).sort_values([2], ascending=False) for i in round_1_score_list]
    #get poition of drag interaction vs all metabolites
    drug_pos = [[n for n, j in enumerate(list(i[0])) if j[:4] == 'drug'] for i in transcipts_1]
    #drug score
    drug_0_score = pd.concat([i[i[0] =='drug_0'] for i in transcipts_1]).rename({0:'Drug', 1:'Transcript', 2:'Score'}, axis = 1)
    #filter all below threshold
    transcipts_1 = [i[i[2] > threshold] for i in transcipts_1]
    #save transcipt_conplex above threshold to one file
    if round == 2:
        inter_import = pd.concat(transcipts_1, ignore_index=True).rename({0:'drug/metabolite',1:'transcipt',2:'conplex_score'},axis = 1)
        #add drug
        test_smiles = pd.read_csv(f'{workdir}/smiles.smi', sep = '\t', header=None)
        smi_ = [test_smiles[test_smiles[1] == i].iloc[0][0] for i in list(inter_import['drug/metabolite'])]
        inter_import['smile'] = smi_
        inter_import[['drug/metabolite','smile','transcipt','conplex_score']].to_csv(f'{workdir}/{name_out}_significant_interactions.tsv',sep = '\t',index = False)
        drug_0_score.to_csv(f'{workdir}/{name_out}_drug_scores.tsv',sep = '\t',index = False)
    transcipts_1 = pd.DataFrame([(j.split('/')[-1].split('_results')[0], i.shape[0], i[2].mean(), k[0]) for i,j, k in zip(transcipts_1, round_1_score_list, drug_pos) if i.shape[0] != 0])
    #chack if any interaction is is above threshold for first roud 
    if round == 1:
        if transcipts_1.shape == (0,0):
            transcipts_1.to_csv(f'{workdir}/round_1.csv', index = False)
            with open(f'{workdir}/round_1.fasta', 'w') as f:
                f.write("all_data_is_filtered_out\n")
            return 'STOP NO TRANSCIPTS ABOVE THRESHOLD'

    transcipts_1 =  transcipts_1.rename({0:'transcipt_name',1:'number_of_iteracting_compounds', 2:'mean_binding_above_threshold', 3:'drug_position'}, axis='columns')
    transcipts_1['if_drug_above_threshold'] = transcipts_1.iloc[:,1] > transcipts_1.iloc[:,3]
    #if protein is mutated it has _2 in name, removes it
    transcipt_names = [i.split('_')[0] for i in list(transcipts_1['transcipt_name'])]
    transcipts_1['protein_name'] = [[seq_exist[seq_exist['transcipt'] == i].iloc[0]['symbol']][0] for i in transcipt_names]
    #save data on first round
    transcipts_1.to_csv(f'{workdir}/round_{round}.csv', index = False)
    #get fasta for second round
    if round == 1:
        #get all transcipts of proteins above threshold
        name_2_filttered = list(transcipts_1['protein_name'])
        transcipts_2 = [list(seq_exist[seq_exist['symbol'] == i]['transcipt']) for i in name_2_filttered]
        transcipts_2 = list(chain(*transcipts_2))
        transcipts_2 = list(np.unique((transcipts_2)))
        #filter out transcripts already ran through complex
        transcipts_2 = list(np.array(transcipts_2)[[i not in list(transcipts_1['transcipt_name']) for i in transcipts_2]])
        fasta_new = [['>'+i, seq_exist[seq_exist['transcipt'] == i]['seq'].iloc[0]] for i in transcipts_2]
        fasta_new = list(chain(*fasta_new))
        #write fasta to run
        with open(f'{workdir}/round_2.fasta', 'w') as f:
            for line in fasta_new:
                f.write(f"{line}\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process and screen fragments.")
    parser.add_argument("--threshold", required=True, type=float, help="Threshold for 1. round.of conplex scores")
    parser.add_argument("--workdir", required=True, help="workdir")
    parser.add_argument("--round", required=True, type=int, help="Round 1 or 2")
    args = parser.parse_args()

    get_round_2(args.threshold, args.workdir, args.round)