#!/usr/bin/env python

import re
import os
import sys
import argparse

## basic argument
parser = argparse.ArgumentParser(description="you should add those parameter")
parser.add_argument('-i', '--input', help = 'the software repeatHMM analysis result', required = True)
parser.add_argument('-o', '--output', help = 'the outdir', default = os.path.abspath("./"))
parser.add_argument('-p', '--pafile', help = 'the pa file pathway', required = True)
argv = vars(parser.parse_args())
inf = os.path.abspath(argv['input'].strip())
oud = os.path.abspath(argv['output'].strip())
paf = os.path.abspath(argv['pafile'].strip())

def checkDir(Dir):
    if not os.path.exists(Dir):
        os.makedirs(Dir)
checkDir(oud)


def trimpa(paf):
    padoc={}
    with open(paf,'r') as pal:
        for palic in pal:
            if palic.startswith("#"):
                continue
            Palicon=palic.strip().split(',')
            padoc[str(Palicon[0])]=[str(str(Palicon[1]).replace("chr", "")),str(Palicon[2]),str(Palicon[3]),str(abs(float(Palicon[5]))),str(Palicon[4])]
    return padoc

def trimre(pad,inf):
    ous="Gene\tChr\tStart\tEnd\tRepeat motif\tRepeat counts\tDistribution\n"
    with open(inf,'r') as inl:
        for inlic in inl:
            if "p2sp" in inlic and "INFO:" in inlic and "allocr" in inlic and "[0, 0]" not in inlic:
                inlicon=inlic.strip().split("p2sp [")
                inlicontmp1=inlicon[1].strip().split("\'")
                genamet=inlicontmp1[1]
                repeatnum=str(inlicontmp1[3].strip().split("allocr:")[1]).replace(",", " ")
                if str(genamet) in pad.keys():
                    ous+=str(genamet)+"\t"+str(pad[genamet][0])+"\t"+str(pad[genamet][1])+"\t"+str(pad[genamet][2])+"\t"+str(pad[genamet][3])+"\t"+str(pad[genamet][4])+"\t"+repeatnum+"\n"
    return ous

if __name__ == "__main__":
    pado = trimpa(paf)
    outr = trimre(pado,inf)
    outfile=oud+"/repeatHMM.result.xls"
    outc=open(outfile,'w')
    outc.write(str(outr))