#!/usr/bin/env python

import os, os.path
from bio.seq.format.Fasta import RefFa, EsbFa
from bio.seq.base import Fasta, write_fasta_o

# read a table file and convert to dict data
def read_table(infile,sep="\t"):
    data = {}
    from bio.base import DictFile
    mydic = DictFile(infile,delimiter="\t")
    for key,val in mydic.items():
        data[key] = val[0]
    return data

spe2cd = read_table("species_name2code.table")
type2cd = read_table("txclass2code.table")

def code_length(l):
    if l <= 200:
        return "0"
    elif l < 501:
        return "1"
    elif l < 1001:
        return "2"
    elif l < 1501:
        return "3"
    elif l < 2001:
        return "4"
    elif l < 2501:
        return "5"
    elif l < 3001:
        return "6"
    elif l < 3501:
        return "7"
    elif l < 4001:
        return "8"
    else:
        return "9"

def refseq(infile,outfile):
    for seq in RefFa(infile):
        myid = seq.ID.split("_")[1]
        # refseq database code for 2
        seq.name = myid+"_"+spe2cd.get(seq.species,"000")+"2"+code_length(len(seq))+type2cd.get(seq.db,"119")
        write_fasta_o(outfile,seq)

def _ensembl(infile,outfile):
    # /home/yduan/data/Ensembl_noncode_release-92/Tupaia_belangeri.TREESHREW.ncrna.fa
    def get_name(inName):
        tmp_array = os.path.basename(inName).split(".")
        return tmp_array[0].replace("_"," ")
    name = get_name(infile)
    specode = spe2cd.get(name,"000")
    for seq in EsbFa(infile):
        typecd = type2cd.get(seq.tx_type,"119")
        # make sure ensembl class code is right.
        # We should remove 119 and make every type a code.
        #if typecd == "119":
        #    print seq.tx_type
        seq.name = seq.ID+"_"+specode+"0"+code_length(len(seq))+typecd
        write_fasta_o(outfile,seq)

def ensembl(infiles,outfile):
    for f in infiles:
        _ensembl(f,outfile)

def _noncode(infile,outfile,tx2cc):
    def get_name(infile):
        basename = os.path.basename(infile)
        basename = basename.split(".")[0]
        return basename.split("_")[1]
    def get_typecd(name):
        noncc = tx2cc.get(name,"0000")
        return type2cd.get(noncc,"119")
    name = get_name(infile)
    specode = spe2cd.get(name,"000")
    for seq in Fasta(infile):
        #print seq.name
        #print get_typecd(seq.name)
        #print specode
        typecd = get_typecd(seq.name)
        seq.name = seq.name+"_"+specode+"1"+code_length(len(seq))+typecd
        write_fasta_o(outfile,seq)

def noncode(infiles,outfile):
    tx2cc = read_table("noncode_tx_cc")
    for f in infiles:
        _noncode(f,outfile,tx2cc)

def main(argv):

    import argparse

    parser = argparse.ArgumentParser(description="Make standard fasta for methods assesment.")
    parser.add_argument('infiles',nargs='+',help="file to be add information, \"-\" for stdin ")
    parser.add_argument('-o','--outfile',nargs='?',help="output file",default=sys.stdout,type=argparse.FileType('w'))
    parser.add_argument('--source',choices=['refSeq','noncode','ensembl'],help='Choose the source database.')
    args = parser.parse_args(argv[1:])

    if args.infiles[0] == '-':
        infiles=sys.stdin
    else:
        infiles=args.infiles

    if not args.source:
        raise KeyError("--source should not be empty.")
    elif args.source == "refSeq":
        refseq(infiles,args.outfile)
        # refseq has some species "000"
    elif args.source == "noncode":
        noncode(infiles,args.outfile)
    elif args.source == "ensembl":
        ensembl(infiles,args.outfile)
    else:
        raise KeyError("I don't know.")

if __name__ == '__main__':

    import sys
    outfile = sys.stdout
    #refseq(sys.stdin,outfile)
    #ensembl(sys.argv[1:],outfile)
    #noncode(sys.argv[1:])
    main(sys.argv)