# This file is part of BYASE.
#
# BYASE is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# BYASE is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with BYASE.  If not, see <https://www.gnu.org/licenses/>.
#
# Author: Lili Dong
#

from typing import Optional, List, Dict
from enum import Enum
import os
from xml.etree.ElementTree import ElementTree

import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.lines import Line2D
from matplotlib.patches import Rectangle
import HTSeq

from ..message import MessageCenter
from .task import Task
from .align import TaskAlign, BAMParam
from .result import TaskResultMeta


matplotlib.use('Agg')


class TaskPlotError(Exception):
    """Task plot error."""
    def __init__(self, task_id: str, msg):
        super().__init__('[PLOT ERROR] [TASK: {}] {}'.format(task_id, msg))


class TaskFig:
    """Task figure."""

    def __init__(self, fig_size, out_path: str):
        self.fig_size = fig_size
        self.out_path = out_path

    def __enter__(self):
        self.fig = plt.figure(figsize=self.fig_size)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fig.savefig(self.out_path)
        self.fig.clear()
        plt.close(self.fig)
        self._apply_patch(self.out_path)

    @staticmethod
    def _apply_patch(svg_path: str):
        """Apply patch to svg file generated by matplotlib to improve compatibility."""
        assert svg_path.endswith('.svg')
        tree = ElementTree(file=svg_path)
        root = tree.getroot()
        non_defs = []
        for child in root:
            if not child.tag.endswith('defs'):
                non_defs.append(child)
        for child in non_defs:
            root.remove(child)
            root.append(child)
        tree.write(svg_path)


_COLOR_BLUE = 'steelblue'
_COLOR_ORANGE = 'chocolate'
_COLOR_YELLOW = 'darkorange'
_COLOR_GREEN = 'forestgreen'
_COLOR_RED = 'indianred'
_COLOR_GREY = 'mediumpurple'
_COLOR_BROWN = '#8c564b'


def _nt_color(nt: str) -> str:
    """Nt color map."""
    if nt == 'A':
        return _COLOR_GREEN
    elif nt == 'C':
        return _COLOR_BLUE
    elif nt == 'G':
        return _COLOR_YELLOW
    elif nt == 'T':
        return _COLOR_RED
    else:
        return 'grey'


class _Coverage:
    """Coverage.

    Attributes:
        allele_num: Allele number.
        isoform_num: Isoform number.
        data: A List of GenomicArray objects that stores coverage.
        _determined: The determined status of each GenomicArray object.
        _determined_map: A dict that maps determined status to GenomicArray object.
    """

    def __init__(self, allele_num: Optional[int], isoform_num: Optional[int]):
        data = []
        determined_map = {}

        if (allele_num is None) or (isoform_num is None):
            determined = [True, False]
        else:
            # Order: (allele_determined, isoform_determined).
            determined = [(True, True), (True, False), (False, True), (False, False)]

        for determined_status in determined:
            ga = HTSeq.GenomicArray('auto', stranded=False, typecode='i')
            determined_map[determined_status] = ga
            data.append(ga)

        self.allele_num = allele_num
        self.isoform_num = isoform_num
        self.data = data
        self._determined = determined
        self._determined_map = determined_map

    def name(self, n: int):
        """The name of the nth GenomicArray object."""
        if n == len(self._determined):
            return 'Total'
        assert n < len(self._determined)
        determined_status = self._determined[n]
        if self.allele_num is None:
            return 'Isoform {}'.format(self.isoform_num + 1 if determined_status else '?')
        if self.isoform_num is None:
            return 'Allele {}'.format(self.allele_num + 1 if determined_status else '?')
        allele_determined, isoform_determined = determined_status
        allele_num = self.allele_num + 1 if allele_determined else '?'
        isoform_num = self.isoform_num + 1 if isoform_determined else '?'
        return 'A{}-I{}'.format(allele_num, isoform_num)

    def add(self, determined_status, ivs: List[HTSeq.GenomicInterval]):
        """Add counts for genomic intervals."""
        ga = self._determined_map[determined_status]
        for iv in ivs:
            ga[iv] += 1

    def calc_total(self):
        """Calculate total coverage."""
        total = HTSeq.GenomicArray('auto', stranded=False, typecode='i')
        for ga in self.data:
            for iv, v in ga.steps():
                if v != 0:
                    total[iv] += v
        self.data.append(total)

    def get_max(self) -> int:
        """Get max coverage."""
        max_coverage = 0
        assert len(self._determined) + 1 == len(self.data)
        total = self.data[-1]
        for _, v in total.steps():
            if v != 0:
                max_coverage = max(max_coverage, v)
        return max_coverage


class TaskPlotType(Enum):
    ASE = 'ase'


class TaskPlot(TaskAlign, TaskResultMeta):
    """Task plot.

    Attributes:
        trace: Result trace.
        trace_stats: Result trace stats.

        coord_offset_map: The coord offset map for display.
        cov_dict: A dict that stores coverage.

        out_dir: The output directory path.
        plot_type: Plot type.
    """

    def __init__(self, task: Task, bam_param: BAMParam,
                 trace: pd.DataFrame, trace_stats: pd.DataFrame,
                 out_dir: str, plot_type: TaskPlotType, mc: MessageCenter):
        super().__init__(task=task, bam_param=bam_param, mc=mc)

        self.trace = trace
        self.trace_stats = trace_stats
        self.out_dir = out_dir
        self.plot_type = plot_type

        self.coord_offset_map = None
        self.cov_dict = None  # type: Optional[Dict[tuple, _Coverage]]

    def _generate_coord_offset_map(self) -> HTSeq.GenomicArray:
        """Generate display coord offset map."""
        all_exons = self.segment.get_all_exons()
        total_len = sum([exon.end - exon.start for exon in all_exons])

        if len(all_exons) <= 1:
            inner_len = None
        else:
            inner_len = float(total_len) / 10.0 / float(len(all_exons) - 1)
            inner_len = max(int(inner_len), 5)

        coord_offset_map = HTSeq.GenomicArray('auto', stranded=False, typecode='i')
        offset = 0
        last_exon_end = None
        for exon in all_exons:
            if last_exon_end is None:
                offset -= exon.start
            else:
                offset += inner_len - (exon.start - last_exon_end)
            coord_offset_map[exon] = offset
            last_exon_end = exon.end
        return coord_offset_map

    def _get_coord(self, pos: int) -> int:
        """Get displayed coord of position."""
        iv = HTSeq.GenomicInterval(self.segment.iv.chrom, pos, pos + 1, '.')
        steps = [step for step in self.coord_offset_map[iv].steps()]
        assert len(steps) == 1
        _, offset = steps[0]
        return pos + offset

    def _load_coverage(self):
        """Load coverage."""
        cov_dict = {}  # type: Dict[tuple, _Coverage]
        for allele_num in list(range(self.ploidy)) + [None]:
            for iso_num in list(range(self.isoforms_count)) + [None]:
                if (allele_num is None) and (iso_num is None):
                    continue
                cov_dict[(allele_num, iso_num)] = _Coverage(allele_num=allele_num, isoform_num=iso_num)

        aln_num = -1
        for align_stat in self.align_stats_iterator():
            aln_num += 1
            if aln_num != 0 and aln_num % 1000 == 0:
                self.mc.handle_progress('{} alignments loaded...'.format(aln_num))

            isoform_determined, allele_determined = False, False
            if len(align_stat.compatible_isoforms) == 1:
                isoform_determined = True
            if len(align_stat.compatible_alleles) == 1:
                allele_determined = True

            ivs = align_stat.matched_ivs
            for allele_num in align_stat.compatible_alleles:
                cov_dict[(allele_num, None)].add(allele_determined, ivs)

            for iso_num in align_stat.compatible_isoforms:
                cov_dict[(None, iso_num)].add(isoform_determined, ivs)

            for allele_num in align_stat.compatible_alleles:
                for iso_num in align_stat.compatible_isoforms:
                    cov_dict[(allele_num, iso_num)].add((allele_determined, isoform_determined), ivs)

        for _, cov in cov_dict.items():
            cov.calc_total()

        return cov_dict

    def _coverage_path(self, allele_num: Optional[int], isoform_num: Optional[int]) -> str:
        """Plot path for coverage."""
        path = 'COV'
        if allele_num is not None:
            path += '_A{}'.format(allele_num)
        if isoform_num is not None:
            path += '_I{}'.format(isoform_num)
        path += '.svg'
        return os.path.join(self.out_dir, path)

    def _histogram_path(self, var_name: str) -> str:
        """Plot path for histogram."""
        path = 'HIST_{}.svg'.format(var_name)
        return os.path.join(self.out_dir, path)

    @property
    def _html_path(self) -> str:
        """Plot path for html."""
        return os.path.join(self.out_dir, 'plot_{}.html'.format(self.plot_type.value))

    def _plot_coverage(self, allele_num: Optional[int], isoform_num: Optional[int], y_lim: int):
        """Plot coverage."""
        out_path = self._coverage_path(allele_num=allele_num, isoform_num=isoform_num)
        with TaskFig(fig_size=(9, 3), out_path=out_path) as task_fig:

            x_lim = (self._get_coord(self.segment.iv.start) - 0.5,
                     self._get_coord(self.segment.iv.end - 1) + 0.5)

            # Plot isofrom.
            ax = task_fig.fig.add_axes([0.05, 0.05, 0.9, 0.1])
            ax.set_axis_off()

            if isoform_num is None:
                exons = self.segment.get_all_exons()
                snps = self.snps
            else:
                exons = self.segment.isoforms[isoform_num].exons
                snps = self.isoform_snps(isoform_num=isoform_num)
            exons = sorted(exons, key=lambda _iv: _iv.start)

            prev_exon_end = None
            for exon in exons:
                s = self._get_coord(exon.start)
                e = self._get_coord(exon.end - 1)
                ax.add_patch(Rectangle((s, 0), e - s, 1, color=_COLOR_GREY, linewidth=0, alpha=0.2))
                if prev_exon_end is not None:
                    ax.add_line(Line2D([prev_exon_end, s], [0.5, 0.5], linestyle=':', color='C4', alpha=0.2))
                prev_exon_end = e

            if allele_num is not None:
                for snp in snps:
                    x = self._get_coord(snp.pos)
                    ax.add_line(Line2D([x, x], [0, 1], color=_nt_color(snp.alleles[allele_num]), alpha=0.6))

            ax.set_xlim(x_lim)

            # Plot coverage.
            ax = task_fig.fig.add_axes([0.05, 0.2, 0.9, 0.7])

            if isoform_num is None:
                title = self.segment.gene_name
                colors = [_COLOR_BLUE, _COLOR_BROWN]
                alphas = [0.8, 0.35]
            elif allele_num is None:
                title = '{} - {}'.format(self.segment.gene_name, self.segment.isoforms[isoform_num].name)
                colors = [_COLOR_GREEN, _COLOR_BROWN]
                alphas = [0.5, 0.35]
            else:
                title = self.segment.isoforms[isoform_num].name
                colors = [_COLOR_BLUE, _COLOR_ORANGE, _COLOR_GREEN, _COLOR_RED]
                alphas = [0.8, 0.65, 0.5, 0.35]
            colors.append(_COLOR_GREY)
            alphas.append(0.2)

            cov = self.cov_dict[(allele_num, isoform_num)]
            for i, ga in enumerate(cov.data):
                x = []
                y = []
                for iv, v in ga.steps():
                    if v != 0:
                        s = self._get_coord(iv.start)
                        e = self._get_coord(iv.end - 1)
                        for k in [s-0.5, s-0.5, e+0.5, e+0.5]:
                            x.append(k)
                        for k in [0, v, v, 0]:
                            y.append(k)
                ax.fill_between(x, 0, y, lw=0, color=colors[i], alpha=alphas[i],
                                zorder=len(cov.data)-i, label=cov.name(i))

            ax.legend(loc='upper right', ncol=len(cov.data) + 1)
            ax.set_xlim(x_lim)
            ax.set_xticks([])
            ax.set_ylim(0, y_lim)

            ax.set_title(title, loc='left', fontsize=11)

    def _plot_histogram(self, var_name: str, title: str):
        """Plot histogram."""
        trace = self.trace[var_name]
        stats = self.trace_stats.loc[var_name, ]
        out_path = self._histogram_path(var_name)

        with TaskFig(fig_size=(3, 3), out_path=out_path) as task_fig:
            with plt.style.context('ggplot'):
                ax = task_fig.fig.add_axes([0.1, 0.25, 0.8, 0.65])

            ax.hist(trace, bins=np.arange(-2, 2, 0.01), density=True, histtype='stepfilled', color=_COLOR_BLUE)

            mean = stats['mean']
            hpd_2_5 = stats['hpd_2.5']
            hpd_97_5 = stats['hpd_97.5']

            ax.axvline(hpd_2_5, color='black', linestyle='--', alpha=0.3)
            ax.axvline(hpd_97_5, color='black', linestyle='--', alpha=0.3)
            ax.axvline(mean, color=_COLOR_RED, alpha=0.6)

            stat_text = 'Mean = {:.3f}\n95% HPD: ({:.3f}, {:.3f})'.format(mean, hpd_2_5, hpd_97_5)
            ax.set_xlabel(stat_text, labelpad=5, fontsize=11)
            ax.yaxis.set_major_locator(MaxNLocator(integer=True))

            ax.set_xticks(np.arange(0, 2, 0.2))
            ax.set_xlim(0, 1)

            ax.set_title(title, loc='left', fontsize=11)

    def _generate_html(self):
        """Generate html."""
        html = '<!DOCTYPE HTML>\n'
        html += '<html>\n'
        html += '<head>\n'
        html += _CSS
        html += '</head>\n'
        html += '<body>\n'
        html += '<table border="0" cellspacing="0" cellpadding="0">\n'

        def _add_td(_html: str, _td_class: str, _img_path: str):
            _html += '<td class="{}">'.format(_td_class)
            _html += '<img src="{}"/>'.format(os.path.basename(img_path))
            _html += '</td>\n'
            return _html

        if self.plot_type is TaskPlotType.ASE:
            for isoform_num in [None] + list(range(self.isoforms_count)):
                html += '<tr class="plot {}">\n'.format('g' if isoform_num is None else 'i{}'.format(isoform_num))
                for allele_num in range(self.ploidy):
                    var = self.get_var_expression(allele_num=allele_num, iso_num=isoform_num)
                    img_path = self._histogram_path(var)
                    td_class = 'plot hist hist_{}'.format(allele_num)
                    html = _add_td(html, td_class, img_path)

                    img_path = self._coverage_path(allele_num=allele_num, isoform_num=isoform_num)
                    td_class = 'plot cov cov_{}'.format(allele_num)
                    html = _add_td(html, td_class, img_path)

                for i, j in self.delta_iterator(self.ploidy):
                    var = self.get_var_diff_expression(allele_num1=i, allele_num2=j, iso_num=isoform_num)
                    img_path = self._histogram_path(var)
                    td_class = 'plot diff diff_{}_{}'.format(i, j)
                    html = _add_td(html, td_class, img_path)

                html += '</tr>\n'
        else:
            assert False

        html += '</table>\n'
        html += _SCRIPTS
        html += '</body>\n'
        html += '</html>\n'

        with open(self._html_path, 'w') as out:
            out.write(html)

    def plot(self) -> str:
        """Plot."""
        if os.path.exists(self._html_path):
            return self._html_path

        self.coord_offset_map = self._generate_coord_offset_map()
        self.cov_dict = self._load_coverage()

        self.mc.handle_progress('Plot histograms...')
        if self.plot_type is TaskPlotType.ASE:
            for allele_num in range(self.ploidy):
                var = self.get_var_expression(allele_num=allele_num, iso_num=None)
                title = 'Gene - Allele {}'.format(allele_num+1)
                self._plot_histogram(var, title)
                for iso_num in range(self.isoforms_count):
                    var = self.get_var_expression(allele_num=allele_num, iso_num=iso_num)
                    title = 'Isoform {} - Allele {}'.format(iso_num+1, allele_num+1)
                    self._plot_histogram(var, title)

            for i, j in self.delta_iterator(self.ploidy):
                var = self.get_var_diff_expression(i, j, iso_num=None)
                title = 'Gene - Diff (Allele {} & {})'.format(i+1, j+1)
                self._plot_histogram(var, title)
                for iso_num in range(self.isoforms_count):
                    var = self.get_var_diff_expression(i, j, iso_num=iso_num)
                    title = 'Isoform {} - Diff (Allele {} & {})'.format(iso_num+1, i+1, j+1)
                    self._plot_histogram(var, title)
        else:
            assert False

        self.mc.handle_progress('Plot coverages...')
        if self.plot_type is TaskPlotType.ASE:
            max_coverage = max([self.cov_dict[(allele_num, None)].get_max() for allele_num in range(self.ploidy)])
            for allele_num in range(self.ploidy):
                self._plot_coverage(allele_num=allele_num, isoform_num=None, y_lim=max_coverage)
        else:
            assert False

        for allele_num in range(self.ploidy):
            for iso_num in range(self.isoforms_count):
                self._plot_coverage(allele_num=allele_num, isoform_num=iso_num, y_lim=max_coverage)

        self._generate_html()
        return self._html_path


_CSS = """
<style>
    body {
        margin: 0;
    }    
    table {
        border-collapse: collapse;
    }
    img {
        vertical-align: bottom;
    }
</style>
"""

_SCRIPTS = """
<script>
    function update_plots(zoom, list) {
        var cells = document.getElementsByClassName("plot");
        for (var i = 0; i < cells.length; i++) {
            cells[i].style.display = "none";
            cells[i].getElementsByTagName("img")[0].style.width=zoom;
        }
        for (var n = 0; n < list.length; n++) {
            cells = document.getElementsByClassName(list[n]);
            for (var i = 0; i < cells.length; i++) {
                cells[i].style.display = "";
            }
        }
    }
</script>
"""
