#!/usr/bin/python
# -*- coding:utf-8 -*-
'''Frameshift mutation, insertion and deletion.


by Wubin Qu <quwubin@gmail.com>,
Copyright @ 2010, All Rights Reserved.
'''

Author  = 'Wubin Qu  <quwubin@gmail.com> CZlab, BIRM, China'
Date    = 'Nov-05-2010 15:57:41'
License = 'GPL v3'
Version = '1.0'

import sys
import os
from optparse import OptionParser
from BifScripts import FastaFormatParser as FFP


def get_opt():
    '''Handle options'''
    usage = '''Usage: %prog [options]

    ./frameshift_mutation.py -i infile.fa -p 263 -b A -t insertion
    
    '''

    version = '%prog Version: ' + '%s [%s]' % (Version, Date)
    parser = OptionParser(usage=usage, version=version)
    parser.add_option('-i', '--infile', dest='infile', help='Input file name. [String]')
    parser.add_option('-o', '--outfile', dest='outfile', help='Output file name. [String]')
    parser.add_option('-t', '--type', dest='type', help='insertion or deletion?. [String]')
    parser.add_option('-p', '--position', dest='position', help='0-index position of the point mutation position for insertion (after the position). [Integer]')
    parser.add_option('-b', '--base', dest='base', help='Nucleotide base for insertion. [String]')
    [options, args] = parser.parse_args()

    if len(args) > 1:
        parser.error('Incorrect argument, add" "-h" for help.')

    if not options.infile:
        parser.error('Input fasta-format file needed, add" "-h" for help.')

    type_list = ['insertion', 'deletion']
    if not options.type or options.type not in type_list:
        parser.error('Frameshift type needed, insertion of deletion?, add" "-h" for help.')

    if not options.position or int(options.position) < 0:
        parser.error('Point mutation position needed, add" "-h" for help.')

    if not options.base:
        parser.error('Inserted nucleotide base is needed, add" "-h" for help.')

    return options

def main ():
    '''Main'''
    options = get_opt()
    records = FFP.parse(open(options.infile))
    if len(records) > 1:
        print 'Currently, only one sequence support, %s has multiple sequences.' % options.infile
        exit()

    seq = records[0]['seq']
    id = records[0]['id']
    desc = records[0]['desc']
    if len(seq) < len(options.position) + 1:
        print 'Error: Mutation position is longer than the sequence'
        exit()
    
    if options.type == 'insertion':
        position = int(options.position) + 1
        mutated_seq = seq[:position] + options.base + seq[position : ]
    else:
        position = int(options.position)
        if seq[position] != options.base:
	    print 'Error: nucleotide base in %s is not %s, actually is: %s' % (position, options.base, seq[position])
	    exit()
    
        mutated_seq = seq[:position] + seq[position + 1 : ]

    output = '>%s %s%s%s' % (id, desc, os.linesep, mutated_seq)
    try:
        fh = open(options.outfile, 'w')
        fh.write(output)
        fh.close()
    except:
        print output

if __name__ == '__main__':
    main()