#!/home/304.3-STRUCTPDEV/.local/virtualenv/MyEnv/bin/python
# -*- coding: utf-8 -*-

import sys
import os
import argparse
from Bio.Seq import Seq

# 26 fevrier 2020, Romain BM
# 10 mars 2020, Jeremy L (modif)
# This script is meant to prepare a fasta file for MEME to halve a palindromic matrix previously generated by MEME
# it reads a MEME output file
# extract the sequences that were used to buil a palindromic matrix
# cut the seq into two, reverse complement the one on the 3 prime side and print them out

print("WARNING: scripts works for full sequences")


parser = argparse.ArgumentParser() 

parser.add_argument("--meme","-m")
parser.add_argument("--fasta","-f")
parser.add_argument("--seqlen","-s", type=int, default=0) 
# this depends on the meme alignment and the size of the symetric motif..
args = parser.parse_args()

fin=open("%s" % args.meme, "r")
fou=open("%s" % args.fasta, "w")

MINSEQLEN=int(args.seqlen)
gate="closed"
cpt=0


def micmac(LINE):
  tmp=LINE.strip("\n").split(" ")
  #print len(tmp)
  #print tmp
  while("" in tmp) : 
    tmp.remove("") 
  seq=''.join(tmp[-3:])
  s1=seq[:len(seq)]
  #s2=seq[len(seq):]
  s2rc=Seq(s1).reverse_complement()
  return [s1, s2rc]

for l in fin:
  if l[0:5]=="MOTIF":
    nseq=int(l.split(" ")[9])
    print("%d sequences were used to built the symetric motif\n" % nseq)

  if "MEME-1 sites sorted by position p-value" in l:
    gate="open"
    seen=0
    kept=0

  if gate=="open" and cpt>3 and cpt<nseq+4:
    listSeq=micmac(l)
    if len(listSeq[0])>MINSEQLEN:
      fou.write(">seq%d\n%s\n" % (seen, listSeq[0]))
      kept+=1
    if len(listSeq[1])+1>MINSEQLEN:
      fou.write(">seq%drc\n%s\n" % (seen, listSeq[1]))
      kept+=1
    cpt+=1
    seen+=1
  elif gate=="open" and cpt < 4:
    cpt+=1

fou.close()
fin.close()
print ("%d seqs were seen again\n" % seen)
print ("%d seqs were kept\n" % kept)

