import argparse
from cpt import OrfFinder
from Bio import SeqIO
from Bio import Seq
from CPT_GFFParser import gffParse, gffWrite
from spaninFuncs import *
import re
import os
import sys

"""
##  Note
NOTE : This was made after I made the i-spanin and o-spanin tools, so there might be some methods that are used differently
and overall some 'differently' written code...
"""

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Get putative protein candidates for u-spanins"
    )
    
    parser.add_argument(
        "fasta_file", type=argparse.FileType("r"), help="Fasta file"
    )  # the "input" argument

    parser.add_argument(
        "--strand",
        dest="strand",
        choices=("both", "forward", "reverse"),
        default="both",
        help="select strand",
    )  # Selection of +, -, or both strands

    parser.add_argument(
        "--table", dest="table", default=11, help="NCBI Translation table", type=int
    )  # Uses "default" NCBI codon table. This should always (afaik) be what we want...

    parser.add_argument(
        "-t",
        "--ftype",
        dest="ftype",
        choices=("CDS", "ORF"),
        default="ORF",
        help="Find ORF or CDSs",
    )  # "functional type(?)" --> Finds ORF or CDS, for this we want just the ORF

    parser.add_argument(
        "-e",
        "--ends",
        dest="ends",
        choices=("open", "closed"),
        default="closed",
        help="Open or closed. Closed ensures start/stop codons are present",
    )  # includes the start and stop codon

    parser.add_argument(
        "-m",
        "--mode",
        dest="mode",
        choices=("all", "top", "one"),
        default="all",  # I think we want this to JUST be all...nearly always
        help="Output all ORFs/CDSs from sequence, all ORFs/CDSs with max length, or first with maximum length",
    )

    parser.add_argument(
        "--switch",
        dest="switch",
        default="all",
        help="switch between ALL putative usps, or a range. If not all, insert a range of two integers separated by a colon (:). Eg: 1234:4321",
    )

    parser.add_argument(
        "--usp_on",
        dest="out_usp_nuc",
        type=argparse.FileType("w"),
        default="_out_usp.fna",
        help="Output nucleotide sequences, FASTA",
    )
    parser.add_argument(
        "--usp_op",
        dest="out_usp_prot",
        type=argparse.FileType("w"),
        default="_out_usp.fa",
        help="Output protein sequences, FASTA",
    )
    parser.add_argument(
        "--usp_ob",
        dest="out_usp_bed",
        type=argparse.FileType("w"),
        default="_out_usp.bed",
        help="Output BED file",
    )
    parser.add_argument(
        "--usp_og",
        dest="out_usp_gff3",
        type=argparse.FileType("w"),
        default="_out_usp.gff3",
        help="Output GFF3 file",
    )
    parser.add_argument(
        "--putative_usp",
        dest="putative_usp_fa",
        type=argparse.FileType("w"),
        default="_putative_usp.fa",
        help="Output of putative FASTA file",
    )
    parser.add_argument(
        "--summary_usp_txt",
        dest="summary_usp_txt",
        type=argparse.FileType("w"),
        default="_summary_usp.txt",
        help="Summary statistics on putative o-spanins",
    )
    parser.add_argument(
        "--putative_usp_gff",
        dest="putative_usp_gff",
        type=argparse.FileType("w"),
        default="_putative_usp.gff3",
        help="gff3 output for putative o-spanins",
    )

    parser.add_argument("--min_size", type=int, default=100, help="minimum size of peptide")
    parser.add_argument("--max_size", type=int, default=200, help="maximum size of peptide")
    parser.add_argument("--lipo_min_start", type=int, default=10, help="minimum start site of lipobox")
    parser.add_argument("--lipo_max_start", type=int, default=30, help="maximum end site of lipobox")
    parser.add_argument("--min_lipo_after", type=int, default=60, help="minumum amount of residues after lipobox")
    parser.add_argument("--max_lipo_after", type=int, default=160, help="maximum amount of residues after lipobox")
    parser.add_argument("--tmd_min_start", type=int, default=75, help="minumum start site of TMD")
    parser.add_argument("--tmd_max_start", type=int, default=200, help="maximum end site of TMD")
    parser.add_argument("--tmd_min_size", type=int, default=15, help="minimum size of TMD")
    parser.add_argument("--tmd_max_size", type=int, default=25, help="maximum size of TMD")

    args = parser.parse_args()

    the_args = vars(parser.parse_args())

    ### usp output, naive ORF finding:
    usps = OrfFinder(args.table, args.ftype, args.ends, args.min_size, args.strand)
    usps.locate(
        args.fasta_file,
        args.out_usp_nuc,
        args.out_usp_prot,
        args.out_usp_bed,
        args.out_usp_gff3,
    )

    args.fasta_file.close()
    args.fasta_file = open(args.fasta_file.name, "r")
    args.out_usp_prot.close()
    args.out_usp_prot = open(args.out_usp_prot.name,"r")

    pairs = tuple_fasta(fasta_file=args.out_usp_prot)
    have_lipo = []
    
    for each_pair in pairs:
        if len(each_pair[1]) <= args.max_size:
            try:
                have_lipo += find_lipobox(pair=each_pair,
                                        minimum=args.lipo_min_start,
                                        maximum=args.lipo_max_start,
                                        min_after=args.min_lipo_after,
                                        max_after=args.max_lipo_after,
                                    )
            except (IndexError, TypeError):
                continue
    
    #print(len(have_lipo))
    #print(have_lipo)

    have_tmd_and_lipo = []
    #print(args.tmd_min_start)
    #print(args.tmd_max_start)
    #print(args.tmd_min_size)
    #print(args.tmd_max_size)

    for each_pair in have_lipo:
        try:
            have_tmd_and_lipo += find_tmd(pair=each_pair,
                                minimum=args.tmd_min_start,
                                maximum=args.tmd_max_start,
                                TMDmin=args.tmd_min_size,
                                TMDmax=args.tmd_max_size,  
                                )
        except (IndexError, TypeError):
            continue
    
    #print(len(have_tmd_and_lipo))
    #print(have_tmd_and_lipo)

    if args.switch == "all":
        pass
    else:
        range_of = args.switch
        range_of = re.search(("[\d]+:[\d]+"), range_of).group(0)
        start = int(range_of.split(":")[0])
        end = int(range_of.split(":")[1])
        have_lipo = parse_a_range(pair=have_tmd_and_lipo, start=start, end=end)
    
    total_have_tmd_and_lipo = len(have_tmd_and_lipo)

    ORF = []
    length = []
    candidate_dict = {k:v for k, v in have_tmd_and_lipo}
    with args.putative_usp_fa as f:
        for desc, s in candidate_dict.items():
            f.write(">" + str(desc))
            f.write("\n" + lineWrapper(str(s).replace("*",""))+"\n")
            length.append(len(s))
            ORF.append(desc)
    #### Extra statistics
    args.out_usp_prot.close()
    all_orfs = open(args.out_usp_prot.name, "r")
    all_isps = open(args.putative_usp_fa.name, "r")
    #record = SeqIO.read(all_orfs, "fasta")
    #print(len(record))
    n = 0
    for line in all_orfs:
        if line.startswith(">"):
            n += 1
    all_orfs_counts = n
    
    c = 0
    for line in all_isps:
        if line.startswith(">"):
            c += 1
    all_isps_counts = c

    if ORF:
        bot_size = min(length)
        top_size = max(length)
        avg = (sum(length)) / total_have_tmd_and_lipo
        n = len(length)
        if n == 0:
          raise Exception("no median for empty data")
        if n % 2 == 1:
          med = length[n // 2]
        else:
          i = n // 2
          med = (length[i - 1] + length[i]) / 2
        with args.summary_usp_txt as f:
            f.write("total potential u-spanins: " +str(total_have_tmd_and_lipo) + "\n")
            f.write("average length (AA): " + str(avg) + "\n")
            f.write("median length (AA): " + str(med) + "\n")
            f.write("maximum orf in size (AA): " + str(top_size) + "\n")
            f.write("minimum orf in size (AA): " + str(bot_size) + "\n")
            f.write("ratio of isps found from naive orfs: " + str(c) + "/" +str(n))

        args.putative_usp_fa = open(args.putative_usp_fa.name, "r")
        gff_data = prep_a_gff3(fa=args.putative_usp_fa, spanin_type="usp", org=args.fasta_file)
        write_gff3(data=gff_data, output=args.putative_usp_gff)
    else:
        with args.summary_usp_txt as f:
            f.write("No Candidate USPs found")
            if have_lipo:
                f.write("\nLipoboxes were found here:\n")
                for each_lipo in have_lipo:
                    f.write('>'+str(each_lipo[0]))
                    f.write("\n" + lineWrapper(each_lipo[1].replace("*",""))+"\n")
            else:
                f.write("\nNo Lipobox(es) were found within search restraints")