'''
Created on 20-Jan-2022

@author: bthakur
'''

import codecs
import inspect
import traceback
import random
import numpy as np
import pandas as pd
from sympy import (latex, Symbol, symbols,
                   Matrix, Identity, nsimplify,
                    Rational,
                    sin, cos, tan,
                    asin, acos, atan,
                    sinh, cosh, tanh,
                    log, exp, sqrt,
                   diff, factor, 
                   Function, Eq,
                   Add, 
                   simplify, solve, factorial,
                   UnevaluatedExpr as ueval
                   )
from sympy import Sum, Indexed, pi, oo, Integral
import sympy
from sympy.core.numbers import One, Integer
from sympy.series import series

from xv.km.util import latex_fraction


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

import matplotlib.pyplot as plt
from xv.util import CategorizedSampleData, list_to_str
from xv.km.util._latex import latex_color


class ComplexFourierHelper:
    
    def __init__(self, **kwargs):
        self.verbose = kwargs.get('verbose', False)
        pass
    
    
    def get_fourier_expression(self):
        '''
        https://medium.com/@kamil2000budaqov/fourier-series-fourier-transform-ae748fcd424a
        
        https://www.bragitoff.com/2021/05/fourier-series-coefficients-and-visualization-python-program/
        
        http://192.168.1.208:8080/user/thakur/notebooks/published-docs/python-language/sympy_1_9/057-002-Concrete_-_SymPy_1-9_documentation.ipynb
        
        '''
        
        from sympy import Sum, Indexed, pi, oo, Integral
        
        
        problem_s = f"""
        Write expression for Complex Fourier series
        """
        
        soln_items = []
        
        t = symbols('t')
        f = Function('f')(t)
        a0 = Symbol("a_0")

        T0 = Symbol("T_0", positive = True)
        omega0_expr = 2 * pi / T0
        omega0 = Symbol(latex(omega0_expr))

        m, n = symbols('m, n', integer = True, positive = True)

        n_theta = omega0  * n * t

        cos_expr = ueval(Indexed('a',n))  * ueval(cos(n_theta))
        sin_expr = ueval(Indexed('b',n))  * ueval(sin(n_theta))

        term_lhs = f
        
        term0 = a0       
        cos_term = Sum(cos_expr, (n, 1, oo))
        sin_term = Sum(sin_expr, (n, 1, oo))
        term_rhs = term0 + cos_term + sin_term

        eq0 = Eq(term_lhs, term_rhs)
        
        soln_item = f"""
        Fourier series is defined as 
        
        $\\displaystyle {latex(eq0)} $
        """
        
        soln_items.append(soln_item.strip())


        soln_items.append("<h3>Part I: Express in complex form</h3>")
        stmt_items, eq_complex = self._get_complex_form(f, m, n, t, a0, T0)
        soln_items.extend(stmt_items)


        soln_items.append("<h3>Part II: Final results</h3>")
        stmt_item = f"""
        $\\displaystyle {latex(eq0)} $

        $\\displaystyle \\implies {latex(eq_complex)} $

        """
        
        soln_items.append(stmt_item.strip())

        
        soln_s = "\n\n".join(soln_items)
        
        answer_s = f"""
        """
        
        result = {
            'problem_s': problem_s.strip(),
            'soln_s': soln_s.strip(),
            'answer_s': answer_s.strip(),
            }
        return result
        pass
    
    def _get_complex_form(self, f, m, n, t, a0, T0):
        soln_items = []


        omega0_expr = 2 * pi / T0
        omega0 = Symbol(latex(omega0_expr))

        n_theta = omega0  * n * t


        cos_expr = ueval(Indexed('a',n))  * ueval(cos(n_theta))
        sin_expr = ueval(Indexed('b',n))  * ueval(sin(n_theta))

        term_lhs = f
        
        term0 = a0       

        
        cos_term = Sum(cos_expr, (n, 1, oo))
        sin_term = Sum(sin_expr, (n, 1, oo))
        term_rhs = term0 + cos_term + sin_term

        eq0 = Eq(term_lhs, term_rhs)

        soln_item = f"""
        $\\displaystyle {latex(eq0)} $
        """
        soln_items.append(soln_item.strip())

        z = exp(sympy.I * n_theta)
        z_expr = Symbol(f"""
        {latex(cos(n_theta))} + {latex(sympy.I * sin(n_theta))}
        """)
        
        z_bar = exp(- sympy.I * n_theta)
        z_bar_expr = Symbol(f"""
        {latex(cos(n_theta))} - {latex(sympy.I * sin(n_theta))}
        """)
        
        z_eq = Eq(z, z_expr)
        z_bar_eq = Eq(z_bar, z_bar_expr)
        
        cos_complex = Symbol(latex_fraction(f"{latex(z_bar)} + {latex(z)}", 2))
        sin_complex = Symbol(latex_fraction(f"{latex(z_bar)} - {latex(z)}", 2))
        
        cos_eq = Eq(cos(n_theta), cos_complex)
        sin_eq = Eq(sin(n_theta), sin_complex)
        
        
        term0 = a0     
        cos_term = Sum(ueval(Indexed('a',n)) * ueval(cos_complex), (n, 1, oo))
        sin_term = Sum(ueval(Indexed('b',n)) * ueval(sin_complex), (n, 1, oo))
        term_rhs = term0 + cos_term + sin_term
        term_rhs_symb = Symbol(f"""
        {latex(term0)} + {latex(cos_term)} + {latex(sin_term)}
        """)

        eq_complex0 = Eq(term_lhs, term_rhs_symb)
        
        soln_item = f"""
        $\\displaystyle {latex(z_eq)} $
        
        $\\displaystyle {latex(z_bar_eq)} $
        
        Hence:
        
        $\\displaystyle {latex(cos_eq)} $
        
        $\\displaystyle {latex(sin_eq)} $
        
        $\\displaystyle \\implies {latex(eq_complex0)} $
        
        """
        
        soln_items.append(soln_item.strip())
        
        
        term_rhs_symb = Symbol(f"""
        {latex(term0)} + {latex(cos_term)} + {latex(sin_term)}
        """)

        complex_term0 = term0
        coeff1 = Indexed('a',n) + sympy.I * Indexed('b',n)
        coeff1 = Symbol(f"""
        {latex_fraction(latex(coeff1), 2)}
        """)

        coeff2 = Indexed('a',n) - sympy.I * Indexed('b',n)
        coeff2 = Symbol(f"""
        {latex_fraction(latex(coeff2), 2)}
        """)
        
        complex_term1 = Sum(ueval(coeff1) * ueval(z_bar), (n, 1, oo))
        complex_term2 = Sum(ueval(coeff2) * ueval(z), (n, 1, oo))
        complex_term_rhs = complex_term0 + complex_term1 + complex_term2
        eq_complex2 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        soln_item = f"""
        $\\displaystyle \\implies {latex(eq_complex2)} $
        """
        
        soln_items.append(soln_item.strip())
        
        c_coeff = Indexed('c',n)
        c_coeff_eq = Eq(c_coeff, coeff1)
        
        c_coeff_bar = Indexed(f'\\bar{{c}}',n)
        c_coeff_bar_eq = Eq(c_coeff_bar, coeff2)
        
        soln_item = f"""
        Let: 
        
        $\\displaystyle {latex(c_coeff_eq)} $
        
        $\\displaystyle {latex(c_coeff_bar_eq)} $
        
        """
        soln_items.append(soln_item.strip())
        

        complex_term1 = Sum(ueval(c_coeff_eq.lhs) * ueval(z_bar), (n, 1, oo))
        complex_term2 = Sum(ueval(c_coeff_bar_eq.lhs) * ueval(z), (n, 1, oo))

        complex_term_rhs = Symbol(f"""
        {latex(complex_term0)} + {latex(complex_term1)} + {latex(complex_term2)} 
        """)
       
        eq_complex3 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        soln_item = f"""
        Then, we can rewrite: 
        
        $\\displaystyle {latex(eq_complex3)} $
        
        """
        soln_items.append(soln_item.strip())
        
        k = Symbol('k')
        
        c_min_k = Indexed(f'\\bar{{c}}', -k)
        z_subs_k = z.subs(n, -k)
        
        complex_term2 = Sum(ueval(c_min_k) * ueval(z_subs_k), (k, -1, -oo))

        complex_term_rhs = Symbol(f"""
        {latex(complex_term0)} + {latex(complex_term1)} + {latex(complex_term2)} 
        """)
       
        eq_complex4 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        
        soln_item = f"""
        Let us use a dummy variable ${latex(k)}$ on the last term on the rhs.
        In the last term on rhs, replacing ${latex(n)}$ with ${latex(-k)}$, we get:
        
        $\\displaystyle {latex(eq_complex4)} $
        
        """
        soln_items.append(soln_item.strip())
        
        
        c_neg_coeff = Indexed(f'\\bar{{c}}', -n)
        z_subs_k = z_subs_k.subs(k, n)
        
        complex_term2 = Sum(ueval(c_neg_coeff) * ueval(z_subs_k), (n, -1, -oo))

        complex_term_rhs = Symbol(f"""
        {latex(complex_term0)} + {latex(complex_term1)} + {latex(complex_term2)} 
        """)
       
        eq_complex5 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        
        soln_item = f"""
        Let us replace the dummy variable ${latex(k)}$ by ${latex(n)}$.
        As it the index of summation, the change of variable will not change the value.
        
        $\\displaystyle {latex(eq_complex5)} $
        
        """
        soln_items.append(soln_item.strip())
        
        
        
        complex_term2 = Sum(ueval(c_coeff) * ueval(z_subs_k), (n, -1, -oo))

        complex_term_rhs = Symbol(f"""
        {latex(complex_term0)} + {latex(complex_term1)} + {latex(complex_term2)} 
        """)
       
        eq_complex6 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        
        soln_item = f"""
        If we define ${latex(c_neg_coeff)}$ such that 
        $\\displaystyle {latex(c_neg_coeff)} = {latex(c_coeff)} 
        $ for ${latex(n)} = -1, -2, -3 \\cdots $
        
        $\\displaystyle {latex(eq_complex6)} $
        
        """
        soln_items.append(soln_item.strip())
        
        complex_term2 = Sum(ueval(c_coeff) * ueval(z_subs_k), (n, -oo, -1))
        
        complex_term_rhs = Symbol(f"""
        {latex(complex_term2)}  + {latex(complex_term0)} + {latex(complex_term1)} 
        """)
       
        eq_complex7 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        
        soln_item = f"""
        Let us rearrange the order of the last term on the rhs, rearrange rhs terms:
        
        $\\displaystyle {latex(eq_complex7)} $
        
        The rhs terms have range from ${latex(-oo)}$ to ${latex(oo)}$, except zero.
        
        """
        soln_items.append(soln_item.strip())
        

        c0_coeff = Indexed('c', 0)

        zero_symb = Symbol('0')
        zero_theta = ueval(omega0)  * ueval(zero_symb) * ueval(t)
        zero_z = exp(sympy.I * zero_theta)

        complex_term0_redef = ueval(c0_coeff) * ueval(zero_z)
        
        complex_term_rhs = Symbol(f"""
        {latex(complex_term2)}  + {latex(complex_term0_redef)} + {latex(complex_term1)} 
        """)
       
        eq_complex8 = Eq(eq_complex0.lhs, complex_term_rhs)
        
        soln_item = f"""
        Let us rewrite the middle term as:
        
        $\\displaystyle 
        {latex(complex_term0)} 
        = {latex(c0_coeff)}
        = {latex(c0_coeff)}{latex(zero_z)}
        $
        
        Then, 
        
        $\\displaystyle {latex(eq_complex8)} $
        
        """
        soln_items.append(soln_item.strip())
        
        complex_rhs = Sum(ueval(c_coeff) * ueval(z_bar), (n, -oo, oo))
        eq_complex9 = Eq(eq_complex0.lhs, complex_rhs)
        
        soln_item = f"""
        The rhs terms have range from ${latex(-oo)}$ to ${latex(oo)}$.
        So, we can rewrite the rhs as:
        
        $\\displaystyle {latex(eq_complex9)} $
        
        """
        soln_items.append(soln_item.strip())
        
        
        return soln_items, eq_complex9
        pass
    
    
    def _expr_to_symbol(self, expr, color = None, evaluate = False):
        symb = Symbol(latex(expr))
        if color is None:
            if evaluate:
                return symb
            return ueval(symb)
        
        colored_symb = Symbol(latex_color(latex(symb), color))
        if evaluate:
            return colored_symb
        return ueval(colored_symb)
                

    def get_wave(self, A, f, theta, phi):
        x = theta
        y = A*np.sin(2*np.pi*f*theta - phi)
        return x, y

    def get_circle(self, A, f, theta, phi):
        x = A*np.cos(2*np.pi*f*theta - phi)
        y = A*np.sin(2*np.pi*f*theta - phi)
        return x, y

    def plot_wave(self, data, time):
        A = data[0]
        f = data[1]
        phi = data[2]
        x, y = self.get_wave(A, f, time, phi)

        ax = plt.gca()
        plt.plot(x, y, color = 'blue')

        ax.set_xlabel('Time', fontsize = 20, color = 'red')
        ax.tick_params(axis='x', colors='red')
        ax.tick_params(axis='y', colors='red')

    def plot_multiple_waves(self, data, time):
        no_of_plots = len(data)
        fig, axes = plt.subplots(no_of_plots, figsize = (16, 9))
        for i in range(no_of_plots):
            ax = axes[i]
            w_data = data[i] 
            A = w_data[0]
            f = w_data[1]
            phi = w_data[2]
            x, y = self.get_wave(A, f, time, phi)
            ax.plot(x, y, color = 'blue', alpha = 1)
            ax.set_ylabel(f'Wave {i + 1}', color = 'blue', fontsize = 15)

            fig.suptitle('Time', fontsize = 20, color = 'red')
            ax.tick_params(axis='x', colors='red')
            ax.tick_params(axis='y', colors='red')
            ax.set_ylim(-15, 15)

    def plot_a_fourier_waves(self, data, time):
        no_of_plots = len(data)
        ys = np.zeros(len(time))
        for i in range(no_of_plots):
    #         ax = axes[i]
            w_data = data[i] 
            A = w_data[0]
            f = w_data[1]
            phi = w_data[2]
            x, y = self.get_wave(A, f, time, phi)
            ys += y
            
        fig = plt.figure()
        ax = plt.gca()
        x = time
        plt.plot(x, y, color = 'blue', alpha = 1)
        ax.set_ylabel(f'Wave {i + 1}', color = 'blue', fontsize = 15)

        fig.suptitle('Time', fontsize = 20, color = 'red')
        ax.tick_params(axis='x', colors='red')
        ax.tick_params(axis='y', colors='red')
        ax.set_ylim(-15, 15)

    def plot_fourier_waves(self, data, time, alpha = 1):
        no_of_plots = len(data) + 2
        fig, axes = plt.subplots(no_of_plots, figsize = (16, 9))
        xs = np.zeros(len(time))
        ys = np.zeros(len(time))
        for i in range(no_of_plots):
            ax = axes[i]
            if i < no_of_plots - 2:
                w_data = data[i] 
                A = w_data[0]
                f = w_data[1]
                phi = w_data[2]
                x, y = self.get_wave(A, f, time, phi)
                ax.plot(x, y, color = 'blue', alpha = alpha)
                xs += x
                ys += y
                ax.set_ylabel(f'Wave {i + 1}', color = 'blue', fontsize = 15)

            elif i == no_of_plots - 2:
                for j in range(no_of_plots-2):
                    w1_data = data[j]
                    A = w1_data[0]
                    f = w1_data[1]
                    phi = w1_data[2]
                    x1, y1 = self.get_wave(A, f, time, phi)
                    ax.plot(x1, y1, color = 'green', alpha = alpha)
                    ax.set_ylabel(f'All waves', color = 'green', fontsize = 15)

            else:
                ax.plot(time, ys, color = 'red', alpha = alpha)
                ax.set_ylabel(f'Fourier Wave', color = 'red', fontsize = 15)
                
                
            ax.set_xlabel('Time', fontsize = 20, color = 'red')
            ax.tick_params(axis='x', colors='red')
            ax.tick_params(axis='y', colors='red')
            ax.set_ylim(-15, 15)

    def plot_circular_wave(self, data, time):
        A = data[0]
        f = data[1]
        phi = data[2]
        x, y = self.get_circle(A, f, time, phi)
        
        plt.figure(figsize = (13, 6))
        ax = plt.gca()
        plt.plot(x, y, color = 'blue')

        ax.set_xlabel('Time', fontsize = 20, color = 'red')
        ax.tick_params(axis='x', colors='red')
        ax.tick_params(axis='y', colors='red')
        ax.set_aspect('equal')    
        
    def plot_multiple_circular_waves(self, data, time):
        no_of_plots = len(data)
        fig, axes = plt.subplots(ncols = no_of_plots, figsize = (16, 9))
        for i in range(no_of_plots):
            ax = axes[i]
            w_data = data[i] 
            A = w_data[0]
            f = w_data[1]
            phi = w_data[2]
            x, y = self.get_circle(A, f, time, phi)
            ax.plot(x, y, color = 'blue')           
                
            ax.set_xlim(-12, 12)
            ax.set_ylim(-12, 12)
            ax.set_aspect('equal')

            ax.set_title(f'Wave {i + 1}', fontsize = 20, color = 'red')
            ax.tick_params(axis='x', colors='red')
            ax.tick_params(axis='y', colors='red')

    def plot_circular_a_fourier_wave(self, data, time, alpha = 1):
        x = 0
        y = 0
        for A, f, phi in data:
            x_temp, y_temp = self.get_circle(A, f, time, phi)
            x += x_temp
            y += y_temp
        plt.plot(x, y, color = 'red', alpha = alpha)
        plt.gca().set_title(f'Fourier Wave', color = 'blue', fontsize = 20)
        plt.gca().set_aspect('equal')

    def plot_circular_fourier_wave(self, data, time, alpha = 1):
        no_of_plots = len(data) + 2
        fig, axes = plt.subplots(ncols = no_of_plots, figsize = (16, 9))
        xs = np.zeros(len(time))
        ys = np.zeros(len(time))
        for i in range(no_of_plots):
            ax = axes[i]
            if i < no_of_plots - 2:
                w_data = data[i] 
                A = w_data[0]
                f = w_data[1]
                phi = w_data[2]
                x, y = self.get_circle(A, f, time, phi)
                ax.plot(x, y, color = 'blue', alpha = alpha)
                xs += x
                ys += y
                ax.set_title(f'Wave {i + 1}', color = 'blue', fontsize = 20)
            elif i == no_of_plots - 2:
                for j in range(no_of_plots-2):
                    w1_data = data[j]
                    A = w1_data[0]
                    f = w1_data[1]
                    phi = w1_data[2]
                    x1, y1 = self.get_circle(A, f, time, phi)
                    ax.plot(x1, y1, color = 'green', alpha = alpha)
                ax.set_title(f'All waves', color = 'green', fontsize = 20)

            else:
                ax.plot(xs, ys, color = 'red', alpha = alpha)
                ax.set_title(f'Fourier wave', color = 'red', fontsize = 20)
                

            ax.set_xlim(-15, 15)
            ax.set_ylim(-15, 15)
            ax.tick_params(axis='x', colors='red')
            ax.tick_params(axis='y', colors='red')
            ax.set_aspect('equal')
            
        pass
    
    pass
#class ComplexFourierHelper:
