'''
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 xv.util import UtilClass

class CoordSysHelper(UtilClass):
    _COORD_SYS_TYPES = [
        'one_dimension',
        'cartesian2D',
        'cartesian3D',
        'polar',
        'spherical',
        'cylindrical',
        'ellipsoidal',
        ]
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._vars = self._init_variables()
        self._one_dimension_bases = self.get_vars(count = 1)
        self._cartesian2D_bases = self.get_vars(count = 2)
        self._cartesian3D_bases = self.get_vars(count = 3)
        self._polar_bases = self._init_polar_bases()
        self._spherical_bases = self._init_spherical_bases()
        self._ellipsoidal_bases = self._init_ellipsoidal_bases()
        self._cylindrical_bases = self._init_cylindrical_bases()
        
        self._funcs = self._init_funcs()
        pass

    def get_coord_sys_types(self):
        return CoordSysHelper._COORD_SYS_TYPES

    def get_bases(self, coord_sys_type):
        sys_types = self.get_coord_sys_types()
        if not coord_sys_type in sys_types:
            self._log(f"Invalid coord_sys_type = {coord_sys_type}")
            return None

        if coord_sys_type == 'one_dimension':
            return self._one_dimension_bases
        
        if coord_sys_type == 'cartesian2D':
            return self._cartesian2D_bases
        
        if coord_sys_type == 'cartesian3D':
            return self._cartesian3D_bases
        
        if coord_sys_type == 'polar':
            return self._polar_bases
        
        if coord_sys_type == 'spherical':
            return self._spherical_bases
        
        if coord_sys_type == 'ellipsoidal':
            return self._ellipsoidal_bases
        
        if coord_sys_type == 'cylindrical':
            return self._cylindrical_bases
        
        return None
        pass
    
    def get_funcs(self, count = 1):
        if count == 1:
            return self._funcs[0]
        return self._funcs[:count]
    
    def get_vars(self, count = 1):
        if count == 1:
            return self._vars[-3:-2]
        if count == 2:
            return self._vars[-3:-1]
        return self._vars[-count:]

    def get_xyz(self):
        return self._vars[-3:]

    def get_uvw(self):
        return self._vars[-6:-3]

    def get_dx(self, x):
        return Symbol(rf"d{latex(x)}")

    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_random_func_data(self):
        '''
        returns random data as dict
        '''
        
        func_data_set = []

        func_data = self.get_func_data_cartesian2D_to_cartesian2D()
        func_data_set.append(func_data)
        
        
        func_data = self.get_func_data_cartesian2D_to_polar()
        func_data_set.append(func_data)
        
        func_data = self.get_func_data_polar_to_cartesian2D()
        func_data_set.append(func_data)
        
        
        func_data = self.get_func_data_cartesian3D_to_spherical()
        func_data_set.append(func_data)
        
        func_data = self.get_func_data_spherical_to_cartesian3D()
        func_data_set.append(func_data)

        func_data = self.get_func_data_cartesian3D_to_ellipsoidal()
        return func_data
        func_data_set.append(func_data)
        
        
        func_data = self.get_func_data_cartesian3D_to_cylindrical()
        func_data_set.append(func_data)
        
        func_data = self.get_func_data_spherical_to_cylindrical()
        func_data_set.append(func_data)
        
        func_data = self.get_func_data_cylindrical_to_spherical()
        func_data_set.append(func_data)
        
        random_data = random.choice(func_data_set)
        return random_data
        pass

    def get_func_data_cartesian2D_to_cartesian2D(self):
        x, y, z = self.get_xyz()
        u, v, w = self.get_uvw()
        
        func_data = {}
        func_data['bases1'] = [x, y]
        func_data['bases2'] = [u, v]
        func_data['func'] = sqrt((x - y) * (x + 6 * y) )
        func_data['mapping_eqs'] = [Eq(u, x - y), Eq(v, x + 6 * y)]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, 7), (0, 7)]

        return func_data
        pass
    

    def get_func_data_cartesian2D_to_polar(self):
        x, y = self.get_bases('cartesian2D')
        r, theta = self.get_bases('polar')
        
        func_data = {}
        func_data['bases1'] = [x, y]
        func_data['bases2'] = [r, theta]
        func_data['func'] = random.choice([
                                1, 
                                sqrt(x**2 + y**2)
                            ])

        func_data['mapping_eqs'] = [Eq(x, r * cos(theta)), 
                                    Eq(y, r * sin(theta))
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, r), (0, 2 * pi)]
        return func_data
        pass
    
    def get_func_data_polar_to_cartesian2D(self):
        r, theta = self.get_bases('polar')
        x, y = self.get_bases('cartesian2D')
        
        func_data = {}
        func_data['bases1'] = [r, theta]
        func_data['bases2'] = [x, y]
        func_data['func'] = random.choice([
                                1, 
                            ])

        func_data['mapping_eqs'] = [Eq(r, sqrt(x**2 + y**2)), 
                                    Eq(theta, atan2(y, x))
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, r), (0, 2 * pi)]
        return func_data
        pass
    
  
    def get_func_data_cartesian3D_to_spherical(self):
        x, y, z = self.get_bases('cartesian3D')
        r, theta, phi = self.get_bases('spherical')
        
        func_data = {}
        func_data['bases1'] = [x, y, z]
        func_data['bases2'] = [r, theta, phi]
        func_data['func'] = 1

        func_data['mapping_eqs'] = [Eq(x, r * sin(theta) * cos(phi)), 
                                    Eq(y, r * sin(theta) * sin(phi)),
                                    Eq(z, r * cos(theta)),
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, r), (0, pi), (0, 2 * pi)]
        return func_data
        pass
    
    
    def get_func_data_spherical_to_cartesian3D(self):
        r, theta, phi = self.get_bases('spherical')
        x, y, z = self.get_bases('cartesian3D')
        
        func_data = {}
        func_data['bases1'] = [r, theta, phi]
        func_data['bases2'] = [x, y, z]
        func_data['func'] = 1

        func_data['mapping_eqs'] = [Eq(r, sqrt(x**2 + y**2 + z**2)), 
                                    Eq(theta, atan2(sqrt(x**2 + y**2), z)),
                                    Eq(phi, atan2(y, x)),
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, x), (0, y), (0, z)]
        return func_data
        pass

    def get_func_data_cartesian3D_to_ellipsoidal(self):
        x, y, z = self.get_bases('cartesian3D')
        r, theta, phi = self.get_bases('ellipsoidal')
        
        a = Symbol('a', constant = True)
        b = Symbol('b', constant = True)
        c = Symbol('c', constant = True)
        
        func_data = {}
        func_data['bases1'] = [x, y, z]
        func_data['bases2'] = [r, theta, phi]
        func_data['func'] = 1

        func_data['mapping_eqs'] = [Eq(x, a * r * sin(theta) * cos(phi)), 
                                    Eq(y, b * r * sin(theta) * sin(phi)),
                                    Eq(z, c * r * cos(theta)),
                                    ]
        
        
        func_data['limits1'] = None
        func_data['limits2'] = [(0, 1), (0, pi), (0, 2*pi)]
        return func_data
        pass
    
    
    def get_func_data_cartesian3D_to_cylindrical(self):
        x, y, z = self.get_bases('cartesian3D')
        r, theta, h = self.get_bases('cylindrical')
        
        func_data = {}
        func_data['bases1'] = [x, y, z]
        func_data['bases2'] = [r, theta, h]
        func_data['func'] = 1

        func_data['mapping_eqs'] = [Eq(x, r * cos(theta)), 
                                    Eq(y, r * sin(theta)),
                                    Eq(z, h),
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, r), (0, 2 * pi), (0, h)]
        return func_data
        pass
    
    def get_func_data_spherical_to_cylindrical(self):
        r, theta, phi = self.get_bases('spherical')
        rho, alpha, h = self.get_bases('cylindrical')
        
        func_data = {}
        func_data['bases1'] = [r, theta, phi]
        func_data['bases2'] = [rho, alpha, h]
        func_data['func'] = 1

        func_data['mapping_eqs'] = [Eq(r, sqrt(rho ** 2 + h ** 2)), 
                                    Eq(theta, atan2(rho, h)),
                                    Eq(phi, alpha),
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, r), (0, pi), (0, 2 * pi)]
        return func_data
        pass
    
    
    def get_func_data_cylindrical_to_spherical(self):
        rho, alpha, h = self.get_bases('cylindrical')
        r, theta, phi = self.get_bases('spherical')
        
        func_data = {}
        func_data['bases1'] = [rho, alpha, h]
        func_data['bases2'] = [r, theta, phi]
        func_data['func'] = 1

        func_data['mapping_eqs'] = [Eq(rho, r * sin(theta)), 
                                    Eq(alpha, phi),
                                    Eq(h, r * cos(theta)),
                                    ]
        func_data['limits1'] = None
        func_data['limits2'] = [(0, r), (0, pi), (0, 2 * pi)]
        return func_data
        pass
    
    def _init_variables(self):
        var_names = list('pqrstuvwxyz')
        return [Symbol(var_name) for var_name in var_names]
    
    def _init_funcs(self):
        var_names = list('fgjk')
        return [Function(var_name) for var_name in var_names]
    
    def _init_polar_bases(self):
        var_symbols = [Symbol('r', positive = True), 
                       Symbol(r'theta')]
        return var_symbols
    
    def _init_spherical_bases(self):
        var_symbols = [Symbol('r', positive = True), 
                       Symbol(r'theta'), 
                       Symbol(r'phi')]
        return var_symbols
    
    def _init_ellipsoidal_bases(self):
        var_symbols = [Symbol('r', positive = True), 
                       Symbol(r'theta'), 
                       Symbol(r'phi')]
        return var_symbols
    
    def _init_cylindrical_bases(self):
        var_symbols = [Symbol('rho', positive = True), 
                       Symbol(r'alpha'), 
                       Symbol('z', positive = True)]
        return var_symbols

    def _log_DELETED(self, sOut, to_print = None):
        if to_print is None:
            to_print = self.verbose
            
        if to_print:
            print(sOut)
    
    def _doit_with_timeout(self, expr, timeout = 2):
        try:
            if not hasattr(expr, 'doit'):
                return None
            
            from xv.util import run_with_timeout
            #solved_values = solve(eq, var_to_find)
            result = run_with_timeout(f'Thread_eq_solver', 
                                      expr.doit, 
                                      None, 
                                      timeout = timeout, 
                                      verbose = self.verbose)
            return result
            
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}:{e}
            expr = {expr}"""
            print(sError)
            print(traceback.format_exc())
            pass
        
        return None
        pass

    
    pass
#class CoordSysHelper:
