'''
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,
                    Rational,
                    sin, cos, tan,
                    asin, acos, atan,
                    sinh, cosh, tanh,
                    log, exp, sqrt,
                   diff, factor, 
                   Function, Eq,
                   Add, Mul, Pow,
                   simplify, solve, factorial,
                   Sum, binomial,
                   oo
                   )

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

from xv.util import toIntOrFloat

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

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

import matplotlib.pyplot as plt
from xv.km.kelements import KProblemTemplate

from xv.util import UtilClass

# multinomial_helper

class MultinomialHelper(UtilClass):
    '''
    This class deals with expressions of more than two variable and their applications.
    For two variables, see BinomialHelper
    '''
    _DEFAULT_TERM_COUNT = 4
    _MAX_DISPLAYABLE_TERM_COUNT = 11
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        #self.verbose = kwargs.get('verbose', False)
        pass

    def get_default_vars(self, count = None):
        if count is None:
            count = random.randint(3, 6)
        import string
        default_vars = symbols(list(string.ascii_lowercase)[:count])
        
        return default_vars
        pass


    def get_random_exponent_dict(self, count = None, m_vars = None):
        if m_vars is None:
            if count is None:
                count = random.randint(3, 6)
            m_vars = self.get_default_vars(count)
        
        exponent_dict = {}
        for m_var in m_vars:
            exponent_dict[m_var] = random.randint(1, 3)
            pass
        return exponent_dict

    
    def get_multinomial_term(self, exponent_dict = None):
        if exponent_dict is None:
            exponent_dict = self.get_random_exponent_dict()
            
        m_vars = list(exponent_dict.keys())
        
        exponent_sum = sum(list(exponent_dict.values()))

        term_dict = {}
        temp_sum = exponent_sum
        for m_var, exponent_val  in exponent_dict.items():
            ueval_coeff = binomial(Symbol(f"{temp_sum}"), Symbol(f"{exponent_val}"))

            expanded_coeff = self.get_expended_binomial(temp_sum, exponent_val)
            coeff = binomial(temp_sum, exponent_val)
            var_expr = m_var ** exponent_val
            
            ueval_term = ueval_coeff * ueval(var_expr)
            expanded_term = Symbol(f"\\;{latex(expanded_coeff)}\\; ") * ueval(var_expr)
            term = coeff * var_expr
            
            term_dict[m_var] = {'ueval_coeff': ueval_coeff,
                                      'expanded_coeff': expanded_coeff,
                                      'coeff': coeff,
                                      'var_expr': var_expr,
                                      'ueval_term': ueval_term,
                                      'expanded_term': expanded_term,
                                      'term': term,
                                      }
            
            temp_sum = temp_sum - exponent_val
            pass
        
        result = {}
        
        result['multinomial_expr'] = ueval(Add(*m_vars)) ** exponent_sum
        
        result['ueval_term'] = Mul(*[ueval(v['ueval_term']) for v in term_dict.values()])
        result['expanded_term'] = Mul(*[ueval(v['expanded_term']) for v in term_dict.values()])
        result['coeff'] = Mul(*[v['coeff'] for v in term_dict.values()])
        result['var_expr'] = Mul(*[v['var_expr'] for v in term_dict.values()])
        result['term'] = Mul(*[v['term'] for v in term_dict.values()])
        
        result['uncomputed'] = term_dict
        
        return result
        pass
    

    def get_expended_binomial(self, n, r):
        if not isinstance(r, int):
            return None
        
        if r < 0:
            return None
       
        if isinstance(n, int):
            if r > n:
                return None
            r = min(r, n-r)

        p = self.get_ueval_factorial(n, r)
        q = self.get_ueval_factorial(r, r)
        
        
        return Symbol(latex(p)) / Symbol(latex(q))
        pass
    
    def get_ueval_factorial(self, start, count = None):
        if count is None:
            count = start

        if not isinstance(count, int):
            return None
        
        ueval_factorial = []
        for i in range(count):
            ueval_factorial.append(ueval(start - i))
        
        return Mul(*ueval_factorial)
        pass
    
    def get_multinomial_expansion(self, m_vars, exponent_sum = None):
        if exponent_sum is None:
            exponent_sum = random.randint(2, 3)

        multinomial_expr = ueval(Add(*m_vars)) ** exponent_sum
        uneval_multinomial_expansion = self._get_multinomial_expansion(m_vars, exponent_sum, evaluate = False)
        multinomial_expansion = self._get_multinomial_expansion(m_vars, exponent_sum, evaluate = True)
            
        result = {}
        
        result['multinomial_expr'] = multinomial_expr
        result['ueval_expansion'] = uneval_multinomial_expansion
        result['expansion'] = multinomial_expansion
        result['simplified'] = simplify(multinomial_expansion)
        
        return result
        pass


    def _get_multinomial_expansion(self, temp_vars, exponent_sum, evaluate = True):
        
        m_exprs = []
        for i in range(len(temp_vars)):
            for exponent in range(exponent_sum, 0, -1):
                
                if evaluate:
                    coeff = binomial(exponent_sum, exponent)
                    m_expr = coeff * temp_vars[i] ** exponent
                else:
                    coeff = binomial(Symbol(f"{exponent_sum}"), Symbol(f"{exponent}"))
                    m_expr = ueval(coeff) * ueval(temp_vars[i]) ** ueval(exponent)
                    
                
                if i == len(temp_vars) - 1:
                    m_exprs.append(m_expr)
                    break
                
                if exponent_sum - exponent <= 0 or not temp_vars[i+1:]:
                    m_exprs.append(m_expr)
                    continue
                
                #print(f""" Calling {temp_vars[i+1:]} ** {exponent_sum - exponent} """)
                m_sub_exprs = self._get_multinomial_expansion(temp_vars[i+1:], 
                                                              exponent_sum - exponent, 
                                                              evaluate = evaluate)
                if m_sub_exprs:
                    if evaluate:
                        m_sub_exprs = ueval(sympy.Add(m_sub_exprs))
                        m_expr = ueval(m_expr) * ueval(m_sub_exprs)
                    else:
                        m_sub_exprs = sympy.Add(m_sub_exprs)
                        m_expr = m_expr * m_sub_exprs
                        pass
                    m_exprs.append(m_expr)
                    pass
                pass
            pass
        
        
        if evaluate:
            m_exprs =  sympy.Add(*m_exprs)
        else:
            m_exprs = [ueval(item) for item in m_exprs]
            m_exprs =  ueval(sympy.Add(*m_exprs))
            pass
            
        return m_exprs
        pass
    

    
    pass
#class MultinomialHelper:
