import argparse
import os
import pandas as pd
from tqdm import tqdm
from utils import load_coords


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='make single mutant csv')
    parser.add_argument("--fasta_dir", type=str, default=None, required=True)
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--pdb_file", type=str, default=None)
    parser.add_argument("--out_file", type=str, default=None)
    parser.add_argument("--start", type=int, default=-1)
    parser.add_argument("--end", type=int, default=int(1e6))
    args = parser.parse_args()

    one_letter = {
            'VAL':'V', 'ILE':'I', 'LEU':'L', 'GLU':'E', 'GLN':'Q',
            'ASP':'D', 'ASN':'N', 'HIS':'H', 'TRP':'W', 'PHE':'F', 'TYR':'Y',
            'ARG':'R', 'LYS':'K', 'SER':'S', 'THR':'T', 'MET':'M', 'ALA':'A',
            'GLY':'G', 'PRO':'P', 'CYS':'C'
            }
    AA = list(one_letter.values())

    if args.fasta_dir is not None:
        proteins = os.listdir(args.fasta_dir)
        for p in tqdm(proteins):
            protein_name = p.split(".")[0]
            seq = open(f"{args.fasta_dir}/{p}").readlines()[1].strip()
            data = {"mutant":[], "DMS_score":[]}
            for idx, s in tqdm(enumerate(seq)):
                if idx + 1 < args.start or idx + 1 > args.end:
                    continue
                for a in AA:
                    if a == s:
                        continue
                    data["mutant"].append(f"{s}{idx+1}{a}")
                    data["DMS_score"].append(0)
            pd.DataFrame(data).to_csv(f"{args.output_dir}/{protein_name}.csv", index=False)
    
    if args.pdb_file is not None:
        output_dir = os.path.dirname(args.out_file)
        os.makedirs(output_dir, exist_ok=True)
        _, seq = load_coords(args.pdb_file, "A")
        data = {"mutant":[], "DMS_score":[]}
        for idx, s in tqdm(enumerate(seq)):
            if idx + 1 < args.start or idx + 1 > args.end:
                continue
            for a in AA:
                if a == s:
                    continue
                data["mutant"].append(f"{s}{idx+1}{a}")
                data["DMS_score"].append(0)
        pd.DataFrame(data).to_csv(args.out_file, index=False)