'''
Created on 20-Jan-2022

@author: bthakur
'''

import codecs
import inspect
import traceback
import random
import numpy as np
from sympy import (latex, Symbol, symbols, MatrixSymbol,
                    Rational, pi,
                    sin, cos, tan,
                    asin, acos, atan, atan2,
                    sinh, cosh, tanh,
                    log, exp, sqrt,
                   diff, factor, 
                   Function, Eq,
                   Add, 
                   simplify, solve, factorial,
                   Sum, binomial,
                   Matrix,
                   Derivative, Integral, Determinant,
                   )

from sympy.core.numbers import One, Integer
from sympy.series import series

from xv.km.util import latex_fraction
from sympy.core.expr import UnevaluatedExpr

from xv.km.ui import SvgHelper
from xv.util.data import CacheUtil
from xv.util import list_to_str

import matplotlib.pyplot as plt

from .._vector import CoordSysHelper

class JacobianHelper(CoordSysHelper):
    _COORD_SYS_TYPES = [
        'one_dimension',
        'cartesian2D',
        'cartesian3D',
        'polar',
        'spherical',
        'cylindrical',
        'ellipsoidal',
        ]
    
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        pass


    def get_jacobian_symbol(self, bases1 = None, bases2 = None, dim = None):
        if not dim is None:
            if dim[0] != dim[1]:
                dim = None
        
        if bases1 is None and bases2 is None:
            if not dim is None:
                return MatrixSymbol('J', *dim)
            else:
                return Symbol('J')
            
        
        if isinstance(bases1, list):
            bases1 = list_to_str([f"{latex(x_var)}" for x_var in bases1], 
                                 sep = " ")
        if isinstance(bases2, list):
            bases2 = list_to_str([f"{latex(x_var)}" for x_var in bases2], 
                                 sep = " ")
            
        if not dim is None:
            return MatrixSymbol(f'{{J}}_{{{bases1}}}^{{{bases2}}}', *dim)
        
        return Symbol(f'{{J}}_{{{bases1}}}^{{{bases2}}}')
        pass


    def get_compatible_coord_sys_pairs(self):
        pairs = [
            ['cartesian2D', 'polar'],
            ['cartesian3D', 'spherical'],
            ['cartesian3D', 'ellipsoidal'],
            ['cartesian3D', 'cylindrical']
            ]
        reverse_pairs = [p[::-1] for p in pairs]
        pairs.extend(reverse_pairs)
        return pairs
    
    def get_jacobian_matrix_data(self, 
                                  bases1 = None, 
                                  bases2 = None, 
                                  mapping_eqs = None,
                                  ):

        try:        
            
            bases1, bases2 = self._parse_bases_pair(bases1, bases2)
            
            (mapping_eqs, 
             mapping_functions, 
             mapping_functions_stmt) = self._parse_mapping_eqs(bases1, bases2, mapping_eqs)
                
            if mapping_functions is None:
                self._log("Failed to get mapping_functions")
                return None
                pass
            
            jacobian_matrix_form1 = []
            jacobian_matrix_form2 = []
            
            for func_idx, func in enumerate(mapping_functions):
                
                row_form1 = []
                row_form2 = []
                for v2 in bases2:
                    row_form1.append(Derivative(bases1[func_idx], v2))
                    row_form2.append(Derivative(func, v2))
                    pass
    
                jacobian_matrix_form1.append(row_form1)
                jacobian_matrix_form2.append(row_form2)
                pass
            
            jacobian_matrix_form1 = Matrix(jacobian_matrix_form1)
            jacobian_matrix_form2 = Matrix(jacobian_matrix_form2)
            
            try:
                jacobian_matrix = jacobian_matrix_form2.doit()
            except:
                jacobian_matrix = jacobian_matrix_form2
            pass
            
            
            mapping_eqs_latex = [f"$\displaystyle {latex(eq)}$" 
                                       for eq in mapping_eqs] 
            mapping_eqs_latex = list_to_str(mapping_eqs_latex, sep = "<br/><br/>")
            
            jacobian_symbol = self.get_jacobian_symbol(bases1 = bases1, 
                                                       bases2 = bases2,
                                                       dim = (len(bases1), len(bases2)))
            result = {
                'bases1': bases1,
                'bases2': bases2,
                'jacobian_symbol': jacobian_symbol,
                'mapping_eqs': mapping_eqs,
                'mapping_eqs_latex': mapping_eqs_latex,
                'mapping_functions': mapping_functions,
                'mapping_functions_stmt': mapping_functions_stmt,
                'jacobian_matrix_form1': jacobian_matrix_form1,
                'jacobian_matrix_form2': jacobian_matrix_form2,
                'jacobian_matrix': jacobian_matrix,
                }
            return result       
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            bases1 = {bases1}
            bases2 = {bases2}
            mapping_eqs = {mapping_eqs}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            
            pass

        return None        
        pass

    def get_jacobian_matrix_statements(self, data):
        '''
        data must be return dict from get_jacobian_matrix_data function of this class
        '''
        bases1 = data['bases1']
        bases2 = data['bases2']
        mapping_eqs_latex = data['mapping_eqs_latex']
        jacobian_symbol = data['jacobian_symbol']
        jacobian_matrix_form1 = data['jacobian_matrix_form1']
        jacobian_matrix_form2 = data['jacobian_matrix_form2']
        jacobian_matrix = data['jacobian_matrix']
        bases1 = data['bases1']
        bases1 = data['bases1']
        bases1 = data['bases1']
        
        problem_s = f"""
            Write Jacobian of $\displaystyle 
            {latex(bases1)} \\rightarrow {latex(bases2)} 
            $ when 
            
            {mapping_eqs_latex} 
            """        
        soln_s = []
        
        soln_s.append(f"""
            Jacobian of $\displaystyle 
                {latex(bases1)} \\rightarrow {latex(bases2)} 
                $

            $\displaystyle = {jacobian_symbol}$
            
            $\displaystyle = {latex(jacobian_matrix_form1)}$
            """.strip())
   
        if str(jacobian_matrix_form1) != str(jacobian_matrix_form2):
            soln_s.append(f"""
                $\displaystyle = {latex(jacobian_matrix_form2)}$
                """.strip())
       
        if str(jacobian_matrix_form2) != str(jacobian_matrix):
            soln_s.append(f"""
                $\displaystyle = {latex(jacobian_matrix)}$
                """.strip())
   
        soln_s = list_to_str(soln_s, sep  = "\n\n")
       
        answer_s = f"""
            $\displaystyle {latex(jacobian_matrix)}$
            """

        result = {
            'problem_s': problem_s,
            'soln_s': soln_s,
            'answer_s': answer_s
            }
        
        return result       
        pass


    def get_transform_data(self, 
                          bases1 = None, 
                          bases2 = None, 
                          mapping_eqs = None,
                          func = None,
                          ):
        try:
            result = self.get_jacobian_matrix_data(bases1 = bases1, 
                                                  bases2 = bases2, 
                                                  mapping_eqs = mapping_eqs,
                                                  )
            if result is None:
                self._log("Failed to get jacobian_matrix_data")
                return None
            
            bases1 = result['bases1']
            bases2 = result['bases2']
            mapping_eqs = result['mapping_eqs']
            mapping_functions = result['mapping_functions']
            
            jacobian_matrix = result['jacobian_matrix']
            jacobian_symbol = result['jacobian_symbol']
            J = jacobian_symbol
            determinant_symbol = Determinant(J)
    
            jacobian_determinant = simplify(Determinant(jacobian_matrix).doit())
            deteminant_statement = self.get_jacobian_deteminant_statement(
                                        jacobian_symbol, jacobian_matrix)
    
            (func_original, 
             func_tranformed, 
             func_tranformed_statement) = self._parse_func(func, bases1, mapping_functions)
            
            func1 = 1
            for x in bases1:
                func1 = func1 * UnevaluatedExpr(self.get_dx(x))
                pass
            func1 = func_original * func1
            
            func2 = 1
            for u in bases2:
                func2 = func2 * UnevaluatedExpr(self.get_dx(u))
                pass
            
            jacobian_transform_form1 = func_tranformed * UnevaluatedExpr(
                                        determinant_symbol) * UnevaluatedExpr(func2)
                                        
            jacobian_transform_form2 = func_tranformed * UnevaluatedExpr(
                                        jacobian_determinant) * UnevaluatedExpr(func2)
            jacobian_transform = simplify(func_tranformed * jacobian_determinant
                                          ) * UnevaluatedExpr(func2)
            
            
            result['func_original'] = func_original
            result['func_tranformed'] = func_tranformed
            result['func1'] = func1
            #result['func2'] = func2
            result['jacobian_determinant'] = jacobian_determinant
            result['deteminant_statement'] = deteminant_statement
            
            result['jacobian_transform_form1'] = jacobian_transform_form1 
            result['jacobian_transform_form2'] = jacobian_transform_form2 
            result['jacobian_transform'] = jacobian_transform 
            
            return result       
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            bases1 = {bases1}
            bases2 = {bases2}
            mapping_eqs = {mapping_eqs}
            func = {func}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            
            pass

        return None        
        pass

    def get_transform_statements(self, data):
        '''
        Data must be dict returned from function get_transform_data of this class 
        '''
        if data is None:
            self._log("data is None")
            return None

        bases1 = data['bases1']
        bases2 = data['bases2']
        func1 = data['func1']
        mapping_eqs_latex = data['mapping_eqs_latex']
        mapping_functions_stmt = data['mapping_functions_stmt']
        deteminant_statement = data['deteminant_statement']
        
        jacobian_transform_form1 = data['jacobian_transform_form1']
        jacobian_transform_form2 = data['jacobian_transform_form2']
        jacobian_transform = data['jacobian_transform']
        
        jacobian_matrix_statements = self.get_jacobian_matrix_statements(data)
       
        problem_s = f"""
            Write $\displaystyle 
            {latex(bases1)} 
            \\rightarrow
            {latex(bases2)} 
            $ transformation of $\displaystyle 
            {latex(func1)}$ when
            
            {mapping_eqs_latex} 
            """
        
        soln_s = []

        soln_s_item = f"""
            <b>Step 1:</b>
            {mapping_functions_stmt}
        
            <b>Step 2:</b>
            {jacobian_matrix_statements['soln_s']}
            
            <b>Step 3:</b>
            {deteminant_statement}
            
            
            <b>Step 4:</b>
            Now, 
            
            $\displaystyle 
            {latex(func1)} 
            $
            
            $\displaystyle
            = {latex(jacobian_transform_form1)}
            $
            """
        soln_s.append(soln_s_item.strip())
        
        if str(jacobian_transform_form1) != str(jacobian_transform_form2):
            soln_s_item = f"""
                $\displaystyle
                = {latex(jacobian_transform_form2)}
                $
                """
            soln_s.append(soln_s_item.strip())
            
        if str(jacobian_transform_form2) != str(jacobian_transform):
            soln_s_item = f"""
                $\displaystyle
                = {latex(jacobian_transform)}
                $
                """
            soln_s.append(soln_s_item.strip())
            
        soln_s = list_to_str(soln_s, sep = "\n\n") 

        answer_s = f"""
            $\displaystyle {latex(func1)} = {latex(jacobian_transform)}
            $
            """
            
        result = {
            'problem_s': problem_s,
            'soln_s': soln_s,
            'answer_s': answer_s
            }
        
        return result       
        pass

    def get_integration_transform_data(self, 
                                  bases1 = None, 
                                  bases2 = None, 
                                  mapping_eqs = None,
                                  func = None,
                                  limits1 = None,
                                  limits2 = None,
                                  ):
        try:
        
            result = self.get_jacobian_matrix_data(bases1 = bases1, 
                                                  bases2 = bases2, 
                                                  mapping_eqs = mapping_eqs,
                                                  )
            if result is None:
                self._log("Failed to get matrix_statements")
                return None
    
            result['limits1'] = limits1
            result['limits2'] = limits2
            
            
            bases1 = result['bases1']
            bases2 = result['bases2']
            mapping_eqs = result['mapping_eqs']
            mapping_functions = result['mapping_functions']
            
            jacobian_matrix = result['jacobian_matrix']
            jacobian_symbol = result['jacobian_symbol']
            J = jacobian_symbol
            determinant_symbol = Determinant(J)
    
            jacobian_determinant = simplify(Determinant(jacobian_matrix).doit())
            deteminant_statement = self.get_jacobian_deteminant_statement(
                                        jacobian_symbol, jacobian_matrix)
    
            result['jacobian_determinant'] = jacobian_determinant
            result['deteminant_statement'] = deteminant_statement
            
    
            (func_original, 
             func_tranformed, 
             func_tranformed_statement) = self._parse_func(func, bases1, mapping_functions)
            
            ##########################
            func1 = func_original
            func1 = self._compute_integral_expr(func_original, bases1, limits = limits1)

    
            temp_func = UnevaluatedExpr(determinant_symbol) * func_tranformed
            jacobian_transform_form1 = self._compute_integral_expr(
                                        temp_func, bases2, limits = limits2)

            temp_func = jacobian_determinant * func_tranformed
            jacobian_transform_form2 = self._compute_integral_expr(
                                        temp_func, bases2, limits = limits2)

            if limits2 is None:
                func4_doits = jacobian_transform_form2
                computed_funcs = []
            else:
                temp_func = jacobian_determinant * func_tranformed
                func4_doits = self._compute_func4_doits(temp_func, 
                                                        bases2, 
                                                        limits2)    
                computed_funcs = self._get_computed_funcs(func4_doits, bases2, limits2)
                pass
            
            jacobian_transform = jacobian_transform_form2.doit()
        
            result['func1'] = func1
            result['func_original'] = func_original
            result['func_tranformed'] = func_tranformed
            
            result['computed_funcs'] = computed_funcs
            
            result['jacobian_transform_form1'] = jacobian_transform_form1 
            result['jacobian_transform_form2'] = jacobian_transform_form2 
            result['jacobian_transform'] = jacobian_transform 
            
            return result       
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            bases1 = {bases1}
            bases2 = {bases2}
            mapping_eqs = {mapping_eqs}
            func = {func}
            limits2 = {limits2}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            
            pass

        return None        
        pass
    
    def _compute_integral_expr(self,
                               func,
                               bases,
                               limits = None):
        for u_idx, u_var in enumerate(bases):
            if limits is None:
                func = Integral(func, u_var)
            else:
                u1, u2 = limits[u_idx]
                func = Integral(func, (u_var, u1, u2))
            pass
        
        return func
        pass
    
    def _compute_func4_doits(self, 
                       func4,
                       bases2,
                       limits2):
        
        self._log(f"Computing func4")
        #func4 = jacobian_determinant * func_tranformed
        func4_doits = []
        
        for u_idx, u_var in enumerate(bases2):
            if not limits2 is None:
                u1, u2 = limits2[u_idx]
                
                if u_idx == 0:
                    func4 = Integral(func4, (u_var, u1, u2))
                else:
                    func4 = Integral(func4_doits[-1], (u_var, u1, u2))
                    
                self._log(f"Applying doit to function {func4}")
                func4_doit = self._doit_with_timeout(func4, timeout = 2)
                if func4_doit is None:
                    self._log(f"Failed to compute doit() for {func4}")
                    return None
                self._log(f"func4_doit =  {func4_doit}")

                func4_doits.append(func4_doit)
            pass
        
        return func4_doits
        pass
    
    def _get_computed_funcs(self, 
                       func4_doits,
                       bases2,
                       limits2):
        
        computed_funcs = []
        if limits2 is None:
            return computed_funcs
        
        for idx in range(len(func4_doits) - 1):
            computed_func = simplify(func4_doits[idx])
            
            for u_idx in range( idx + 1, len(bases2)):
                u_var = bases2[u_idx]
                u1, u2 = limits2[u_idx]
                
                computed_func = Integral(computed_func, (u_var, u1, u2))
                pass
            
            computed_funcs.append(computed_func)
            pass
        
        return computed_funcs
        pass
    

    def get_integration_transform_statements(self, data):
        if data is None:
            self._log("data is None")
            return None
        
        bases1 = data['bases1']
        bases2 = data['bases2']
        limits1 = data['limits1']
        limits2 = data['limits2']
        func1 = data['func1']
        mapping_eqs_latex = data['mapping_eqs_latex']
        mapping_functions_stmt = data['mapping_functions_stmt']
        deteminant_statement = data['deteminant_statement']
        
        
        computed_funcs = data['computed_funcs']
        
        jacobian_transform_form1 = data['jacobian_transform_form1']
        jacobian_transform_form2 = data['jacobian_transform_form2']
        jacobian_transform = data['jacobian_transform']
        
        jacobian_matrix_statements = self.get_jacobian_matrix_statements(data)
       
        
        problem_s = []
        problem_s_item = f"""
            Compute $\displaystyle 
            {latex(func1)}$ 
            
            using $\displaystyle 
            {latex(bases1)} 
            \\rightarrow
            {latex(bases2)} 
            $ transformation, when
            
            {mapping_eqs_latex} 
            """
        problem_s.append(problem_s_item.strip())

        if not limits1 is None:
            problem_s_item = f"""
                $\displaystyle 
                {latex(bases1)} 
                $ limits are $\displaystyle 
                {latex(limits1)} 
                $
                """
            problem_s.append(problem_s_item.strip())
            pass
        
        if not limits2 is None:
            problem_s_item = f"""
                $\displaystyle 
                {latex(bases2)} 
                $ limits are $\displaystyle 
                {latex(limits2)} 
                $
                """
            problem_s.append(problem_s_item.strip())
            pass

        problem_s = list_to_str(problem_s, sep = "<br/><br/>")
        
        soln_s = []

        soln_s_item = f"""
            <b>Step 1:</b>
            {mapping_functions_stmt}
        
            <b>Step 2:</b>
            {jacobian_matrix_statements['soln_s']}
            
            <b>Step 3:</b>
            {deteminant_statement}
            
            
            <b>Step 4:</b>
            Now, 
            
            $\displaystyle 
            {latex(func1)} 
            $
            
            $\displaystyle
            = {latex(jacobian_transform_form1)}
            $
            """
        soln_s.append(soln_s_item.strip())
        
        if str(jacobian_transform_form1) != str(jacobian_transform_form2):
            soln_s_item = f"""
                $\displaystyle
                = {latex(jacobian_transform_form2)}
                $
                """
            soln_s.append(soln_s_item.strip())
            
        if str(jacobian_transform) != str(jacobian_transform_form2):
            computed_funcs = [f"$\displaystyle = {latex(computed_func)}$" 
                           for computed_func in computed_funcs]
            
            computed_funcs = list_to_str(computed_funcs, sep = "<br/><br/>")
            soln_s.append(computed_funcs.strip())
            
            soln_s_item = f"""
                $\displaystyle 
                = {latex(jacobian_transform)}
                $
                """
            soln_s.append(soln_s_item.strip())
            pass
            
        soln_s = list_to_str(soln_s, sep = "\n\n") 

        answer_s = f"""
            $\displaystyle {latex(jacobian_transform)}
            $
            """
            
        result = {
            'problem_s': problem_s,
            'soln_s': soln_s,
            'answer_s': answer_s
            }
        
        return result       
        ##########################################
        
        soln_s = []
        result['problem_s'] = problem_s
        result['soln_s'] = soln_s
        result['answer_s'] = answer_s
        
        return result       
        pass
    
    

    def get_jacobian_deteminant_statement(self, jacobian_symbol, jacobian_matrix):
        jacobian_determinant_form1 = Determinant(jacobian_symbol)
        jacobian_determinant_form2 = Determinant(jacobian_matrix)
        jacobian_determinant_form3 = jacobian_determinant_form2.doit()
        jacobian_determinant = simplify(jacobian_determinant_form3)
            
        soln_s = []
        
        soln_s.append(f"""
            Determinant of $\displaystyle {latex(jacobian_symbol)}$
            
            $\displaystyle = {latex(jacobian_determinant_form1)}$
            
            $\displaystyle = {latex(jacobian_determinant_form2)}$
            """.strip())
   
        if str(jacobian_determinant_form2) != str(jacobian_determinant_form3):
            soln_s.append(f"""
                $\displaystyle = {latex(jacobian_determinant_form3)}$
                """.strip())
   
        if str(jacobian_determinant_form3) != str(jacobian_determinant):
            soln_s.append(f"""
                $\displaystyle = {latex(jacobian_determinant)}$
                """.strip())
   
        soln_s = list_to_str(soln_s, sep  = "\n\n")
        
        return soln_s       
        pass
    

    def _parse_func(self, func, bases, mapping_functions):
        if func is None:
            func = 1
            
        if not hasattr(func, 'subs'):
            func_original = func
            func_tranformed = func_original
            func_tranformed_statement = None
            return func_original, func_tranformed, func_tranformed_statement

    

        func_original = func
        subs_dict = dict(zip(bases, mapping_functions))
        func_tranformed = simplify(func.subs(subs_dict))
        
        func_tranformed_statement = f"""
        Transform the function into new variables:
         $\displaystyle 
        {latex(func)} 
        = {latex(func_tranformed)} 
        $
        """
        
        return func_original, func_tranformed, func_tranformed_statement
        pass
    
    def _mapping_eqs_to_mapping_functions(self, mapping_eqs, bases): 
        try:
            bases_str = list_to_str([f"$\displaystyle {latex(x_var)}$ " 
                                     for x_var in bases])
            
            mapping_eqs_str = [f"$\displaystyle {latex(eq)}$ "
                                     for eq in mapping_eqs]
            mapping_eqs_str = list_to_str(mapping_eqs_str, sep = "\n\n")
            
            soln_set = solve(mapping_eqs, bases)
            
            if not isinstance(soln_set, dict):
                mapping_eqs_str = list_to_str(mapping_eqs, sep = "\n\n")
                self._log(f"""
                Failed to get solution for the following equations.
                bases = {bases}
                
                {mapping_eqs_str}
                """)
                return None, ""
            
            mapping_functions = list(soln_set.values())
            
            soln_set_str = [f"$\displaystyle {latex(k1)} = {latex(v1)}$ " 
                            for k1, v1 in soln_set.items()]
            soln_set_str = list_to_str(soln_set_str, sep = "\n\n")
            
            
            mapping_functions_stmt = f"""
            The equations are:
            
            {mapping_eqs_str}
    
            Solve the equations for {bases_str}:
            
            {soln_set_str}
            """
            
            self._log(f"mapping_functions = {mapping_functions}")
            return mapping_functions, mapping_functions_stmt.strip()
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            bases = {bases}
            mapping_eqs = {mapping_eqs}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            
            pass

        return None        
        pass


    def _parse_bases_pair(self, bases1 = None, bases2 = None):
        if not bases1 is None and not bases2 is None:
            return bases1, bases2
            
        coord_sys_pairs = self.get_compatible_coord_sys_pairs()
        coord_sys_pair = random.choice(coord_sys_pairs)
        bases1 = self.get_bases(coord_sys_pair[0])
        bases2 = self.get_bases(coord_sys_pair[1])
        return bases1, bases2
        pass
    
    def _parse_mapping_eqs(self, bases1, bases2, mapping_eqs = None):
        if mapping_eqs is None:
            self._log("mapping_eqs is None. Setting default mapping_eqs")
            mapping_fs = self.get_funcs(count = len(bases1))
            mapping_eqs = zip(bases1, mapping_fs)
            mapping_eqs = [Eq(x_var, func(*bases2)) 
                           for x_var, func in mapping_eqs]
            pass
        
        (mapping_functions, 
         mapping_functions_stmt) = self._mapping_eqs_to_mapping_functions(mapping_eqs, 
                                                                         bases1)
        return mapping_eqs, mapping_functions, mapping_functions_stmt
    
    pass
#class JacobianHelper:
