#!/usr/bin/python3
# A program for finding dna fasta files that correspond to the protein fasta files containing single-copy orthologs
# USAGE: ./proteinOrthos2dnaOrthos_TA.py --input PATH2OrthoFiles --cdsFile PATH2combinedCDSfile
# Author: Taruna Aggarwal
# Affiliation: University of New Hampshire, Durham, NH, USA
# Date: 01/20/2016
# Purpose is parse through a protein orthogroup fasta file generated by Orthofinder; identify parts of the headers
# that contain original protein file names + some identifier for the protein file unique to the header; look for the
# identifier in the CDS file header; write the header + file name + species name and its corresponding sequence to a
# new file

import os
import argparse

parser = argparse.ArgumentParser(description="This script extracts coding DNA sequences for the orthogroups generated by Orthofinder."
                                             "The script assumes one-to-one relation between protein and cds files.")
parser.add_argument('--input', help="PATH to parsed orthogroup files", required=True)
parser.add_argument('--cdsFile', help="PATH to a file containing concatenated coding DNA sequences", required=True)
args = parser.parse_args()

# create an empty dictionary that will keep the file names as key and the identifiers as values

cds_dict = {}

for currentFile in os.listdir(args.input):
    cds_names = []
    if currentFile.endswith(".fa"):
        working_file = open(args.input + "/" + currentFile, "r")
        for currentLine in working_file:
            currentLine = currentLine.rstrip()
            if currentLine.startswith(">"):
                try:
                    default_name = currentLine.split('renamed_')[1]
                except IndexError:
                    pass
                cds_names.append(default_name)
        if currentFile not in cds_dict:
            key = currentFile
            cds_dict[key] = list(set(cds_names))

combined_cds_file = open(args.cdsFile, "r")

for filename, value in cds_dict.items():
    new_file = open("CDS_"+filename, "w")
    flag = False
    for line in combined_cds_file:
        line = line.rstrip()
        if line.startswith(">"):
            for item in value:
                # print(index)
                if line.endswith(item):
                    # print(line)
                    new_file.write("{0}_{1}\n".format(line, filename[:-3]))
                    flag = True
                    break
                else:
                    flag = False
        elif flag:
            new_file.write("{0}\n".format(line))
    combined_cds_file.seek(0)
    # print("{0}\n{1}\n}".format(filename, value)

working_file.close()
combined_cds_file.close()
new_file.close()