'''
Created on 12-Jul-2021

@author: bthakur

'''
import inspect
import traceback

import random
from xv.km.util import latex_fraction, latex_color
from xv.util import toIntOrFloat

from sympy import (Rational, sqrt, symbols, Symbol, limit, latex, Function,
                   Limit, oo, Derivative, Integral,
                   simplify,
                   exp, log, expand_log,
                   Mul, Add, Pow,
                   Sum, binomial, factorial
                   )

from sympy import UnevaluatedExpr as ueval

from sympy import (acos, acot, acoth, acsc, asec, asin, 
                   asinh, atan, atan2, atanh, 
                   cos, cosh, cot, coth, csc, sec, 
                   sin, sinc, sinh, tan, tanh)



# _derivative_solver
# DerivativeSolver

from xv.util import UtilClass

class DerivativeSolver(UtilClass):
    
    def __init__(self, expr, x, var_symbs, **kwargs):
        super().__init__(**kwargs)
        self.expr = expr
        self.x = x
        self.var_symbs = var_symbs
        
        self._soln = None
        self._soln_stmts = None
        self._r_dict = None
        self.r_to_r_dict = None
        self._r_exprs = None
        pass

    def _init_soln_data(self):
        self._soln = []
        self._soln_stmts = []

        self._r_dict = {}
        self.r_to_r_dict = {}
        self._r_exprs = {}
        self._soln_result = []
        pass
    
    def get_soln(self):
        #init data 
        self._init_soln_data()
        
        # solve derivative
        self.get_sub_soln()
        #self._reset_r_to_r_dict()
        
        soln_result = list(reversed(self._soln_result))
        return soln_result
        pass
    
    def get_reverse_soln(self):
        #init data 
        self._init_soln_data()
        
        # solve derivative
        self.get_sub_soln()
        #self._reset_r_to_r_dict()
        
        # return result
        #reverse_soln = self._soln_stmts
        reverse_soln = self._soln_result
        return reverse_soln
        pass

    def _compute_org_ueval(self, expr, x):
        ltx = self._modified_latex
        
        try:
            d_dx = latex_fraction("d", f"d{latex(x)}")

            cls = type(expr)
            org_ueval_args = [ueval(org_ueval_arg) for org_ueval_arg in expr.args]
            org_ueval_expr = cls(*org_ueval_args)
            org_ueval_expr_dx_latex = f"{d_dx} \\left( {ltx(org_ueval_expr, mul_symbol='dot')} \\right)"
            return (org_ueval_args, org_ueval_expr, org_ueval_expr_dx_latex)
        
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            expr = {expr}
            expr.args = {expr.args}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            return None
        
        pass

    def _update_r_symb(self, arg_idx, arg, x, r_args, temp_r_dict):
        r_symb = None
        for _r_dict_k, _r_dict_v in self._r_dict.items():
            if _r_dict_v == arg:
                r_symb = _r_dict_k
                break
            
        if r_symb is None:
            r_symb = Function(f"R_{{{len(self._r_dict)}}}")(x)
        temp_r_dict[r_symb] = arg
        self._r_dict[r_symb] = arg
        
        r_args[arg_idx] = r_symb
        return r_symb
        pass

    def get_sub_soln(self, expr = None):
        ltx = self._modified_latex
        
        x = self.x
        var_symbs = self.var_symbs
        if expr is None:
            expr = self.expr

        r_args = [arg for arg in expr.args]
        mod_args = [arg for arg in expr.args]


        expr_dx = Derivative(ueval(expr), x)
        #expr_dx_doit = expr_dx.doit()

        d_dx = latex_fraction("d", f"d{latex(x)}")
        expr_dx_latex = f"{d_dx} \\left( {latex(expr, mul_symbol='dot')} \\right)"


        cls = type(expr)
        org_vars = self._compute_org_ueval(expr, x)
        if org_vars is None:
            return None
        org_ueval_args, org_ueval_expr, org_ueval_expr_dx_latex = org_vars


        stmt = f"""
        Solving expr $\\displaystyle 
            = {expr_dx_latex} = {org_ueval_expr_dx_latex}
            $ 
        """
        self._soln_stmts.append(stmt.strip())
        
        
        temp_r_dict = {}
        for arg_idx, arg in enumerate(expr.args):
            if not hasattr(arg, 'free_symbols'):
                stmt = f"""
                arg $\\displaystyle 
                    = {latex(arg, mul_symbol='dot')}
                    $ has no free_symbols
                """
                self._soln_stmts.append(stmt.strip())
                continue
                
            if not set(arg.free_symbols).intersection(set(var_symbs)):
                stmt = f"""
                arg $\\displaystyle 
                    = {latex(arg, mul_symbol='dot')}
                    $ has no free_symbols
                """
                self._soln_stmts.append(stmt.strip())
                continue
            
            if arg == x:
                stmt = f"""
                expr and x are same $\\displaystyle 
                    = {latex(arg)} = {latex(x)}
                    $ 
                """
                self._soln_stmts.append(stmt.strip())
                continue
                    
            if isinstance(arg, Function):
                if (arg in var_symbs or 
                    hasattr(arg, 'name')):
                    stmt = f"""
                    arg $\\displaystyle 
                        = {latex(arg)} 
                        $ is a function and in {var_symbs} or has name attr
                    """
                    self._soln_stmts.append(stmt.strip())
                    continue
                    
            
            stmt = f"""
            arg $\\displaystyle 
                = {latex(arg, mul_symbol='dot')}
                $ has free_symbols = {arg.free_symbols}
            """
            self._soln_stmts.append(stmt.strip())
            
            #Add new var to dicts or reuse existing one
            r_symb = self._update_r_symb(arg_idx, arg, x, r_args, temp_r_dict)

            stmt = f"""
            Added var $\\displaystyle 
                {latex(r_symb)}
                = {latex(temp_r_dict[r_symb], mul_symbol='dot')}
                $ 
            """
            #self._soln_stmts.append(stmt.strip())
            
            
            stmt = f"""
            Calling self.get_sub_soln $\\displaystyle 
                \\left( {latex(expr, mul_symbol='dot')} \\right)
                $ 
            """
            #self._soln_stmts.append(stmt.strip())
            
            mod_arg = self.get_sub_soln(arg)
            mod_args[arg_idx] = ueval(mod_arg)

            #self._r_dict[Derivative(r_symb, x)] = Derivative(mod_arg, x)
            
            pass


        rev_temp_r_dict = {v:ueval(k) for k,v in temp_r_dict.items()}
        expr_subs = expr.subs(rev_temp_r_dict)

        
        stmt = f"""
        Rewriting expr as $\\displaystyle 
            = {latex(expr, mul_symbol='dot')}
            = {latex(expr_subs, mul_symbol='dot')}
            $ 
            using dict
        """
        #self._soln_stmts.append(stmt.strip())

        dict_stmts = self._print_dict(rev_temp_r_dict, 'rev_temp_r_dict')
        if dict_stmts:
            self._soln_stmts.extend(dict_stmts)
            
        dict_stmts = self._print_dict(self._r_dict, 'self._r_dict')
        if dict_stmts:
            self._soln_stmts.extend(dict_stmts)
            
       
        cls = type(expr)
        mod_args = [ueval(mod_arg) for mod_arg in mod_args]
        mod_expr = cls(*mod_args)
        
        mod_expr_dx = Derivative(mod_expr, x)
        mod_expr_dx_doit = mod_expr_dx.doit()
        mod_expr_dx_latex = f"{d_dx} \\left( {ltx(mod_expr, mul_symbol='dot')} \\right)"
        
        r_expr = cls(*r_args)
        r_expr_dx = Derivative(r_expr, x)
        r_expr_dx_doit = r_expr_dx.doit()
        r_expr_dx_latex = f"{d_dx} \\left( {latex(r_expr, mul_symbol='dot')} \\right)"
        r_expr_dx_doit_latex = f"{latex(r_expr_dx_doit, mul_symbol='dot')}"
        
        
        
        
        ueval_r_dict = {k:ueval(v) for k,v in self._r_dict.items()}
        r_expr_dx_doit_subs = r_expr_dx_doit.subs(ueval_r_dict, evaluate = False)
        r_expr_dx_doit_subs_latex = f"{ltx(r_expr_dx_doit_subs, mul_symbol='dot')}"

        try:
            r_expr_dx_doit_mod = r_expr_dx_doit.copy()
            for symb, symb_v in self._r_dict.items():
                t_dict = {Derivative(symb, x): Derivative(symb_v, x)}
                
                r_expr_dx_doit_mod = r_expr_dx_doit_mod.subs(t_dict)
    
                t_symb_dict = {symb: symb_v}
                r_expr_dx_doit_mod = r_expr_dx_doit_mod.subs(t_symb_dict)
                
                pass
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            r_expr_dx_doit = {r_expr_dx_doit}
            type(r_expr_dx_doit).__name__ = {type(r_expr_dx_doit).__name__}
            """
            #self._log(sError)
            #self._log(traceback.format_exc())
            
            r_expr_dx_doit_mod = r_expr_dx_doit
            pass
            
        
        soln_result_item = {
            'x': x,
            'expr': expr,
            'cls': type(expr).__name__,
            'expr_subs': expr_subs, # subs expr with reverse temp_r_dict
            'org_ueval_args': org_ueval_args,
            'r_args': r_args,
            'mod_args': mod_args,
            'temp_r_dict': temp_r_dict,
            'r_dict': self._r_dict.copy(),
            }
        self._soln_result.append(soln_result_item)



        
        stmt = f"""
        Rewriting expr $\\displaystyle 
            = {org_ueval_expr_dx_latex} 
            = {r_expr_dx_doit_latex}
            = {r_expr_dx_doit_subs_latex}
            = {mod_expr_dx_latex}
            = {latex(mod_expr_dx_doit, mul_symbol='dot')}
            $ 
        """
        self._soln_stmts.append(stmt.strip())


        stmt = f"""
        Rewriting expr $\\displaystyle 
            = {expr_dx_latex} 
            = {org_ueval_expr_dx_latex} 
            = {r_expr_dx_doit_latex}
            = {r_expr_dx_doit_subs_latex}
            = {mod_expr_dx_latex}
            = {latex(mod_expr_dx_doit, mul_symbol='dot')}
            $ 
        """
        #self._soln_stmts.append(stmt.strip())
        
        return mod_expr
        pass
    
    
    def _print_dict(self, data, name):
        dict_stmts = []
        stmt = f"""
        {name}:
        """
        dict_stmts.append(stmt.strip())
        for k,v in data.items():
            stmt = f"""
            $\\displaystyle 
                {latex(k, mul_symbol='dot')}
                = {latex(v, mul_symbol='dot')}
                $ 
                
            """
            dict_stmts.append(stmt.strip())
            pass
        return dict_stmts
        pass
    
    def _reset_r_to_r_dict(self):
        ltx = self._modified_latex
        
        self.r_to_r_dict = {}
        expr = self._soln_result[-1]['expr']
        r_dict = self._soln_result[-1]['r_dict']
        x = self._soln_result[-1]['x']
        
        if not r_dict:
            return

        keys = list(r_dict.keys())
        keys.sort(key = lambda k : f"{k}", reverse = False)
        for k in keys:
            v = r_dict[k]
            print(f"k = {k} = {v} ")
            
            if not self.r_to_r_dict:
                self.r_to_r_dict[k] = v
            else:
                for rk in self.r_to_r_dict.keys():
                    self.r_to_r_dict[rk] = self.r_to_r_dict[rk].subs(v, k)
                    pass
                self.r_to_r_dict[k] = v
            pass
        
        print(f"""
        self.r_to_r_dict = {self.r_to_r_dict}
        expr = {expr}
        """)
        
        keys = list(self.r_to_r_dict.keys())
        keys.sort(key = lambda k : f"{k}", reverse = True)
        for k in keys:
            
            v = self.r_to_r_dict[k]
            expr = expr.subs(v, k)
            d_expr_dx_doit = Derivative(expr, x).doit()
            print(f"""
            k = {k} = {v} 
            expr = {expr}
            d_expr_dx_doit = {d_expr_dx_doit}
            """)
            
        pass
    
    def convert_soln_to_stmts(self, soln_result):
        ltx = self._modified_latex
        
        # reset _r_exprs
        self._r_exprs = {}

        soln_result_stmts = []
        
        if self.verbose:
            soln_result_stmts.append("""<p style="color:red">Start of soln_result</p>""")
        
        
        # print r_dict
        r_dict = soln_result[-1]['r_dict']
        if self.verbose and r_dict:
            r_dict_stmts = self._print_r_dict(r_dict)
            if r_dict_stmts:
                soln_result_stmts.extend(r_dict_stmts)
                pass
            
        
        #for sub_soln in reversed(soln_result[-2:]):
        for sub_soln_idx, sub_soln in enumerate(soln_result):

            soln_items = self._convert_sub_soln_to_stmts(soln_result, sub_soln_idx)
            
            if not soln_items:
                continue
            
            soln_result_stmts.extend(soln_items)
            pass
        
        if self.verbose:
            soln_result_stmts.append("""<p style="color:red">End of soln_result</p>""")
        
        return soln_result_stmts
        pass
    
    def _convert_sub_soln_to_stmts(self, soln_result, sub_soln_idx):
        ltx = self._modified_latex
        
        sub_soln = soln_result[sub_soln_idx]
        
        x = self.x
        dx = Symbol(f"d{latex(x)}")
        d_dx = latex_fraction("d", f"d{latex(x)}")
                
        soln_items = []
        if self.verbose:
            soln_items.append("""<p style="color:blue">Start of sub_soln</p>""")

        expr = sub_soln['expr']
        org_ueval_args = sub_soln['org_ueval_args']
        mod_args = sub_soln['mod_args']
        r_args = sub_soln['r_args']
        r_dict = sub_soln['r_dict']
        temp_r_dict = sub_soln['temp_r_dict']
        
        r_dict_expr_key = self._get_r_dict_expr_key(r_dict, expr)
        if r_dict_expr_key: 
            dR_latex = f"{d_dx} {r_dict_expr_key.name}"
        else:
            dR_latex = None
            

        # print sub_soln dict
        sub_soln_items = self._print_sub_soln_items(sub_soln)
        if sub_soln_items:
            #soln_items.extend(sub_soln_items)
            pass

        org_ueval_expr = type(expr)(*org_ueval_args)
        stmt0 = f"""
        {d_dx} \\left( {ltx(org_ueval_expr, mul_symbol='dot')} \\right)
        """

        
        r_expr = type(expr)(*r_args)
        r_expr_der = Derivative(r_expr, x)
        r_expr_der_doit = r_expr_der.doit()
        
        
        stmt_r_dict = []
        if temp_r_dict:
            temp_r_dict_latex = self._get_dict_as_latex_list(temp_r_dict)
            temp_r_dict_latex = ",\\;".join(temp_r_dict_latex)
            
            keys = list(temp_r_dict.keys())
            r_expr_der_latex = ltx(r_expr_der, mul_symbol='dot', funct_symbs = keys)
            r_expr_der_doit_latex = ltx(r_expr_der_doit, mul_symbol='dot', funct_symbs = keys)
            stmt = f"""
                \\text{{Let: }} {temp_r_dict_latex} \\\\ 
                
                \\\\ 
                 
                \\text{{Then, }} {r_expr_der_latex}
                = {r_expr_der_doit_latex}  \\\\ 
            """
            stmt_r_dict.append(stmt)
            
            if sub_soln_idx == 0:
                self._r_exprs[f"expr{len(self._r_exprs)}"] = r_expr_der
            else:
                _subs_r_expr = self._r_exprs['expr0'].subs(temp_r_dict)
                self._r_exprs[f"expr{len(self._r_exprs)}"] = _subs_r_expr
            pass

        stmt_r_dict = "\n\n".join(stmt_r_dict)
        stmt_r_dict = f"""
        \\begin{{cases}}
        {stmt_r_dict}
        \\end{{cases}}
        """
        
        from itertools import cycle
        colors = cycle(['red', 'green', 'blue'])
        
        r_expr_der_doit_latex = ltx(r_expr_der_doit, mul_symbol='dot')
        for k,v in r_dict.items():
            for v_pow in [1, 2, 3]:
                v_latex = f""" \\left( {ltx(v ** v_pow, mul_symbol='dot')} \\right) """
                v_latex = latex_color(v_latex, next(colors))
                r_expr_der_doit_latex = r_expr_der_doit_latex.replace(ltx(k ** v_pow, mul_symbol='dot'), 
                                                            v_latex)
            pass
        
        stmt_r_expr = f"""
            {r_expr_der_doit_latex}
        """

         
        mod_expr = type(expr)(*mod_args)
        stmt_df_dx = f"""
            {ltx(Derivative(mod_expr, x).doit(), mul_symbol='dot')}
        """
        
        if sub_soln_idx == 0:
            first_sub_sol_stmts = self._get_first_sub_sol_stmts(
                stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict, temp_r_dict)
            soln_items.extend(first_sub_sol_stmts)
        else:
            non_first_sub_sol_stmts = self._get_non_first_sub_sol_stmts(
                stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict, temp_r_dict, 
                sub_soln_idx, dR_latex)
            soln_items.extend(non_first_sub_sol_stmts)
            pass
        
        
        if self.verbose:
            soln_items.append("""<p style="color:blue">End of sub_soln</p>""")
        
        return soln_items
        pass
    
    def _get_r_dict_expr_key(self, r_dict, expr):
        for k, v in r_dict.items():
            if v == expr:
                return k
        return None
        pass

    def _get_first_sub_sol_stmts(self, 
                stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict,  temp_r_dict):
        
        if not temp_r_dict:
            return self._get_first_sub_sol_stmts_no_temp_r_dict(
                        stmt0, stmt_r_expr, stmt_df_dx)
        
        return self._get_first_sub_sol_stmts_with_temp_r_dict(
                        stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict)
        pass

    def _get_first_sub_sol_stmts_no_temp_r_dict(self, 
                                                stmt0, stmt_r_expr, stmt_df_dx):
        soln_items = []
        
        stmt0_mod = f"""
        Let 
        $\\displaystyle 
        {{D}} =  {stmt0.strip()}
        $
        """ 
        soln_items.append(stmt0_mod.strip())

        stmt_mod = f"""
        $\\displaystyle 
        = {stmt_r_expr.strip()}
        $
        """ 
        soln_items.append(stmt_mod.strip())
        
        if stmt_df_dx.strip() != stmt_r_expr.strip():
            stmt_mod = f"""
            $\\displaystyle 
            = {stmt_df_dx.strip()}
            $
            """ 
            soln_items.append(stmt_mod.strip())
            
        return soln_items
         
        pass

    def _get_first_sub_sol_stmts_with_temp_r_dict(self,
                            stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict):
        soln_items = []
        
        stmt0_mod = f"""
        Let 
        $\\displaystyle 
        {{D}} =  {stmt0.strip()}
        $
        """ 
        soln_items.append(stmt0_mod.strip())

        stmt_mod = f"""
        $\\displaystyle 
        = {stmt_r_expr.strip()} \\qquad {stmt_r_dict.strip()}
        $
        """ 
        soln_items.append(stmt_mod.strip())

        if stmt_df_dx.strip() != stmt_r_expr.strip():
            stmt_mod = f"""
            $\\displaystyle 
            = {stmt_df_dx.strip()}
            $
            """ 
            soln_items.append(stmt_mod.strip())
            
        return soln_items
        pass



    def _get_non_first_sub_sol_stmts(self, 
        stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict,  temp_r_dict, sub_soln_idx, dR_latex):
        
        if not temp_r_dict:
            return self._get_non_first_sub_sol_stmts_no_temp_r_dict(
                    stmt0, stmt_df_dx, sub_soln_idx, dR_latex)
        
        return self._get_non_first_sub_sol_stmts_with_temp_r_dict(
                    stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict, sub_soln_idx, dR_latex)
        pass

    def _get_non_first_sub_sol_stmts_no_temp_r_dict(self, 
                        stmt0, stmt_df_dx, sub_soln_idx, dR_latex):
        soln_items = []
        
        if dR_latex:
            stmt_mod = f"{dR_latex} = {stmt0.strip()} = {stmt_df_dx.strip()}"
        else:
            stmt_mod = f"{stmt0.strip()} = {stmt_df_dx.strip()}"

        stmt_mod = f"""
        <h4>sub-step: {sub_soln_idx}</h4>
        $\\displaystyle 
        {stmt_mod}
        $
        """ 
        soln_items.append(stmt_mod.strip())
        
        return soln_items
         
        pass

    def _get_non_first_sub_sol_stmts_with_temp_r_dict(self,
                    stmt0, stmt_r_expr, stmt_df_dx, stmt_r_dict, sub_soln_idx, dR_latex):
        soln_items = []

        if dR_latex:
            stmt0_mod = f"{dR_latex} = {stmt0.strip()}"
        else:
            stmt0_mod = f"{stmt0.strip()}"

        
        stmt0_mod = f"""
        <h4>sub-step: {sub_soln_idx}</h4>
        $\\displaystyle 
        {stmt0_mod}
        $
        """ 
        soln_items.append(stmt0_mod.strip())

        stmt_mod = f"""
        $\\displaystyle 
        = {stmt_r_expr.strip()} \\qquad {stmt_r_dict.strip()}
        $
        """ 
        soln_items.append(stmt_mod.strip())
        pass

        if stmt_df_dx.strip() != stmt_r_expr.strip():
            stmt_mod = f"""
            $\\displaystyle 
            = {stmt_df_dx.strip()}
            $
            """ 
            soln_items.append(stmt_mod.strip())
            
        return soln_items
        pass



    def _get_dict_as_latex_list(self, temp_r_dict):
        ltx = self._modified_latex

        dict_stmts = []
        for k,v in temp_r_dict.items():
            stmt = f"""
                {ltx(k, mul_symbol='dot')}
                = {ltx(v, mul_symbol='dot')}
            """
            dict_stmts.append(stmt.strip())
            pass
        return dict_stmts

    def _print_r_dict(self, r_dict):
        ltx = self._modified_latex

        r_dict_stmts = []
        stmt = f"""
        r_dict:
        """
        r_dict_stmts.append(stmt.strip())
        for k,v in r_dict.items():
            stmt = f"""
            $\\displaystyle 
                {latex(k, mul_symbol='dot')}
                = {ltx(v, mul_symbol='dot')}
                $ 
                
            """
            r_dict_stmts.append(stmt.strip())
            pass
        return r_dict_stmts
        
   
    def _print_sub_soln_items(self, sub_soln):
        soln_items = []
        for k, v in sub_soln.items():
            try:
                stmt = f"""
                {k} = {v}
                """
                soln_items.append(stmt.strip())
            
            except Exception as e:
                sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
                k = {k}
                v = {v}
                """
                self._log(sError)
                self._log(traceback.format_exc())
                pass
            
            pass
        return soln_items
        pass
    


    

    def _modified_latex(self, expr, funct_symbs = None, mul_symbol = None):
        expr = self._convert_function_to_symbol(expr, funct_symbs = funct_symbs)
        expr_latex = latex(expr, mul_symbol = mul_symbol)
        
        expr_latex = expr_latex.replace("\\partial", "d")
        
        return expr_latex
        pass

    def _convert_function_to_symbol(self, expr, funct_symbs = None):
        if funct_symbs is None:
            funct_symbs = [symb for symb in self.var_symbs 
                           if isinstance(symb, Function)]

        if not funct_symbs:
            return expr
        
        temp_dict = {}
        for funct_symb in funct_symbs:
            temp_dict[funct_symb] = Symbol(funct_symb.name)
            pass
        
        if not temp_dict:
            return expr

        expr_subs = expr.subs(temp_dict)
        
        return expr_subs
        pass


    
    def _log_DELETED(self, sOut, to_print = None):
        if to_print is None:
            to_print = self.verbose
            
        if to_print:
            print(sOut)
            
    pass
#class DerivativeSolver:
