##
# \file    cadabra2_defaults.py
# \ingroup pythoncore
# Cadabra2 pure Python functionality.
#
# This is a pure-python initialisation script to set the  path to
# sympy and setup printing of Cadabra expressions.  This script is
# called both by the command line interface 'cadabra2' as well as by
# the GUI backend server 'cadabra-server'.

import sys
import cadabra2
from cadabra2 import *
from importlib.machinery import PathFinder, ModuleSpec, SourceFileLoader
from importlib.abc import MetaPathFinder
from cdb_appdirs import user_config_dir, user_data_dir
import datetime
import atexit
import rlcompleter

__cdbkernel__=cadabra2.__cdbkernel__
__cdbkernel__.completer=rlcompleter.Completer()

import os
os.environ.setdefault('PATH', '')

PY3 = sys.version_info[0] == 3
if PY3:
   unicode = str

if "OFF" == "OFF":
   discr = "\\discretionary{}{}{} "
else:
   discr = ""
   
class PackageCompiler(MetaPathFinder):
        @classmethod
        def find_module(cls, fullname, path):
                return find_spec(fullname, path, None)

        @classmethod
        def find_spec(cls, fullname, path, target=None):
                #log_("Finding {}, path=[{}]".format(fullname, ', '.join(f'"{p}"' for p in path) if path is not None else ""))
                # Top-level import if path=None
                if path is None or path == "":
                        path = sys.path
                # Get unqualified package name
                if '.' in fullname:
                        parents = fullname.split('.')
                        name = parents.pop()
                else:
                        name = fullname
                        parents = []
                # Go through path and try to find a notebook.
                for entry in path:
                   have_cnb = os.path.isfile(os.path.join(entry, name + ".cnb"))
                   have_cdb = os.path.isfile(os.path.join(entry, name + ".cdb"))
                   have_ipynb = os.path.isfile(os.path.join(entry, name + ".ipynb"))
                   if have_cnb or have_cdb or have_ipynb:
                       # Notebook was found. Create a version of it in the temporary directory, then
                       # return a ModuleSpec object to allow Python to load that file
                       pkg_path = os.path.join(user_config_dir(), "cadabra_packages", *parents)
                       # Create the path if it doesn't exist
                       if not os.path.exists(pkg_path):
                               os.makedirs(pkg_path)
                       if have_cnb:                                             
                          compile_package__(os.path.join(entry, name + ".cnb"), os.path.join(pkg_path, name + ".py"))
                       elif have_cdb:                                             
                          compile_package__(os.path.join(entry, name + ".cdb"), os.path.join(pkg_path, name + ".py"))
                       else:                                             
                          compile_package__(os.path.join(entry, name + ".ipynb"), os.path.join(pkg_path, name + ".py"))
                       return ModuleSpec(
                                        fullname, 
                                        SourceFileLoader(fullname, os.path.join(pkg_path, name + ".py")), 
                                        origin=os.path.join(pkg_path, name + ".py"))
                
                # Return none if no notebook was found
                return None

# Prepend to sys.meta_path, so that all imports will first be checked in
# case they are notebooks that need compiling
sys.meta_path.insert(0, PackageCompiler)

# Add current directory to Python module import path.
sys.path.append(".")

#sys.path.insert(0,'/home/kasper/Development/git.others/sympy') 

# Attempt to import sympy; if not, setup logic so that the
# shell does not fail later.

try:
    import sympy
except:
    class Sympy:
        """!@brief Stub object for when Sympy itself is not available.
        
        @long When Sympy is not available, this object contains some basic
        functionality to prevent things from breaking elsewhere.
        """
        __version__="unavailable"

    sympy = Sympy()

if sympy.__version__ != "unavailable":
#    from sympy import factor
#    from sympy import integrate
#    from sympy import diff
    from sympy import symbols
    from sympy import latex
#    from sympy import sin, cos, tan, sqrt, trigsimp
    from sympy import Matrix as sMatrix

# Whether running in command-line mode or as client-server, there always
# needs to be a Server object known as 'server' through which interaction
# with the display routines is handled. The 'display' function will
# call the 'server.send' method.

if 'server' in globals():
    mopen="\\begin{dmath*}{}";
    mclose="\\end{dmath*}";
else:
    mopen=''
    mclose=''
    class Server:
        """!@brief Object to handle advanced display in a UI-independent way.

        @long Cadabra makes available to Python a Server object, which
        contains functions to send output to the user. When running
        from the command line this simply prints to the screen, but it
        can talk to a remote client to display images and maths.
        """
        
        def send(self, data, typestr, parent_id, cell_id, last_in_sequence):
            """ Send a message to the client; 'typestr' indicates the cell type,
            'parent_id', if non-null, indicates the serial number of the parent
            cell. If 'cell_id' is not 0, put the output in the cell with that id.
            """
            print(data)
            return 0

        def architecture(self):
            return "terminal"

        def test(self):
            print("hello there!")

        def handles(self, otype):
            if(otype=="plain"):
                return True
            return False

        def totals(self):
            return __cdb_progress_monitor__.totals()
            
    server = Server()

# Import matplotlib and setup functions to prepare its output
# for sending as base64 to the client. Example use:
#
#   import matplotlib.pyplot as plt
#   p = plt.plot([1,2,3],[1,2,5],'-o')
#   display(p[0])
#

have_matplotlib=True
try:
    import matplotlib
    import matplotlib.artist
    import matplotlib.figure
    matplotlib.use('Agg')
except ImportError:
    have_matplotlib=False

def save_history(history_path):
    try:
        readline.write_history_file(history_path)
    except:
        pass

try:
    import readline
    if not hasattr(readline, 'redisplay'):
      readline.redisplay = lambda: None
    history_path = os.path.join(user_data_dir(), "cadabra_history")
    if os.path.exists(history_path):
        readline.read_history_file(history_path)
        readline.set_history_length(1000)
    atexit.register(save_history, history_path)
except:
    pass

import io
import base64

## @brief Generic display function which handles local as well as remote clients.
#
# The 'display' function is a replacement for 'str', in the sense that
# it will generate human-readable output. However, in contrast to
# 'str', it knows about what the front-end ('server') can display, and
# will adapt the output to that. For instance, if
# server.handles('latex_view') is true, it will generate LaTeX output,
# while it will generate just plain text otherwise.
# 
# Once it has figured out which display is accepted by 'server', it
# will call server.send() with data depending on the object type it is
# being fed. Data types the server object can support are:
# 
# - "latex_view": text-mode LaTeX string.
# - "image_png":  base64 encoded png image.
# - "image_svg":  svg image.
# - "verbatim":   ascii string to be displayed verbatim.

def display(obj, cell_id=0, delay_send=False):
    """
    Generalised 'print' function which knows how to display objects in the 
    best possible way on the used interface, be it a console or graphical
    notebook. In particular, it knows how to display Cadabra expressions
    in typeset form whenever LaTeX functionality is available. Can also be
    used to display matplotlib plots.

    When using a Cadabra front-end (command line or notebook), an expression
    with a trailing semi-colon ';' will automatically be wrapped in a 
    'display' function call so that the expression is displayed immediately.
    """

    if type(obj)==list and len(obj)>0 and (isinstance(obj[0], matplotlib.figure.Figure) or isinstance(obj[0], matplotlib.figure.Artist)):
       # matplotlib has the annoying habit of returning plots as lists of figures.
       # print("list of figures", file=sys.stderr)
       return display(obj[0], cell_id, delay_send)
    
    if 'matplotlib' in sys.modules and isinstance(obj, matplotlib.figure.Figure):
        # print("figure 1", file=sys.stderr)
        imgstring = io.BytesIO()
        obj.savefig(imgstring,format='svg')
        imgstring.seek(0)
        b64 = base64.b64encode(imgstring.getvalue())
        return server.send(b64, "image_svg", 0, cell_id, False)
        # FIXME: Use the 'handles' query method on the Server object
        # to figure out whether it can do something useful
        # with a particular data type.

    elif 'matplotlib' in sys.modules and isinstance(obj, matplotlib.artist.Artist):
        # print("figure 2", file=sys.stderr)
        f = obj.get_figure()
        imgstring = io.BytesIO()
        f.savefig(imgstring,format='svg')
        imgstring.seek(0)
        b64 = base64.b64encode(imgstring.getvalue())
        return server.send(b64, "image_svg", 0, cell_id, False)

    elif hasattr(obj,'_backend'):
        # print("figure 3", file=sys.stderr)
        if hasattr(obj._backend,'fig'):
            f = obj._backend.fig
            imgstring = io.BytesIO()
            f.savefig(imgstring,format='svg')
            imgstring.seek(0)
            b64 = base64.b64encode(imgstring.getvalue())
            return server.send(b64, "image_svg", 0, cell_id, False)

    elif 'vtk' in sys.modules and isinstance(obj, vtk.vtkRenderer):
        # Vtk renderer, see http://nbviewer.ipython.org/urls/bitbucket.org/somada141/pyscience/raw/master/20140917_RayTracing/Material/PythonRayTracingEarthSun.ipynb
        pass
                    
#    elif isinstance(obj, numpy.ndarray):
#        server.send("\\begin{dmath*}{}"+str(obj.to_list())+"\\end{dmath*}", "latex")

    elif isinstance(obj, Ex):
        if server.handles('latex_view'):
            if delay_send:
                return obj._latex_()
            else:
                ret = mopen+obj._latex_()+mclose
                id=server.send(ret, "latex_view", 0, cell_id, False)
                # print(id)
                # Make a child cell of the above with input form content.
                cell_id = server.send(obj.input_form(), "input_form", id, cell_id, False)
                return cell_id
        else:
            server.send(unicode(obj), "plain", 0, cell_id, False)

    elif isinstance(obj, ExNode):
        if server.handles('latex_view'):
            if delay_send:
                return obj._latex_()
            else:
                ret = mopen+obj._latex_()+mclose
                id=server.send(ret, "latex_view", 0, cell_id, False)
                # print(id)
                # Make a child cell of the above with input form content.
                cell_id = server.send(obj.input_form(), "input_form", id, cell_id, False)
                return cell_id
        else:
            server.send(unicode(obj), "plain", 0, cell_id, False)
            return None

    elif isinstance(obj, LaTeXString):
        if server.handles('latex_view'):
            ret = mopen+obj._latex_()+mclose
            if delay_send:
                return ret
            else:
                cell_id = server.send(ret , "latex_view", 0, cell_id, False)
                return cell_id
        else:
            server.send(unicode(obj), "plain", 0, cell_id, False)
            return None
            
    elif isinstance(obj, Property):
        if server.handles('latex_view'):
            ret = mopen+obj._latex_()+mclose
            if delay_send:
                return ret
            else:
                return server.send(ret , "latex_view", 0, cell_id, False)
                # Not yet available.
                # server.send(obj.input_form(), "input_form", 0, False)
        else:
            server.send(unicode(obj), "plain", 0, cell_id, False)
            return None
            
    elif type(obj)==list:
        if server.handles('latex_view'):
           # print(f"list {obj}", file=sys.stderr)
           atleast_one_unstolen=False
           cdb_out="$\\big[$"
           first=True
           for elm in obj:
               if first==False:
                   cdb_out+=","+discr
               else:
                   first=False
               nxt = display(elm, cell_id, delay_send=True)
               if nxt != None:
                  cdb_out+= "$"+str(display(elm, cell_id=cell_id, delay_send=True))+"$"
                  atleast_one_unstolen=True
           cdb_out+="$\\big]$";
           if delay_send:
              return cdb_out
           if atleast_one_unstolen:
              return server.send(cdb_out, "latex_view", 0, cell_id, False)
        else:
           cdb_out = "["
           first=True           
           for elm in obj:
               if not first:
                   cdb_out+=", "
               else:
                   first=False                  
               cdb_out += str(elm)
           cdb_out += "]"
           server.send(unicode(cdb_out), "plain", 0, cell_id, False)
           return None
           # FIXME: send input_form version.
        
    elif hasattr(obj, "__module__") and hasattr(obj.__module__, "find") and obj.__module__.find("sympy")!=-1:
        if delay_send:
           return latex(obj)
        else:
           if server.handles('latex_view'):
               return server.send("\\begin{dmath*}{}"+latex(obj)+"\\end{dmath*}", "latex_view", 0, cell_id, False)
           else:
               server.send(latex(obj), "plain", 0, cell_id, False)
               return None
        
    else:
        # Failing all else, just dump a str representation to the notebook, asking
        # it to display this verbatim.
        # server.send("\\begin{dmath*}{}"+str(obj)+"\\end{dmath*}", "latex")
        if delay_send:
            if server.handles('latex_view'):
               return "\\verb|"+str(obj)+"|"
            else:
               return str(obj)
        else:
            if server.handles('latex_view'):
               return server.send(unicode(obj), "verbatim", 0, cell_id, False)
            else:
               server.send(unicode(obj), "plain", 0, cell_id, False)
               return None
    
__cdbkernel__.server=server
__cdbkernel__.display=display

class Console(object):
        """
        The interactive console works in the same Python context as 
        the notebook cells to allow evaluation of expressions for 
        debugging/logging purposes
        """
        def log(self, *objs):
            """Sends a string representation of objs to the console"""
            if server.architecture() == "terminal":
                print(*objs)
            elif server.architecture() == "client-server":
                server.send(" ".join(str(obj) for obj in objs), "csl_out", 0, cell_id, False)

        def clear(self):
            """Clears the output of the console window"""
            if server.architecture() == "client-server":
                server.send("", "csl_clear", 0, cell_id, False)

        def warn(self, msg):
            if server.architecture() == "terminal":
                print("Warning: " + msg)
            elif server.architecture() == "client-server":
                server.send(msg, "csl_warn", 0, cell_id, False)

console = Console()
__cdbkernel__.configure_warnings(callback=console.warn)
    
# Set display hooks to catch certain objects and print them
# differently. Should probably eventually be done cleaner.

def _displayhook(arg):
    global remember_display_hook
    if isinstance(arg, Ex):
        print(unicode(arg))
    elif isinstance(arg, Property):
        print(unicode(arg))
    else:
        remember_display_hook(arg)

remember_display_hook = sys.displayhook
sys.displayhook = _displayhook

# Default post-processing algorithms. These are not pre-processed
# so need to have the '__cdbkernel__' argument.

def post_process(__cdbkernel__, ex):
    collect_terms(ex)