#!/usr/bin/python3
# 2021 Leo Kaindl

"""
Makes result HTML interactive by enabling motif search.
Compute neutrality/diversity stats for selected population.
"""

import cgi
import cgitb
import itertools

import numpy as np
import pickle
import re
import pandas as pd
import toytree

from os import path, makedirs
from lxml import etree, html
from lxml.html import builder as E
from Bio import SeqIO
from toytree.utils import ToytreeError


def _exclude(ex_motifs, leaves):
    # filter leaves with exclusion motifs
    if ex_motifs != '':
        # compile exclusion regex.
        queries = ['.*' + re.escape(word.strip()) + '.*' for word in ex_motifs.split(',')]
        if '.*.*' in queries and len(queries) > 1:
            queries.remove('.*.*')
        regex = re.compile(r'|'.join(queries))

        # drop excluded tips
        for dropped_tip in list(filter(regex.match, leaves)):
            leaves.remove(dropped_tip)
    return leaves


def _to_files(leaves):
    # translate to file paths
    with open('metadata.tsv', 'r') as tsv:
        # prep an empty list for each leaf
        sam_to_file = {i: list() for i in range(len(leaves))}
        # iterate over the table and for matching IDs append file paths
        [sam_to_file[leaves.index(sample)].append(_file) for sample, _file
         in [line.split('\t')[0:3:2]  # get first and third column
             for line in tsv.readlines()] if sample in leaves]
    return ['\n'.join(files) for sample, files in sorted(sam_to_file.items())]


def _h(a):
    # harmonic number
    b = 0
    for c in range(1, a):
        b += 1 / c
    return b


def _qh(a):
    # quadratic harmonic number
    b = 0
    for c in range(1, a):
        b += 1 / np.power(c, 2)
    return b


def _diversity_stats(gene_now, records, limits, poly):
    # translate dictionary of sequences to numpy int array
    seqs = np.array([seq for seq in records.values()])
    n_sequences = len(records)
    bases = ['A', 'C', 'G', 'T', '-']
    weird_chars = np.setdiff1d(np.unique(seqs), bases)
    codes = dict(zip(bases + list(weird_chars), range(len(bases) + len(weird_chars))))
    coded_seqs = np.vectorize(codes.get)(seqs)

    conserved, biallelic, polyallelic, unknown, gaps = 0, 0, 0, 0, 0
    drop = set()

    # iterate over MSA columns, search for segregating sites
    for j in range(coded_seqs.shape[1]):
        col = coded_seqs[:, j]
        if len(set(col)) == 1:
            # only one char at site
            it = col[0]
            if it > 4:
                # all-unknown-site. can happen as we selected a subset of sequences
                unknown += 1
                drop.add(j)
            elif it == 4:
                # all-gap-site
                gaps += 1
                drop.add(j)
            else:
                conserved += 1
        else:  # more than one character in the whole column
            if max(col) > 4 and len(col[col > 4]) > limits[1] * n_sequences:
                # there are too many unknown chars in the column
                unknown += 1
                drop.add(j)
            elif max(col) == 4 and len(col[col >= 4]) > limits[0] * n_sequences:
                # there are too many gaps at the site
                gaps += 1
                drop.add(j)
            else:
                # if gaps or unknown characters are at the site, replace them with the most common nucleotide
                try:
                    col[col >= 4] = np.bincount(col[col < 4]).argmax()
                except ValueError:
                    # there are no nucleotides
                    drop.add(j)
                    continue

                # ordinary site treatment
                num_diff_chars = len(set(col))
                if num_diff_chars > 2:
                    polyallelic += 1
                    if not poly:
                        drop.add(j)
                elif num_diff_chars == 2:
                    biallelic += 1
                else:
                    conserved += 1
    # print('GENE:%s\tallsites:%d\tconserved:%d\tbiallelic:%d\tpolyallelic:%d\tunknown:%d\tgaps:%d\tpoly:%s\ndrop:%d\n'
    #       % (gene_now, coded_seqs.shape[1], conserved, biallelic, polyallelic, unknown, gaps, poly, len(drop)))

    seg_sites = biallelic
    n_sites = coded_seqs.shape[1] - unknown - gaps
    if poly:
        seg_sites += polyallelic
    else:
        n_sites -= polyallelic

    if seg_sites == 0:
        return '<tr><td>%s</td><td>%d</td><td>0</td><td>--</td><td>--</td><td>--</td><td>--</td>' \
               '<td>--</td><td>%d</td><td>%d</td></tr>' % (gene_now, n_sites, gaps, unknown)

    # crop to allowed sites
    coded_seqs = coded_seqs[:, list(set(range(coded_seqs.shape[1])) - drop)]

    k = 0
    # count pairwise differences
    for combi in itertools.combinations(range(n_sequences), 2):
        k += np.sum(coded_seqs[combi[0]] != coded_seqs[combi[1]])

    # binomial coefficient, k = k^ in formula (10) from Tajima1989
    k = k * 2 / n_sequences / (n_sequences - 1)
    pi = k / n_sites

    a1 = _h(n_sequences)
    theta_w = seg_sites / a1 / n_sites  # per site from Yang2014

    # Tajima's D Formula:
    # d = k - seg_sites/a1 = k - theta_W
    # D = d / sqrt(Var(d))
    try:
        a2 = _qh(n_sequences)
        e1 = 1 / a1 * ((n_sequences + 1) / (3 * n_sequences - 3) - 1 / a1)
        b2 = (2 * (n_sequences * n_sequences + n_sequences + 3)) / (9 * n_sequences * (n_sequences - 1))
        c2 = b2 - ((n_sequences + 2) / (n_sequences * a1)) + (a2 / (a1 * a1))
        e2 = c2 / (a1 * a1 + a2)
        tajima = (k - theta_w) / np.sqrt(e1 * seg_sites + e2 * seg_sites * (seg_sites - 1))
    except (ZeroDivisionError, RuntimeWarning) as z:
        tajima = float('inf')

    n_genotypes = len(np.unique(coded_seqs, axis=0))

    return '<tr><td>%s</td><td>%d</td><td>%d</td><td>%.5f</td><td>%.5f</td>' \
           '<td>%.5f</td><td>%.5f</td><td>%d</td><td>%d</td><td>%d</td></tr>' \
           % (gene_now, n_sites, seg_sites, k, pi, theta_w, tajima, n_genotypes, gaps, unknown)


cgitb.enable()
# cgi.test()
print('Content-type: text/html\n\n')

# get form data
form = cgi.FieldStorage()
call = form['submit'].value
motifs = form['motifs'].value
thresholds = float(form['gap'].value), float(form['unknown'].value)
poly_allowed = True if form['poly'].value == 'True' else False
as_files = True if form['out_type'].value == 'files' else False
try:
    ex_motifs = form['exclude'].value
except KeyError:
    ex_motifs = ''
with open(form['pickle'].value, 'rb') as pickle_fh:
    tree = pickle.load(pickle_fh)

try:
    makedirs('queries', exist_ok=True)

    # compile motifs regex. empty motifs is caught before
    queries = ['.*' + re.escape(word.strip()) + '.*' for word in motifs.split(',')]
    regex = re.compile(r'|'.join(queries))

    # find matching tips
    leaves = list(filter(regex.match, tree.get_tip_labels()))

    # filter with exclusion pattern
    leaves = _exclude(ex_motifs, leaves)

    if len(leaves) == 0:
        raise ValueError('no matches')

    if call == 'subtree':
        mrca = tree.get_mrca_idx_from_tip_labels(names=leaves)
        leaves = tree.get_tip_labels(mrca)
        leaves = _exclude(ex_motifs, leaves)
        leaves = [leaf.split(' ')[0] if 'strain' not in leaf else leaf for leaf in leaves]
        files = _to_files(leaves)
        txt = '\n'.join(reversed(files)) if as_files else ', '.join(reversed(leaves))
        mrca = '%d match%s, subtree MRCA idx: %s' % (len(leaves), 'es' if len(leaves) > 1 else '', mrca)
    elif call == 'match':
        leaves = [leaf.split(' ')[0] if 'strain' not in leaf else leaf for leaf in leaves]
        files = _to_files(leaves)
        txt = '\n'.join(files) if as_files else ', '.join(leaves)
        mrca = '%d match%s' % (len(leaves), 'es' if len(leaves) > 1 else '')
    else:
        raise ValueError('illegal call')

except (ToytreeError, ValueError) as ex:
    txt = str(ex)
    mrca = ''
    leaves = []

# save result in file
_path = path.join('queries', call + '_' + motifs + '_' + ex_motifs)
_path += '_files.txt' if as_files else '_samples.txt'
with open(_path, 'w') as txt_fh:
    txt_fh.write(txt.replace(', ', '\n') + '\n')

# diversity stats
if len(leaves) > 1:

    # check if a reference is among the selected leaves:
    picked_refs = set()
    picked_ids = set()
    for leaf in leaves:
        if 'strain' in leaf:
            picked_refs.add(leaf[leaf.index('strain'):])
        else:
            picked_ids.add(leaf)

    # if a reference is selected
    if picked_refs:
        # read in tabular data
        df = pd.read_csv(form['metadata'].value, sep='\t', dtype={'id': str})
        df.set_index('id', inplace=True)
        df = df.reference_species[pd.notna(df.reference_species)]  # pandas.core.series.Series

        ref_ids = set()
        for k, v in df.items():
            strain = v[v.index('strain'):]
            if strain in picked_refs:
                picked_refs.remove(strain)
                ref_ids.add(k)

        # merge picked_ids and ref_ids
        picked_ids |= ref_ids

    # per-gene metrics and Tajima's D in sliding-window
    pop_seqs = {record.id.split(' ')[0]: record.seq
                for record in SeqIO.parse(form['msa_path'].value, 'fasta')
                if record.id.split(' ')[0] in picked_ids}

    # get gene lengths from form
    g_lens = [entry.split(':') for entry in form['g_lens'].value.split('_')]
    pos = 0
    rows = ''

    # calculate diversity/neutrality statistics per gene
    for entry in g_lens:
        gene, glen = entry[0], int(entry[1])
        cut_seqs = {seq_id: seq[pos:pos + glen] for seq_id, seq in pop_seqs.items()}
        rows += _diversity_stats(gene, cut_seqs, thresholds, poly_allowed)
        pos += glen

    # overall if sensible
    if len(g_lens) > 1:
        rows += _diversity_stats('overall', pop_seqs, thresholds, poly_allowed)

    # write raw html table
    table = '<hr/><h4>Neutrality + Diversity Statistics</h4><p>Exclusion threshold for gap sites ' \
            'was <code>%.2f</code>, and <code>%.2f</code> for unknown sites.</p>' \
            '<table border="1" class="dataframe"><thead><tr style="text-align: justify;"><th></th>' \
            '<th>valid sites</th><th>S</th><th>k</th><th>π</th><th>Watterson\'s θ</th>' \
            '<th>Tajima\'s D</th><th>Genotypes</th><th>gaps</th><th>unknown</th></tr></thead>' \
            '<tbody>%s</tbody></table>' % (thresholds[0], thresholds[1], rows)
else:
    table = ''

# read in original page
ET = etree.parse('result.html', parser=etree.HTMLParser())
port = int(ET.findall('//meta[@name="port"]')[0].get('content'))

# fetch div of motif search results
div = ET.findall('//div[@id="motif_result"]')[0]
div.clear()

# insert matches, mrca info, file link and stats into HTML
div.append(E.PRE(E.CODE(txt)))
div.append(E.P(mrca, E.SPAN('see ', E.A('file', href='http://localhost:%s/%s' % (port, _path)), style='float:right')))
div.append(etree.fromstring('<div>%s</div>' % table))

# edit logo and icon path
ET.xpath('//img[@id="logo"]')[0].set('src', 'http://localhost:%d/.img/favicon.png' % port)
ET.xpath('//link[@id="favicon"]')[0].set('href', 'http://localhost:%d/.img/favicon.ico' % port)
msa_path = ET.xpath('//img[@id="msa_viz"]')
if len(msa_path) > 0:
    msa_path[0].set('src', 'http://localhost:%d/rectangular_msa.png' % port)

# re-set metadata.tsv link
ET.xpath('//a[@id="tsv_link"]')[0].set('href', '../../metadata.tsv')
ET.xpath('//iframe[@id="mview_link"]')[0].set('src', '../../msa_mview.html')
# paste into input fields
ET.xpath('//input[@id="motifs"]')[0].set('value', motifs)
ET.xpath('//input[@id="exclude"]')[0].set('value', ex_motifs)

# create a markup for selected nodes
markup = ';font-weight: bold; font-style: italic; text-decoration: underline 0.15rem'

for g in ET.findall('//svg[@class="toyplot-canvas-Canvas"]/g/g/g[@class="'
                    'toyplot-mark-Text"]/g/g[@class="toyplot-Datum"]/text'):
    st = g.get('style').replace(markup, '')
    # remove old markup
    g.set('style', st)
    # account for differently named refs in leaves list
    tip_id = g.text
    tip_id = tip_id.split(' ')[0] if 'strain' not in tip_id else tip_id
    if tip_id in leaves:
        # insert new markup
        g.set('style', st + markup)

# convert to raw text
print(html.tostring(ET, encoding='unicode'))
