#!/usr/bin/env python
# coding: utf-8

"""Interact with matplotlib and other plotters

"""

import os.path
import datetime
now = datetime.datetime.now
import logging
import subprocess
import sys

import numpy
import matplotlib
import matplotlib.pyplot
from . import config
from . import io
from . import meta
from . import tools

def plotdir():
    """Returns todays plotdir.

    Configuration 'plotdir' must be set.  Value is expanded with strftime.
    """
    return datetime.date.today().strftime(config.get_config('plotdir'))

def print_or_show(fig, show, outfile, in_plotdir=True, tikz=None,
                  data=None, store_meta="", close=True):
    """Either print or save figure, or both, depending on arguments.

    Taking a figure, show and/or save figure in the default directory,
    obtained with :func:plotdir.  Creates plot directory if needed.

    :param fig: Figure to store.  
    :type fig: matplotlib.Figure object
    :param show: Show figure or not
    :type show: boolean
    :param outfile: File to write figure to, or list of files.  If the
        string ends in a '.', write to x.png and x.pdf.
    :type outfile: string or list of strings
    :param in_plotdir: If true, write to default plot directory.  If
        false, write to currect directory or use absolute path.
    :type in_plotdir: boolean
    :param tikz: Try to write tikz code with matplotlib2tikz.  Requires
        that the latter is installed.
    :type tikz: boolean
    :param data: Store associated data in .dat file (useful for pgfplots).
        May be a list of ndarrays, which results in multiple numbered datafiles.
    :type data: ndarray or list thereof
    :param store_meta: Also store other info.  This is a string that will
        be written to a file.  If not set or set to None, it will just
        write the pyatmlab version.  The file will use the same basename
        as the outfile, but replacing the extention by "info".  However,
        this only works if outfile is a string and not a list thereof.
        To write nothing, pass an empty string.
    :type store_meta: str.
    :param close: If true, close figure.  Defaults to true.
    :type close: bool.
    """

    if outfile is not None:
        outfiles = [outfile] if isinstance(outfile, str) else outfile
        if isinstance(outfile, str):
            if outfile.endswith("."):
                outfiles = [outfile+ext for ext in ("png", "pdf")]
                infofile = outfile + "info"
            else:
                outfiles = [outfile]
                infofile = None

        # interpret as sequence
        for outf in outfiles:
            if in_plotdir and not '/' in outf:
                outf = os.path.join(plotdir(), outf)
            logging.info("Writing to file: {}".format(outf))
            if not os.path.exists(os.path.dirname(outf)):
                os.makedirs(os.path.dirname(outf))
            fig.canvas.print_figure(outf)
        if close:
            matplotlib.pyplot.close(fig)
        pr = subprocess.run(["pip", "freeze"], stdout=subprocess.PIPE) 
        info = " ".join(sys.argv) + "\n" + pr.stdout.decode("utf-8") + "\n"
        info += tools.get_verbose_stack_description()
        if infofile is not None and info:
            if in_plotdir and not "/" in infofile:
                infofile = os.path.join(plotdir(), infofile)
            with open(infofile, "w", encoding="utf-8") as fp:
                fp.write(info)
    if show:
        matplotlib.pyplot.show()

    if tikz is not None:
        import matplotlib2tikz
        print(now(), "Writing also to:", os.path.join(plotdir(), tikz))
        matplotlib2tikz.save(os.path.join(plotdir(), tikz))
    if data is not None:
        if not os.path.exists(io.plotdatadir()):
            os.makedirs(io.plotdatadir())
        if isinstance(data, numpy.ndarray):
            data = (data,)
        # now take it as a loop
        for (i, dat) in enumerate(data):
            outf = os.path.join(io.plotdatadir(),
                "{:s}{:d}.dat".format(
                    os.path.splitext(outfiles[0])[0], i))
            fmt = ("%d" if issubclass(dat.dtype.type, numpy.integer) else
                    '%.18e')
            if len(dat.shape) < 3:
                numpy.savetxt(outf, dat, fmt=fmt)
            elif len(dat.shape) == 3:
                io.savetxt_3d(outf, dat, fmt=fmt)
            else:
                raise ValueError("Cannot write {:d}-dim ndarray to textfile".format(
                    len(dat.shape)))


@@ -187,16 +187,30 @@
 port sys
+%0Aimport pickle
 %0A%0Aimport
@@ -2486,16 +2486,58 @@
  %22info%22%0A
+                figfile = outfile + %22pkl%22%0A
         
@@ -2615,25 +2615,24 @@
  = None%0A
-%0A
         
 # interp
@@ -2627,432 +2627,31 @@
     
-# interpret as sequence%0A        for outf in outfiles:%0A            if in_plotdir and not '/' in outf:%0A                outf = os.path.join(plotdir(), outf)%0A            logging.info(%22Writing to file: %7B%7D%22.format(outf))%0A            if not os.path.exists(os.path.dirname(outf)):%0A                os.makedirs(os.path.dirname(outf))%0A            fig.canvas.print_figure(outf)%0A        if close:%0A            matplotlib.pyplot.close(fig)
+        figfile = None%0A
 %0A   
@@ -3003,16 +3003,82 @@
 fofile)%0A
+            logging.info(%22Writing info to %7B:s%7D%22.format(infofile))%0A
         
@@ -3170,48 +3170,782 @@
     
-if show:%0A        matplotlib.pyplot.show(
+    if figfile is not None:%0A            if in_plotdir and not %22/%22 in figfile:%0A                figfile = os.path.join(plotdir(), figfile)%0A            logging.info(%22Writing figure object to %7B:s%7D%22.format(figfile))%0A            with open(figfile, %22wb%22) as fp:%0A                pickle.dump(fig, fp, protocol=4)%0A        # interpret as sequence%0A        for outf in outfiles:%0A            if in_plotdir and not '/' in outf:%0A                outf = os.path.join(plotdir(), outf)%0A            logging.info(%22Writing to file: %7B%7D%22.format(outf))%0A            if not os.path.exists(os.path.dirname(outf)):%0A                os.makedirs(os.path.dirname(outf))%0A            fig.canvas.print_figure(outf)%0A    if show:%0A        matplotlib.pyplot.show()%0A%0A    if close:%0A        matplotlib.pyplot.close(fig
 )%0A%0A 

Write out figure object and close figure at the end