"""
Draw a dendrogram of sequences in a FASTA file.
"""
import logging
import numpy as np
import matplotlib
import dnaio
from scipy.spatial import distance
from scipy.cluster import hierarchy
from ..utils import distances


logger = logging.getLogger(__name__)


def add_arguments(parser):
    arg = parser.add_argument
    arg('--mark', '--db', metavar='FASTA',
        help='Path to a FASTA file with a set of "known" sequences. Sequences '
        'in the main file that do *not* occur here will be marked with (new). '
        'If not given, no sequences will be marked (use this to compare two '
        'databases).')
    arg('--method', choices=('single', 'complete', 'weighted', 'average'),
        default='average',
        help='Linkage method. Default: "average" (=UPGMA)')
    arg('fasta', help='Path to input FASTA file')
    arg('plot', help='Path to output PDF or PNG')


class PrefixComparer:
    def __init__(self, sequences):
        self._sequences = [ s.upper() for s in sequences ]

    def __contains__(self, other):
        for seq in self._sequences:
            if seq.startswith(other) or other.startswith(seq):
                return True
        return False


def main(args):
    with dnaio.open(args.fasta) as fr:
        sequences = list(fr)
    logger.info('Plotting dendrogram of %s sequences', len(sequences))
    if args.mark:
        with dnaio.open(args.mark) as fr:
            mark = PrefixComparer(record.sequence for record in fr)
        labels = []
        n_new = 0
        for record in sequences:
            if record.sequence not in mark:
                extra = ' (new)'
                n_new += 1
            else:
                extra = ''
            labels.append(record.name + extra)
        logger.info('%s sequence(s) marked as "new"', n_new)
    else:
        labels = [s.name for s in sequences]

    import seaborn as sns
    import matplotlib.pyplot as plt
    sns.set()
    sns.set_style("white")
    font_size = 297 / 25.4 * 72 / (len(labels) + 5)
    font_size = min(16, max(6, font_size))
    height = font_size * (len(labels) + 5) / 72
    fig = plt.figure(figsize=(210 / 25.4, height))
    matplotlib.rcParams.update({'font.size': 4})
    ax = fig.gca()
    sns.despine(ax=ax, top=True, right=True, left=True, bottom=True)
    sns.set_style('whitegrid')
    if len(sequences) >= 2:
        m = distances([s.sequence for s in sequences])
        y = distance.squareform(m)
        mindist = int(y.min())
        logger.info('Smallest distance is %s. Found between:', mindist)
        for i,j in np.argwhere(m == y.min()):
            if i < j:
                logger.info('%s and %s', labels[i], labels[j])
        l = hierarchy.linkage(y, method=args.method)
        hierarchy.dendrogram(l, labels=labels, leaf_font_size=font_size, orientation='right', color_threshold=0.95*max(l[:,2]))
    else:
        ax.text(0.5, 0.5, 'no sequences', fontsize='xx-large')
    ax.grid(False)
    fig.set_tight_layout(True)
    fig.savefig(args.plot)
