#
# CHANGES
#
# 2023, Lewis Blake, Arjo Segers
#   Formatted using "black".
#

# -------------------------------------------------
# help
# -------------------------------------------------

"""
**********************
``utopya_post`` module
**********************

The :py:mod:`utopya_post` module provides classes to post process results.
All classes are accesible through the top level :py:mod:`utopya` module.

* The :py:class:`.UtopyaHtmlFile` class could be used to fill the
  body of an html file.

* The :py:class:`UtopyaBarPlot` class facilitates creation of (stacked) bar plots.

The configuration of the objects is done using rcfile settings.


Class hierarchy
===============

The classes provided by this module have been derived with the following hierchy:

* :py:class:`.UtopyaBase`

  * :py:class:`.UtopyaHtmlFile`

  * :py:class:`.UtopyaPlot`

    * :py:class:`.UtopyaBarPlot`

Classes
=======

"""


# -------------------------------------------------
# modules
# -------------------------------------------------

# modules:
import logging

# tools:
import utopya_base


############################################################
###
### HTML
###
############################################################


class UtopyaHtmlFile(utopya_base.UtopyaBase):

    """
    Storage for content of HTML file.
    A list attribute ``lines`` is created to hold the body of the content.

    Use the :py:meth:`Write` method to write the content to a file.
    """

    def __init__(self):

        """
        Initialize empty content.
        """

        # init ancestor:
        utopya_base.UtopyaBase.__init__(self)

        # no content yet:
        self.lines = []

    # enddef __init__

    # *

    def Append(self, line):

        """
        Add line to content.
        """

        # add:
        self.lines.append(line)

    # enddef Append

    # *

    def Prepend(self, line):

        """
        Add line to content.
        """

        # add:
        self.lines.insert(0, line)

    # enddef Prepend

    # *

    def Extend(self, html):

        """
        Add lines from other :py:class:`UtopyaHtmlFile` object.
        """

        # add:
        self.lines.extend(html.lines)

    # enddef Extend

    # *

    def Write(self, filename, title=None, info=False):

        """
        Write content to HTML file.

        This routine will add add the header lines,
        including a title if provided.

        Enable the ``info`` flag to write an informative
        message on the created file.
        """

        # modules:
        import os

        # target directory:
        dname = os.path.dirname(filename)
        # defined?
        if len(dname) > 0:
            # create if necessary:
            if not os.path.isdir(dname):
                os.makedirs(dname)
        # endif

        # create:
        f = open(filename, "w")
        # add header:
        f.write("<html>\n")
        f.write("<header>\n")
        if title is not None:
            f.write("<title>%s</title>\n" % title)
        f.write("</header>\n")

        # add body:
        f.write("<body>\n")
        for line in self.lines:
            f.write("%s\n" % line)
        f.write("</body>\n")

        # end:
        f.write("</html>\n")

        # info?
        if info:
            logging.info("")
            logging.info("Browse to:")
            logging.info("  file://%s" % os.path.abspath(filename))
            logging.info("")
        # endif

    # enddef Write

    # *

    def GetImgLink(self, img_src, img_attr=""):

        """
        Return html code to display an image with given attributes;
        a link is created to the full size version.

        Example::

           self.GetImgLink( 'test.png', 'width="400"' )
        """

        # html line:
        line = '<a href="%s"><img src="%s" border="0" %s></a>' % (img_src, img_src, img_attr)

        # ok
        return line

    # enddef GetImgLink


# endclass UtopyaHtmlFile


############################################################
###
### plots
###
############################################################


class UtopyaPlot(utopya_base.UtopyaBase):

    """
    Base class for figures.
    An attribute ``fig`` of the :py:class:`matplotlib.pyplot.figure`
    class is created, initialized with the optional keyword arguments.
    """

    def __init__(self, **kwargs):

        """
        Initialize base class.
        """

        # modules:
        import matplotlib.pyplot as plt

        # initialize ancestor:
        utopya_base.UtopyaBase.__init__(self)

        # new figure:
        self.fig = plt.figure(**kwargs)

    # enddef __init__

    # *

    def Close(self):

        """
        Close file.
        """

        # modules:
        import matplotlib.pyplot as plt

        # close:
        plt.close(self.fig)

    # enddef __init__

    # *

    def Export(self, filename):

        """
        Export figure to provided filename.
        """

        # tools:
        import os

        # destination directory:
        dname = os.path.dirname(filename)
        # defined?
        if len(dname) > 0:
            # create if necessary:
            if not os.path.isdir(dname):
                os.makedirs(dname)
        # endif

        # export:
        self.fig.savefig(filename)

    # enddef Export


# endclass UtopyaPlot

# *


class UtopyaBarPlot(UtopyaPlot):

    """
    Create bar plot with timing values.
    A vertical bar is drawn for one or more experiment;
    different values are stacked.

    Example usage::

        # initialize plot:
        p = UtopyaBarPlot()

        # add timing values for 2 experiments:
        p.Add( [2.0,3.0], label='base' )
        p.Add( [1.0,4.0], label='extra' )

        # create plot:
        p.Finish( figname='bars.png' )

    The optional arguments of the :py:meth:`.finish` method can
    be used for fine-tuning the result.

    """

    def __init__(self, **kwargs):

        """
        Initialize figure.
        """

        # initialize ancestor:
        UtopyaPlot.__init__(self)

        # storage for values and associated labels:
        self.values = []
        self.labels = []

    # enddef

    # *

    def Add(self, y, label=None):

        """
        Define list of values for the bars.
        Elements might be None if not available.

        The optional label will be used in a legend.
        """

        # storage:
        self.values.append(y)
        self.labels.append(label)

    # endif

    # *

    def Finish(self, figname=None, title=None, xlabs=[], ydefault=None, yref=None):

        """
        Create bar plot.

        Optional arguments:

        * ``figname`` : save result to file
        * ``title`` : figure title
        * ``xlabs`` : list with names on x-axis for the bars
        * ``ydefault`` : default value to be used for a bar if all values are None.
        * ``yref`` : reference value, if defined the y-axis is in percentage
             of this value
        """

        # modules:
        import numpy

        # settings:
        width = 0.8

        # number of added values:
        nvalue = len(self.values)
        # no values ?
        if nvalue == 0:
            # default specified ?
            if ydefault is None:
                # nothing to be done ..
                return
            else:
                # number of bars:
                nx = len(ydefault)
            # endif
        else:
            # number of bars:
            nx = len(self.values[0])
        # endif

        # dummy ax:
        xx = numpy.arange(1, nx + 1)

        # absolute or relative ?
        relative = (yref != None) and (yref != 0.0)
        if relative:
            ylab = "time [%]"
        else:
            ylab = "time"
        # endif

        # colors:
        colors = ["red", "green", "blue", "yellow", "cyan", "magenta", "brown", "steelblue"]

        # new figure:
        ax = self.fig.add_axes([0.1, 0.1, 0.5, 0.8])
        # current bottom:
        b = numpy.zeros(nx)
        # storage for handles and labels:
        pp = []
        ll = []
        # values ? otherwise use the default:
        if nvalue == 0:
            # current values:
            y = numpy.array(ydefault)
            # replace None by 0.0
            for k in range(len(y)):
                if y[k] == None:
                    y[k] = 0.0
            # endfor
            # absolute or relative ?
            if relative:
                yy = y / yref * 100.0  # %
            else:
                yy = y
            # endif
            # pick color:
            color = colors[0]
            # add bar plot:
            p = ax.bar(xx, yy, width=width, color=color)
        else:
            # current color index:
            icolor = 0
            # loop over all values in verse order:
            jj = range(len(self.values) - 1, -1, -1)
            for j in jj:
                # current values:
                y = numpy.array(self.values[j])
                label = self.labels[j]
                # replace None by 0.0
                for k in range(len(y)):
                    if y[k] == None:
                        y[k] = 0.0
                # endfor
                # absolute or relative ?
                if relative:
                    yy = y / yref * 100.0  # %
                    bb = b / yref * 100.0  # %
                else:
                    yy = y
                    bb = b
                # endif
                # pick color:
                color = colors[icolor]
                icolor = (icolor + 1) % len(colors)
                # add:
                p = ax.bar(xx, yy, bottom=bb, width=width, color=color, label=label)
                # next bottom:
                b = b + y
                # store handle and label:
                pp.append(p[0])
                ll.append(label)
            # endfor
            # add legend, revert order:
            pp.reverse()
            ll.reverse()
            ax.legend(pp, ll, loc="upper left", bbox_to_anchor=(1.0, 1.0))
        # endif   # values or default

        # x axis:
        ax.set_xticks(xx)
        ax.set_xticklabels(xlabs)
        # y axis:
        ax.set_ylabel(ylab)
        if relative:
            ylim = ax.get_ylim()
            ylim = [0.0, max(110, ylim[1] * 1.1)]
            ax.set_ylim(ylim)
        # title bar:
        if title is not None:
            ax.set_title(title)

        # save ?
        if figname is not None:
            # export to target file:
            self.Export(figname)
        # endif

        # clear:
        self.Close()

    # enddef


# enddef  # UtopyaBarPlot


# -------------------------------------------------
# end
# -------------------------------------------------
