"""
Plot histograms of differences to reference V gene

For each gene, a histogram is plotted that shows how often a sequence was
assigned to that gene at a certain percentage difference.
"""
import sys
import logging
import numpy as np
import seaborn as sns
from matplotlib.backends.backend_pdf import FigureCanvasPdf, PdfPages
from matplotlib.figure import Figure

from ..table import read_table

logger = logging.getLogger(__name__)


def add_arguments(parser):
    arg = parser.add_argument
    arg('--minimum-group-size', '-m', metavar='N', default=None, type=int,
        help="Plot only genes with at least N assigned sequences. "
        "Default: 0.1%% of assigned sequences or 100, whichever is smaller.")
    arg('--max-j-shm', metavar='VALUE', type=float, default=None,
        help='Use only rows with J%%SHM >= VALUE')
    arg('--multi', metavar='PDF', default=None,
        help='Plot individual error frequency histograms (for each V gene) to this PDF file')
    arg('--boxplot', metavar='PDF', default=None,
        help='Plot a single page with box(en)plots of V SHM for multiple genes')
    arg('table', metavar='FILTERED.TAB.GZ', help='Table with parsed IgBLAST results')


def plot_difference_histogram(group, gene_name, bins=np.arange(20.1)):
    """
    Plot a histogram of percentage differences for a specific gene.
    """
    exact_matches = group[group.V_SHM == 0]
    cdr3s_exact = len(set(s for s in exact_matches.CDR3_nt if s))
    js_exact = len(set(exact_matches.J_gene))

    fig = Figure(figsize=(100/25.4, 60/25.4))
    ax = fig.gca()
    ax.set_xlabel('Percentage difference')
    ax.set_ylabel('Frequency')
    fig.suptitle('Gene ' + gene_name, y=1.08, fontsize=16)
    ax.set_title('{:,} sequences assigned'.format(len(group)))

    ax.text(0.25, 0.95,
        '{:,} ({:.1%}) exact matches\n  {} unique CDR3\n  {} unique J'.format(
            len(exact_matches), len(exact_matches) / len(group),
            cdr3s_exact, js_exact),
        transform=ax.transAxes, fontsize=10,
        bbox=dict(boxstyle='round', facecolor='white', alpha=0.5),
        horizontalalignment='left', verticalalignment='top')
    _ = ax.hist(list(group.V_SHM), bins=bins)
    return fig


def main(args):
    table = read_table(args.table, usecols=['V_gene', 'J_gene', 'V_SHM', 'J_SHM', 'CDR3_nt'])
    if not args.multi and not args.boxplot:
        print('Don’t know what to do', file=sys.stderr)
        sys.exit(2)

    # Discard rows with any mutation within J at all
    logger.info('%s rows read', len(table))
    if args.max_j_shm is not None:
        # Discard rows with too many J mutations
        table = table[table.J_SHM <= args.max_j_shm][:]
        logger.info('%s rows remain after keeping only those with J%%SHM <= %s',
            len(table), args.max_j_shm)

    if args.minimum_group_size is None:
        total = len(table)
        minimum_group_size = min(total // 1000, 100)
        logger.info('Skipping genes with less than %s assignments', minimum_group_size)
    else:
        minimum_group_size = args.minimum_group_size

    # Genes with high enough assignment count
    all_genes = table['V_gene'].unique()
    genes = sorted(table['V_gene'].value_counts().loc[lambda x: x >= minimum_group_size].index)
    gene_set = set(genes)

    logger.info('%s out of %s genes have enough assignments', len(genes), len(all_genes))
    if args.multi:
        with PdfPages(args.multi) as pages:
            for gene, group in table.groupby('V_gene'):
                if gene not in gene_set:
                    continue
                fig = plot_difference_histogram(group, gene)
                FigureCanvasPdf(fig).print_figure(pages, bbox_inches='tight')

        logger.info('Wrote %r', args.multi)

    if args.boxplot:
        aspect = 1 + len(genes) / 32
        g = sns.catplot(x='V_gene', y='V_SHM', kind='boxen', order=genes, data=table,
            height=2 * 2.54, aspect=aspect, color='g')
        # g.set(ylim=(-.1, None))
        g.set(ylabel='% V SHM (nt)')
        g.set(xlabel='V gene')
        g.set_xticklabels(rotation=90)
        g.savefig(args.boxplot)
        logger.info('Wrote %r', args.boxplot)
