import argparse
import pandas as pd
import numpy as np
from itertools import chain
from pathlib import PosixPath as Path

seq_exist = pd.read_csv('/app/MANE_all_transcipts.csv')

def get_round_2(threshold, workdir, round, drug_csv : Path):
    workdir = Path(workdir)
    
    # Get output names
    # Get patient fasta name
    # fasta_files = list(workdir.glob('*_variants_transcript_id_mutations.fasta'))
    # if not fasta_files:
    #     raise FileNotFoundError("No variants transcript mutations fasta file found")
    # fasta_name = fasta_files[0].stem.replace('_variants_transcript_id_mutations', '')
    fasta_name = "patient_0"
    
    # Get test drug
    drug_name = drug_csv.stem
    
    name_out = f"{drug_name}_{fasta_name}"
    
    # Get all work dir complex files
    results_files = list(workdir.glob('*_results.tsv'))
    if not results_files:
        raise FileNotFoundError("No results TSV files found")
    
    round_1_score_list = [str(f) for f in results_files]
    
    # Read 1st round
    transcipts_1 = [
        pd.read_csv(i, sep='\t', header=None).sort_values([2], ascending=False)
        for i in round_1_score_list
        ]
    
    # Get position of drug 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 transcript_complex 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(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(
            workdir / f'{name_out}_significant_interactions.tsv', sep='\t', index=False
        )
        drug_0_score.to_csv(workdir / f'{name_out}_drug_scores.tsv', sep='\t', index=False)
    
    transcipts_1 = pd.DataFrame([
        (Path(j).stem.replace('_results', ''), 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
    ])
    
    # Check if any interaction is above threshold for first round 
    if round == 1:
        if transcipts_1.shape == (0,0):
            transcipts_1.to_csv(workdir / 'round_1.csv', index=False)
            with open(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')
    
    path_out = workdir / f'round_{round}.csv'
    
    if len(transcipts_1) == 0:
        with path_out.open("w") as f:
            f.write("NO TRANSCIPTS ABOVE THRESHOLD")
            return 'STOP NO TRANSCIPTS ABOVE THRESHOLD'
        
    
    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'])]
    temp = [
        seq_exist[seq_exist['transcipt'] == i]
        for i in transcipt_names
    ]
    transcipts_1['protein_name'] =[
        i.iloc[0]['symbol'] if len(i) > 0 else ""
        for i in temp
    ]
    
    # Save data on first round
    transcipts_1.to_csv(path_out, index=False)
    
    # Get fasta for second round
    if round == 1:
        # Get all transcripts 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(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, type=Path, help="workdir")
    parser.add_argument("--round", required=True, type=int, help="Round 1 or 2")
    parser.add_argument("--drug-csv", required=True, type=Path, help="csv file with drug smiles")
    args = parser.parse_args()

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