'''
Created on 12-Jul-2021

@author: bthakur
'''
import random
import inspect
import traceback
import re
from sympy import (Eq, symbols, Symbol, sqrt, latex, factor,
                   Rational, solve, log, simplify, E, exp,  
                   sin, cos, I, series, pi,
                   UnevaluatedExpr as ueval
                   )
from sympy.abc import theta
import sympy

from xv.km import CustomLogarithm as Log
from xv.km.util import latex_fraction, latex_exponent
from _collections import OrderedDict

# _arithmetic_series
# ArithmeticSeries

class ArithmeticSeries:
    def __init__(self, a, r, n, **kwargs):
        self.verbose = kwargs.get('verbose', False)
        self.a = a
        self.r = r
        self.n = n

        self.alpha = symbols(r'\alpha') # 1+r
        self.beta = symbols(r'\beta') # 1/(1+r)
        pass

    def get_alpha_value(self, r = None, to_latex = False):
        if r is None:
            r = self.r

        if not to_latex:
            return 1 + r
        
        r_latex = self.get_var_latex(r)
        return f"1 + {r_latex}"
    
    def get_beta_value(self, r = None, to_latex = False):
        if r is None:
            r = self.r
        
        if not to_latex:
            return 1 / (1 + r)
        
        r_latex = self.get_var_latex(r)
        return latex_fraction(1, f"1 + {r_latex}")
    
    def get_term(self, index, to_latex = False, bracket = True):
        if not to_latex:
            return self.a + self.r * index
        
        r_latex = self.get_var_latex(self.r)
        a_latex = self.get_var_latex(self.a)
        
        index_latex = self.get_var_latex(index)
        
        term_latex = f"{a_latex} + {r_latex} \\cdot {{{index_latex}}}"
        
        if not bracket:
            return term_latex

        term_latex = f"\\left( {term_latex} \\right)"

        return term_latex    
        pass
    
    
    
    
    def get_terms(self, count = None, to_latex = False):
        if count is None:
            count = self.n
        
        is_second_lot = True
        first_lot_count = 3
        if isinstance(count, int):
            if count < 6:
                is_second_lot = False
                first_lot_count = count
                pass
        
        r_latex = self.get_var_latex(self.r)
        a_latex = self.get_var_latex(self.a)
        count_2_latex = self.get_var_latex(count-2)
        count_1_latex = self.get_var_latex(count-1)
        
        terms = { i: self.a + self.r * i  for i in range(first_lot_count)}
        terms_latex = { i:f"{a_latex} + {r_latex} \\cdot {{{i}}}"  for i in range(first_lot_count)}
        
        if is_second_lot:
            terms[count-2] = self.a + self.r * (count-2)
            terms_latex[count-2] = f"{a_latex} + {r_latex} \\cdot {count_2_latex}" 

            terms[count-1] = self.a + self.r * (count-1)
            terms_latex[count-1] = f"{a_latex} + {r_latex} \\cdot {count_1_latex}" 
        
        if to_latex:
            terms_latex = {k:f"\\left( {v} \\right)" for k, v in terms_latex.items()}
            
            return terms_latex
        
        return terms
        pass
    
    def get_computed_terms(self, count = None):
        computed_terms = self.get_terms(to_latex =  False)
        computed_terms = {k:simplify(v) for k, v in computed_terms.items()}
        return computed_terms
    
    def get_summtion_latex(self, a = None, r = None, n = None):
        if a is None:
            a = self.a
            
        if r is None:
            r = self.r
            
        if n is None:
            n = self.n

        r_latex = self.get_var_latex(r)
        a_latex = self.get_var_latex(a)
            
        s = f"""
        {{\\large\\sum\\limits_{{k=0}}^{{{latex(n-1)}}}}} 
        \\left( {{ {a_latex} + {r_latex} \\cdot {{k}} }}  \\right)
        """
        return s.strip()    
    
    
    
    def get_series_computation_stmt(self):
        
        gs = self
        
        terms = gs.get_terms(to_latex =  True)
        computed_terms = gs.get_computed_terms()
        series = gs.get_series(terms, to_str = False)
        series_str = gs.get_series(terms, to_str = True)
        summtion_latex = gs.get_summtion_latex()

        computed_terms_str = gs.get_series(computed_terms, to_str = True)
        

        
        r_latex = gs.get_var_latex(gs.r)
        n_latex = gs.get_var_latex(gs.n)
        
        soln_s = f"""
        
        ${{t}}_{{0}} = {latex(gs.a)} = {gs.get_term(0, to_latex = True)} $ 
        
        ${{t}}_{{1}} 
        = {{t}}_{{0}} + {r_latex}
        = {gs.get_term(0, to_latex = True)} + {r_latex}
        = {gs.get_term(1, to_latex = True)} $ 
        
        ${{t}}_{{2}} 
        = {{t}}_{{1}} + {r_latex}
        = {gs.get_term(1, to_latex = True)} + {r_latex}
        = {gs.get_term(2, to_latex = True)} $ 
        
        $\\cdots$
        
        ${{t}}_{{{latex(gs.n - 1)}}} 
        = {{t}}_{{{latex(gs.n - 2)}}} + {r_latex}
        = {gs.get_term(gs.n - 2, to_latex = True)} + {r_latex}
        = {gs.get_term(gs.n - 1, to_latex = True)} $ 
        
        ${{t}}_{{{latex(gs.n)}}} 
        = {{t}}_{{{latex(gs.n - 1)}}} + {r_latex}
        = {gs.get_term(gs.n - 1, to_latex = True)} + {r_latex}
        = {gs.get_term(gs.n, to_latex = True)} $ 
        
        Therefore, the series is
        $\\displaystyle {series_str}$
        
        $\\displaystyle  = {computed_terms_str}$
        
        It can also be written as ${summtion_latex}$
        """

                
        result = {'terms': terms,
                  'series': series,
                  'series_str': series_str,
                  'computed_terms_str': computed_terms_str,
                  'summtion_latex': summtion_latex,
                  'soln_s': soln_s,
            }
        
        return result
    
    
    def get_series_sum_stmt(self):
        
        gs = self
        
        terms = gs.get_terms(to_latex =  True)
        terms = OrderedDict(terms)
        series = gs.get_series(terms, to_str = False)
        series_str = gs.get_series(terms, to_str = True)
        series_r_str = gs.get_series(terms, to_str = True, reversed_order = True)
        
        sum_str = gs.get_sum(to_latex = True)
        
        summtion_latex = gs.get_summtion_latex()
        
        summed_term = ueval(2 * gs.a) + ueval(gs.n - 1) * ueval(gs.r)
        sum_value = (2 * gs.a + (gs.n - 1) * gs.r) *  gs.n / 2
        
        sum_latex = latex_fraction(f"{latex(summed_term)} \\cdot {gs.n}", 2)
        
        soln_s = f"""
        $S = {series_str} \\qquad \\cdots (i) $ 
        
        Reverse the order on RHS
        
        $S = {series_r_str} \\qquad \\cdots (ii) $ 
        
        Add both $(i)$ and  $(ii)$. Note that each terms is equal to $
        {latex(summed_term)}$ and there are ${gs.n}$ terms. Hence, we can write:
        
        $2S = \\left( {latex(summed_term)} \\right) \\cdot {gs.n} $ 
        
        $\\displaystyle S = {sum_latex} $ 
        
        $\\displaystyle \\implies S = {latex(sum_value)}$
        """
                
        result = {'terms': terms,
                  'series': series,
                  'series_str': series_str,
                  'summtion_latex': summtion_latex,
                  'sum_str': sum_str,
                  'soln_s': soln_s,
            }
        
        return result

    def get_sum(self, n = None, to_latex = False):
        gs = self
        if n is None:
            n = self.n
            
        
        summed_term = ueval(2 * gs.a) + ueval(n - 1) * ueval(gs.r)
        sum_value = (2 * gs.a + (n - 1) * gs.r) *  n / 2
        
        sum_latex = latex_fraction(f"{latex(summed_term)} \\cdot {n}", 2)
        
        if not to_latex:
            return sum_value

        return sum_latex
        pass
    
    def get_series(self, terms, to_str = False, reversed_order = False):
        include_dots = self._include_dots(terms)
        
        terms = list(terms.values())
        if include_dots:
            terms.insert(-2, r"\cdots")
        
        s = [f"{t}" for t in terms]
        
        if reversed_order:
            s = s[::-1]
        
        if not to_str:
            return s
        return ' + '.join(s)
        pass
    
    
    def _include_dots(self, terms):
        keys =  list(terms.keys())
        if len(keys) != 5:
            return False
        try:
            if keys[-3] + 1 != keys[-2]:
                return True
            else:
                return False
        except:
            pass
        return True
        pass
    
    def get_var_latex(self, v):
        if len(f"{v}") == 1:
            v_latex = f"{v}"
            return v_latex
        
        v_latex = f"\\left({latex(v)}\\right)"
        if isinstance(v, int):
            if v > 0:
                v_latex = f"{v}"
            pass

        return v_latex  
    
    pass
#class ArithmeticSeries:
