"""
===============================================================================
Use pysam model to calculate the AS from bam file
===============================================================================
Author : Shujia Huang
Date   : 2014-03-25 14:29:08
"""
import sys
import re
import optparse
import os
import string
import pysam 
import matplotlib.pyplot as plt
import numpy as np
import AlterAlign as ATA

def IsSNP(refbase, alleles):

    isSnp = True
    for ale in alleles: 
        if len(ale) > 1 or len(refbase) > 1: isSnp = False

    return isSnp

def main(opt): 

    vcfInfile = opt.vcfInfile
    bamInfile = opt.bamInfile
    faInfile  = opt.refInfile
    sampleID  = opt.sample
    refId     = opt.refChrId
    outPrefix = opt.outPrefix
    mapq      = opt.mapq
    newID     = opt.exc
    fa = ATA.LoadFaSeq(faInfile, refId)

    print >> sys.stderr, '# [INFO] Fa Loading finish ***'
    if bamInfile[-4:] == '.bam':
        samInHandle = pysam.Samfile(bamInfile, 'rb')
    else :
        samInHandle = pysam.Samfile(bamInfile, 'r')

    #samOutHandle = pysam.Samfile(outPrefix + '.bam', 'wb', template=samInHandle)
    vcfOutHandle = open(outPrefix + '.vcf', 'w')

    print >> sys.stderr, '# [INFO] Now Scaning the VCF and doing alternate align ... ...'
    if vcfInfile[-3:] == '.gz':
        if refId == "ALL":
            I = os.popen('gzip -dc %s' % vcfInfile)
        else:
            I = os.popen('/home/siyang/Bin/software_pip/tabix-0.2.6/tabix -h %s %s' % (vcfInfile, refId))
    else :
        I = open(vcfInfile)

    frist = True
    while 1:
        lines = I.readlines(100000)
        if not lines : break;

        for line in lines :
        
            line = line.strip('\n')
            col  = line.split()
            if re.search(r'^##', line): 
                if frist and re.search(r'^##FORMAT=', line): 
                    frist = False
                    vcfOutHandle.write('##FORMAT=<ID=AA,Number=4,Type=Integer,Description="Information of Alternate Align. Format: Ref_perfect,Alt_Perfect,Both_Perfect,Both_Imperfect">\n')
                vcfOutHandle.write('%s\n' % line)
                continue
            elif re.search(r'^#CHROM', line):
                if len(col) < 10 :
                    print >> sys.stderr, '# [ERROR] The input vcf file (%s) does not contain the "Sample" fields!\n' % vcfInfile
                    sys.exit(1)
                sam2col = {sam:i+9 for i, sam in enumerate(col[9:])}
                if sampleID not in sam2col: 
                    print >> sys.stderr, '# [ERROR] The input sample id (%s) is not match in Vcf file(%s)\n' % (sampleID, vcfInfile)

                if len(newID) > 0: 
                    vcfOutHandle.write('%s\t%s\n' % ('\t'.join(col[:9]), newID))
                else: 
                    vcfOutHandle.write('%s\t%s\n' % ('\t'.join(col[:9]), sampleID))
                continue

            if refId != 'ALL' and refId != col[0]: continue
            if col[4] == '.': continue # ignore REFCALL or INTERGAP

            idx = sam2col[sampleID]
            fi  = col[idx].split(':')

            gt  = fi[0].split('/')
            if '|' in fi[0]: gt = fi[0].split('|')
            gtIdx = 1 # Default is the first ALT Sequence
            if len(gt) == 2 and gt[1] != '.': 
                gtIdx = string.atoi(gt[1])

            col[4] = col[4].split(',')[gtIdx-1] # Match to the identity sample
            isAltAlign = False
            zr,za,zc,zi = 0,0,0,0
            if not IsSNP(col[3], [col[4]]):
                # Not SNP, INTERGAP ...
                isAltAlign = True
                zr,za,zc,zi = ATA.Align(samInHandle,
                                        fa,
                                        col[0], 
                                        string.atoi(col[1]), 
                                        col[2], 
                                        col[3], 
                                        col[4][1:], #col[4][0] is reference
                                        mapq)
            # Ignore the position which is meanless
            if not isAltAlign and col[idx] == './.': continue

            fm = {t:i for i, t in enumerate(col[8].split(':'))} # Get Format
            if col[idx] != './.' and len(fi) != len(fm): 
                raise ValueError('[ERROR] The format of "FORMAT"' + 
                                 'fields is not match sample '    + 