#!/usr/bin/Rscript

#plot motifs in the format of cisfinder output
#and save motifs as meme format
library(seqLogo)
args <- commandArgs(T)
input_motif <- args[1]
output_dir <- args[2]

SaveMotif <- function(motif_file, output_dir){
    tmp = file(motif_file, 'r')
    
    #save motif as meme format
    out = file(paste(output_dir, "/motif.meme", sep=""), 'w')
    write("MEME version 4\n", out)
    write("ALPHABET= ACGT\n", out)
    write("strands: + -\n", out)
    write("Background letter frequencies\n", out)
    write("A 0.25 C 0.25 G 0.25 T 0.25\n", out)

    line = readLines(tmp, n=1)
    while(length(line) > 0){
        if(startsWith(line, ">")){
            motif_id <- strsplit(line, "\t")[[1]][1]
            #save pfm
            motif_id <- substr(motif_id, 2, nchar(motif_id)) #remove ">"
            write(paste(">", motif_id, sep=""), paste(output_dir, "/", motif_id, ".pfm", sep=""))
            motif <- c()
            while(nchar(line) > 0){
                line = readLines(tmp, n=1)
                freq = as.numeric(strsplit(line, "\t")[[1]][-1]) #remove index
                motif <- rbind(motif, freq)
            }
            motif <- data.frame(motif, row.names=NULL)
            if(is.na(motif[1,1])){next}
            else{
                write(paste("MOTIF", motif_id, sep=" "), out)
                write(paste("letter-probability matrix: alength= 4 w= ", dim(motif)[1], sep=""), out)
                #save pfm
                rownames(motif) <- seq(0,dim(motif)[1]-1)
                write.table(motif, paste(output_dir, "/", motif_id, ".pfm", sep=""), append=T, quote=F, col.names=F, row.names=T, sep="\t")
                MotifLogo(motif, paste(output_dir, "/", motif_id, ".seqlogo.pdf", sep=""), out)
            }
        }
        else{line = readLines(tmp, n=1)}
    }
    close(tmp)
    close(out)
}

MotifLogo <- function(pfm, figure_filename, meme_file){
    #shape of pfm: length*4
    pfm_norm = pfm/rowSums(pfm)
    
    #pdf(figure_filename, width=8, height=6)
    pdf(figure_filename)
    seqLogo(t(pfm_norm))
    dev.off()
    
    #save pwm
    write.table(pfm_norm, sub(pattern = ".seqlogo.pdf", replacement = ".pwm", figure_filename), quote=F, sep="\t", col.names=F, row.names=F)
    write.table(pfm_norm, meme_file, append=T, quote=F, col.names=F, row.names=F)
    write("\n", meme_file)
}

###main function
SaveMotif(input_motif, output_dir)