# Copyright 2015 by Gert Hulselmans.  All rights reserved.
# This file is part of the Biopython distribution and governed by your
# choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
# Please see the LICENSE file that should have been included as part of this
# package.

"""Parse Cluster Buster position frequency matrix files."""

from Bio import motifs


class Record(list):
    """Class to store the information in a Cluster Buster matrix table.

    The record inherits from a list containing the individual motifs.
    """

    def __str__(self):
        return "\n".join(str(motif) for motif in self)


def read(handle):
    """Read motifs in Cluster Buster position frequency matrix format from a file handle.

    Cluster Buster motif format: http://zlab.bu.edu/cluster-buster/help/cis-format.html
    """
    motif_nbr = 0
    record = Record()
    nucleotide_counts = {'A': [], 'C': [], 'G': [], 'T': []}
    motif_name = ""

    for line in handle:
        line = line.strip()
        if line:
            if line.startswith('>'):

                if motif_nbr != 0:
                    motif = motifs.Motif(alphabet="GATC", counts=nucleotide_counts)
                    motif.name = motif_name
                    record.append(motif)

                motif_name = line[1:].strip()
                nucleotide_counts = {'A': [], 'C': [], 'G': [], 'T': []}
                motif_nbr += 1
            else:
                if line.startswith('#'):
                    continue

                matrix_columns = line.split()

                if len(matrix_columns) == 4:
                    [nucleotide_counts[nucleotide].append(float(nucleotide_count))
                     for nucleotide, nucleotide_count in zip(['A', 'C', 'G', 'T'], matrix_columns)]

    motif = motifs.Motif(alphabet='GATC', counts=nucleotide_counts)
    motif.name = motif_name
    record.append(motif)

    return record


def write(motifs):
    """Return the representation of motifs in Cluster Buster position frequency matrix format."""
    lines = []
    for m in motifs:
        line = '>{0}\n'.format(m.name)
        lines.append(line)
        for ACGT_counts in zip(m.counts['A'], m.counts['C'], m.counts['G'], m.counts['T']):
            lines.append('{0:0.0f}\t{1:0.0f}\t{2:0.0f}\t{3:0.0f}\n'.format(*ACGT_counts))

    # Finished; glue the lines together.
    text = ''.join(lines)

    return text
