'''this module is used to append up/down-stream sequence to the target sequence
version1: retrive extending sequences based on BED file generated by JCVI
JCVI:
python -m jcvi.formats.gff bed --type=gene --key=gene GeSeqJob.gff3 -o GeSeq.bed
Note: 
1) JCVI will transform 1-based GFF3 to 0-based BED, which is also suitable for TBtools
2) due to trans-splicing, it might report the unusual range like:
"12588   307844" which indicates that the total gene length is 295256.
'''

from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq


def ParsePlastid(filename):
    '''this file is used to parse the plastid genome,
    Note: only consider single-loop plastid'''
    seq_id = ''
    sequence = ''
    with open(filename, 'rt') as input:
        for record in SeqIO.parse(input, 'fasta'):
            seq_id = record.id
            sequence = str(record.seq)
    return seq_id, sequence


def RetrieveAbnormal(filename, sequence, N, outputname):
    '''Read BED file of abnormal gene and retrieve the expanded range of sequence,
    "expanded" means that up- and down-stream with N bp
    Note: Because this module apply "SeqIO.write" to conduct the saving process,
    it might not be smooth to paste the target sequence to tools like sublime text to further down-stream analysis
    '''
    abnormals = {}
    with open(filename, 'rt') as input:
        for line in input: # "fragment", "copy-abnormal" filter were not set
            ab = line.strip('\n').split('\t')
            start_pos = int(ab[1]) - int(N)
            end_pos = int(ab[2]) + int(N)
            
            if ab[3] in ['nad1', 'nad2', 'nad5'] or "fragment" in ab[3] or "copy_abnormal" in ab[3]: # ignore
                continue

            if ab[-1] == '+':    
                # if seq is on the "+" strand, its down-stream will be appended
                # "-" strand, its up-stream will be appended
                abnormals[ f'{ab[3]}\t{ab[1]}:{ab[1]}\t{ab[2]}:{end_pos}\t{ab[-1]}' ] = sequence[int(ab[1]):end_pos]  # {header}:sequence
            else:
                abnormals[ f'{ab[3]}\t{ab[1]}:{start_pos}\t{ab[2]}:{ab[2]}\t{ab[-1]}' ] = Seq(sequence[start_pos:int(ab[2])]).reverse_complement()
    Records = []
    for record in list(abnormals.items()):
        tmpRecord = SeqRecord(Seq(record[1]), id=record[0], description='')
        Records.append(tmpRecord)
    SeqIO.write(Records, outputname, 'fasta')


def main():
    import optparse

    usage = """python -m plastidUtilis.SequenceAppend -f <input_fasta> -i <input_BED> -o <output_filename>"""
    parser = optparse.OptionParser(usage)
    parser.add_option("-f", dest="fasta", help="input plastid/mitochondia genome fasta",
                    metavar="FILE", action="store", type="string")
    parser.add_option("-i", dest="bed", help="input BED of abnormal gene/fragment",
                    metavar="FILE", action="store", type="string")
    parser.add_option("-n", "--number", help="the number of flanking bp",
                    metavar="NUMBER", action="store", type="int")
    parser.add_option("-o", dest="output", help="output fasta name",
                    metavar="OUT", action="store", type="string")
    (options, args) = parser.parse_args()

    _, sequence = ParsePlastid(options.fasta)
    RetrieveAbnormal(options.bed, sequence, options.number, options.output)


if __name__ == "__main__":
    main()
