#! /usr/bin/env python3
################################################################
import re as myre
import PythonLatex.latex_structure as tex
import types
import sympy as sym
import hashlib
import dill
################################################################
from sympy.printing.latex import LatexPrinter
from sympy import latex
from sympy.core.function import Derivative
################################################################


class CustomStrPrinter(LatexPrinter):

    def nbIdenticalDx(self, expr):
        first_dx = expr.args[1:][0]
        count = 0
        for dx in expr.args[1:]:
            if dx == first_dx:
                count += 1
        if count == len(expr.args[1:]):
            return count, first_dx

        return 0, None

    def _print_Derivative(self, expr):
        # print expr
        # print LatexPrinter._print_Derivative(self,expr)

        count, dx = self.nbIdenticalDx(expr)
        f = expr.args[0]
        if len(f.args) == 0:
            return LatexPrinter._print_Derivative(self, expr)

        dxf = f.args[0]
        f_nargs = len(f.args)
        # case where ' applies
        if f_nargs == 1 and dx == dxf:
            return (latex(expr.args[0].func) +
                    "'" * len(expr.args[1:]) +
                    "(" + latex(dxf) + ")")

        return LatexPrinter._print_Derivative(self, expr)

    def _print_Subs(self, expr):
        if not isinstance(expr.args[0], Derivative):
            return LatexPrinter._print_Subs(self, expr)

        # print expr
        deriv = expr.args[0]
        _dxs = deriv.args[1:]
        # print _dxs
        p1 = expr.args[1]
        p2 = expr.args[2]
        # print len(_dxs)

        _subst = {}
        _subst2 = {}
        for i in range(0, len(p1)):
            _subst[p1[i]] = p2[i]
            if hasattr(p2[i], 'args'):
                _p2 = type(p2[i])
            _subst2[p1[i]] = _p2

        _dxs = [e.subs(_subst2) for e in _dxs]

        new_f = deriv.args[0].subs(_subst)
        # print 'CCCCCCCCC'
        # print new_f
        # print _dxs
        # print _subst

        tmp = Derivative(new_f, *_dxs)
        # print tmp
        return latex(tmp)
        # return LatexPrinter._print_Subs(self,expr)

################################################################


class PyCode(object):
    """
    """

    def checkLine(self, line):

        m = myre.match(r"\\end{python}", line)
        if m:
            return Block()

        return Block.checkLine(self, line)

    def checkValidity(self):
        for c in self.content:
            if isinstance(c, types.InstanceType):
                raise Exception("python block cannot have subblocks")

    def __str__(self):
        return ""

    def generateHashSource(self):
        txt_content = self.content.encode()
        self.hash_source = hashlib.sha1(txt_content).hexdigest()
        gl = globals()
        self.global_keys_before = gl.keys()
        self.local_keys_before = self.__dict__.keys()

    def needReexecute(self):
        if self.alias == 'header':
            return True
        if 'loaded_hash_source' in self.__dict__:
            return not self.loaded_hash_source == self.hash_source
        return True

    def saveBlockOutputInfo(self):
        if self.alias == 'header':
            return
        global_keys = globals().keys()
        global_keys_to_save = set(global_keys) - set(self.global_keys_before)
        global_keys_to_save = list(global_keys_to_save)
        global_entries_to_save = dict(
            [(e, globals()[e]) for e in global_keys_to_save])

        local_keys = self.__dict__.keys()
        local_keys_to_save = set(local_keys) - set(self.local_keys_before)
        local_keys_to_save = list(local_keys_to_save)
        local_entries_to_save = dict(
            [(e, self.__dict__[e]) for e in local_keys_to_save])

        try:
            for k in global_entries_to_save:
                try:
                    print('try to save ', k,
                          ' ', latex(global_entries_to_save[k]))
                    global_entries_to_save[k] = dill.dumps(
                        global_entries_to_save[k])
                    print('saved ', k)
                except Exception as e:
                    print(e)

            self.global_block_result = dill.dumps(
                global_entries_to_save)
            self.local_block_result = dill.dumps(local_entries_to_save)
            self.global_hash_result = hashlib.sha1(
                self.global_block_result).hexdigest()
            self.local_hash_result = hashlib.sha1(
                self.local_block_result).hexdigest()
        except Exception as e:
            print("cannot save block: {0}\n{1}".format(self.alias, e))
            self.global_block_result = None
            self.global_hash_result = None
            self.local_block_result = None
            self.local_hash_result = None
            self.hash_source = None

        filename = self.alias + '.pytex.out'
        f = open(filename, 'w')
        # print "save symbols {0}: hash_source {1}".format(
        # self.alias,self.hash_source)
        raw_save = [self.hash_source, self.global_hash_result,
                    self.global_block_result,
                    self.local_hash_result,
                    self.local_block_result]
        dill.dump(raw_save, f)
        f.close()

    def loadBlockOutputInfo(self):
        return
        if self.alias == 'header':
            return
        filename = self.alias + '.pytex.out'
        try:
            f = open(filename, 'r')
            self.loaded_hash_source, self.global_hash_result,
            self.global_block_result, self.local_hash_result,
            self.local_block_result = dill.load(f)
            # print "loaded symbols"
            # print "AAAAAAAAAAA " + str(self.loaded_hash_source)
            globals().update(dill.loads(self.global_block_result))
            self.__dict__.update(dill.loads(self.local_block_result))
            # print pickle.loads(self.global_block_result).keys()
            # print self.__dict__.keys()
            f.close()
            print('loaded {0}'.format(filename))
        except Exception:
            # print "Could not load block " + self.alias
            # print e
            pass

    def evalBlock(self):
        txt_content = self.content
        self.generateHashSource()
        self.loadBlockOutputInfo()
        need_reexecute = self.needReexecute()
        if need_reexecute or reexec:
            print("execute block: {0}".format(self.alias))
            gl = globals()
            gl.update({'self': self})
            try:
                exec(txt_content, gl)
            except Exception as e:
                mesg = """{0}:0: For block '{1}' there was an execution problem
{2}
Block was:
************
{3}
""".format(parsed_latex_file, self.alias, e, txt_content)
                print(mesg)

                if not continue_flag:
                    raise RuntimeError('have to stop')

            del gl['self']
            # self.saveBlockOutputInfo()

    def getContent(self, varname):
        try:
            c = self.__dict__[varname]
        except Exception as ex:
            print(ex)
            raise Exception("cannot retreive variable " +
                            varname + " from block " + self.alias)

        if (myre.search('sympy', str(type(c)))):
            c = sym.latex(c)
        return c

    def __init__(self, name, alias, content):

        self.name = name
        self.alias = alias
        self.content = content
        if not self.name == 'python':
            raise Exception(
                'cannot create a PyCode block from block type ' +
                self.name)
        if self.alias is None:
            raise Exception("not valid python block")

        self.block_eval = None


################################################################

def evaluatePythonBlock(b):
    if isinstance(b, tex.LatexEnvironment) and b.name == 'python':
        b.hide = True
        b = PyCode(b.name, b.option, str(b.toks[1]))

        globals()[b.alias] = b
        b.evalBlock()


def replaceAliases(b, text):

    if isinstance(b, tex.LatexEnvironment) and b.name == 'python':
        return text

    if type(text) == str:
        regex = r'\\py{(.*?)}'
        m = myre.findall(regex, text)
        if not m:
            return text
        # print text
        for expression in m:
            # print('replaceAliasses: ', expression)
            try:
                d = dict()
                exec("ans = " + expression, globals(), d)
            except Exception as e:
                print("Block " + str(b.name) +
                      " Cannot parse expression '" +
                      str(expression) +
                      "'\n" + str(e))
                if continue_flag:
                    continue
                raise(e)

            ans = d['ans']
            rep = '\\py{' + expression + '}'

            try:
                ans_str = str(latex(ans))

            except Exception as e:
                print("AAAAAAAA " + str(e))
                return text

            # print "ans is now " + str(ans)
            globals().update({'ans': ans})
            # print "replace " + rep + " with " + cont

            text = text.replace(rep, ans_str, 1)
        return text

################################################################


def collectScript(b, text):
    if isinstance(b, tex.LatexEnvironment) and b.name == 'python':
        # print('AAAA ', b)
        globals()['total_python_script'] += "{0} = CodeEval()\n".format(
            b.option)
        globals()['total_python_script'] += text.replace(
            'self', b.option)
        return text

    if type(text) == str:
        # print text
        regex = r'\\py{(.*?)}'
        m = myre.findall(regex, text)
        if not m:
            return text
        # print m
        for expression in m:
            globals()['total_python_script'] += (
                "print('printing expression: {0}')\n".format(
                    str(expression).replace("'", "\\'")))
            globals()['total_python_script'] += (
                "ans = " + expression + "\n")
            globals()['total_python_script'] += (
                "print(str(latex(ans)) + '\\n\\n')\n")
        return text

################################################################


continue_flag = False
reexec = True


def interpretPython(filename, c_flag, reexec_flag):

    continue_flag = c_flag
    reexec = reexec_flag
    tex_struct = tex.LatexStructure()

    # tex_struct.buildLatexBlocks(filename)
    tex_struct.parseLatexFile(filename)

    total_python_script = """
from PythonLatex.generate_latex import *
class CodeEval:
    pass

"""
    globals().update({'total_python_script': total_python_script})
    globals().update({'reexec': reexec})
    globals().update({'continue_flag': continue_flag})
    globals().update({'parsed_latex_file': filename})

    tex_struct.pathInBlock(text_functor=collectScript)

    total_python_script = globals()['total_python_script']
    tmp_file = open('python_script.py', 'w')
    tmp_file.write(total_python_script)
    tmp_file.close()

    tex_struct.pathInBlock(begin_functor=evaluatePythonBlock,
                           text_functor=replaceAliases)


    # pyblocks = tex_struct.getBlocksFromType("python")
    #
    # for b in pyblocks:
    #     total_python_script += b.content[0]

    return str(tex_struct)