"""Genomic track visualization utilities."""

import warnings

from itertools import product

import matplotlib.pyplot as plt

import numpy as np

import seaborn as sns

from janggu.utils import NMAP

from janggu.utils import PMAP

from janggu.utils import _to_list

def plotGenomeTrack(tracks, chrom, start, end, figsize=(10, 5), plottypes=None):

    """plotGenomeTrack shows plots of a specific interval from cover objects data.

    It takes one or more cover objects as well as a genomic interval consisting

    of chromosome name, start and end and creates

    a genome browser-like plot.

    Parameters

    ----------

    tracks : janggu.data.Cover, list(Cover), janggu.data.Track or list(Track)

        One or more track objects.

    chrom : str

        chromosome name.

    start : int

        The start of the required interval.

    end : int

        The end of the required interval.

    figsize : tuple(int, int)

        Figure size passed on to matplotlib.

    plottype : None or list(str)

        Plot type indicates whether to plot coverage tracks as line plots,

        heatmap, or seqplot using 'line' or 'heatmap', respectively.

        By default, all coverage objects are depicted as line plots if plottype=None.

        Otherwise, a list of types must be supplied containing the plot types for each

        coverage object explicitly. For example, ['line', 'heatmap', 'seqplot'].

        While, 'line' and 'heatmap' can be used for any type of coverage data,

        'seqplot' is reserved to plot sequence influence on the output. It is

        intended to be used in conjunction with 'input_attribution' method which

        determines the importance of paricular sequence letters for the output prediction.

    Returns

    -------

    matplotlib Figure

        A matplotlib figure illustrating the genome browser-view of the coverage

        objects for the given interval.

        To depict and save the figure the native matplotlib functions show()

        and savefig() can be used.

    """

    tracks = _to_list(tracks)

    for track in tracks:

        if not isinstance(track, Track):

            warnings.warn('Convert the Dataset object to proper Track objects.'

                          ' In the future, only Track objects will be supported.',

                          FutureWarning)

            if plottypes is None:

                plottypes = ['line'] * len(tracks)

            assert len(plottypes) == len(tracks), \

                "The number of cover objects must be the same as the number of plottyes."

            break

    def _convert_to_track(cover, plottype):

        if plottype == 'heatmap':

            track = HeatTrack(cover)

        elif plottype == 'seqplot':

            track = SeqTrack(cover)

        else:

            track = LineTrack(cover)

        return track

    tracks_ = []

    for itrack, track in enumerate(tracks):

        if isinstance(track, Track):

            tracks_.append(track)

        else:

            warnings.warn('Convert the Dataset object to proper Track objects.'

                          ' In the future, only Track objects will be supported.',

                          FutureWarning)

            tracks_.append(_convert_to_track(track, plottypes[itrack]))

    tracks = tracks_

    headertrack = 2

    trackheights = 0

    for track in tracks:

        trackheights += track.height

    spacer = len(tracks) - 1

    grid = plt.GridSpec(headertrack + trackheights + spacer,

                        10, wspace=0.4, hspace=0.3)

    fig = plt.figure(figsize=figsize)

    # title and reference track

    title = fig.add_subplot(grid[0, 1:])

    title.set_title(chrom)

    plt.xlim([0, end - start])

    title.spines['right'].set_visible(False)

    title.spines['top'].set_visible(False)

    title.spines['left'].set_visible(False)

    plt.xticks([0, end-start], [start, end])

    plt.yticks(())

    y_offset = 1

    for track in tracks:

        y_offset += 1

        track.add_side_bar(fig, grid, y_offset)

        track.plot(fig, grid, y_offset, chrom, start, end)

        y_offset += track.height

    return (fig)

class Track(object):

    """General track

    Parameters

    ----------

    data : Cover object

        Coverage object

    height : int

        Track height.

    """

    def __init__(self, data, height):

        self.height = height

        self.data = data

    @property

    def name(self):

        """Track name"""

        return self.data.name

    def add_side_bar(self, fig, grid, offset):

        """Side-bar"""

        # side bar indicator for current cover

        ax = fig.add_subplot(grid[(offset): (offset + self.height), 0])

        ax.set_xticks(())

        ax.spines['right'].set_visible(False)

        ax.spines['top'].set_visible(False)

        ax.spines['bottom'].set_visible(False)

        ax.set_yticks([0.5])

        ax.set_yticklabels([self.name])

    def get_track_axis(self, fig, grid, offset, height):

        """Returns axis object"""

        return fig.add_subplot(grid[offset:(offset + height), 1:])

    def get_data(self, chrom, start, end):

        """Returns data to plot."""

        return self.data[chrom, start, end][0, :, :, :]

class LineTrack(Track):

    """Line track

    Visualizes genomic data as line plot.

    Parameters

    ----------

    data : Cover object

        Coverage object

    height : int

        Track height. Default=3

    linestyle : str

        Linestyle for plot

    marker : str

        Marker code for plot

    color : str

        Color code for plot

    linewidth : float

        Line width.

    """

    def __init__(self, data, height=3, linestyle='-', marker='o', color='b',

                 linewidth=2):

        super(LineTrack, self).__init__(data, height)

        self.height = height * len(data.conditions)

        self.linestyle = linestyle

        self.linewidth = linewidth

        self.marker = marker

        self.color = color

    def plot(self, fig, grid, offset, chrom, start, end):

        """Plot line track."""

        coverage = self.get_data(chrom, start, end)

        offset_ = offset

        trackheight = self.height//len(self.data.conditions)

        def _get_xy(cov):

            xvalue = np.where(np.isfinite(cov))[0]

            yvalue = cov[xvalue]

            return xvalue, yvalue

        for i, condition in enumerate(self.data.conditions):

            ax = self.get_track_axis(fig, grid, offset_, trackheight)

            offset_ += trackheight

            if coverage.shape[1] == 2:

                #both strands are covered separately

                xvalue, yvalue = _get_xy(coverage[:, 0, i])

                ax.plot(xvalue, yvalue,

                        linewidth=self.linewidth,

                        linestyle=self.linestyle,

                        color=self.color, label="+", marker='+')

                xvalue, yvalue = _get_xy(coverage[:, 1, i])

                ax.plot(xvalue, yvalue,

                        linewidth=self.linewidth,

                        linestyle=self.linestyle,

                        color=self.color, label="-", marker=1)

                ax.legend()

            else:

                xvalue, yvalue = _get_xy(coverage[:, 0, i])

                ax.plot(xvalue, yvalue, linewidth=self.linewidth,

                        color=self.color,

                        linestyle=self.linestyle,

                        marker=self.marker)

            ax.set_yticks(())

            ax.set_xticks(())

            ax.set_xlim([0, end-start])

            if len(self.data.conditions) > 1:

                ax.set_ylabel(condition, labelpad=12)

            ax.spines['right'].set_visible(False)

            ax.spines['top'].set_visible(False)

class SeqTrack(Track):

    """Sequence Track

    Visualizes sequence importance.

    Parameters

    ----------

    data : Cover object

        Coverage object

    height : int

        Track height. Default=3

    """

    def __init__(self, data, height=3):

        super(SeqTrack, self).__init__(data, height)

    def plot(self, fig, grid, offset, chrom, start, end):

        """Plot sequence track"""

        if len(self.data.conditions) % len(NMAP) == 0:

            alphabetsize = len(NMAP)

            MAP = NMAP

        elif len(self.data.conditions) % len(PMAP) == 0:  # pragma: no cover

            alphabetsize = len(PMAP)

            MAP = PMAP

        else:  # pragma: no cover

            raise ValueError(

                "Coverage tracks seems not represent biological sequences. "

                "The last dimension must be divisible by the alphabetsize.")

        for cond in self.data.conditions:

            if cond[0] not in MAP:

                raise ValueError(

                    "Coverage tracks seems not represent biological sequences. "

                    "Condition names must represent the alphabet letters.")

        coverage = self.get_data(chrom, start, end)

        # project higher-order sequence structure onto original sequence.

        coverage = coverage.reshape(coverage.shape[0], -1)

        coverage = coverage.reshape(coverage.shape[:-1] +

                                    (alphabetsize,

                                     int(coverage.shape[-1]/alphabetsize)))

        coverage = coverage.sum(-1)

        ax = self.get_track_axis(fig, grid, offset, self.height)

        x = np.arange(coverage.shape[0])

        y_figure_offset = np.zeros(coverage.shape[0])

        handles = []

        for letter in MAP:

            handles.append(ax.bar(x, coverage[:, MAP[letter]],

                                  bottom=y_figure_offset,

                                  color=sns.color_palette("hls",

                                                          len(MAP))[MAP[letter]],

                                  label=letter))

            y_figure_offset += coverage[:, MAP[letter]]

        ax.legend(handles=handles)

        ax.set_yticklabels(())

        ax.set_yticks(())

        ax.set_xticks(())

        ax.set_xlim([0, end-start])

class HeatTrack(Track):

    """Heatmap Track

    Visualizes genomic data as heatmap.

    Parameters

    ----------

    data : Cover object

        Coverage object

    height : int

        Track height. Default=3

    """

    def __init__(self, data, height=3):

        super(HeatTrack, self).__init__(data, height)

    def plot(self, fig, grid, offset, chrom, start, end):

        """Plot heatmap track."""

        ax = self.get_track_axis(fig, grid, offset, self.height)

        coverage = self.get_data(chrom, start, end)

        im = ax.pcolor(coverage.reshape(coverage.shape[0], -1).T)

        if coverage.shape[-2] == 2:

            ticks = [':'.join([x, y]) for y, x \

                     in product(['+', '-'], self.data.conditions)]

        else:

            ticks = self.data.conditions

        ax.set_yticklabels(ticks)

        ax.set_xticks(())

        ax.set_yticks(np.arange(0, len(ticks) + 1, 1.0))

        ax.set_xlim([0, end-start])