'''
Created on 01-Apr-2021
@author: bthakur
'''

import random
import inspect
import traceback

from sympy import (symbols, Symbol, latex)
from sympy.core.expr import UnevaluatedExpr
import numpy as np


from xv.km.util import latex_fraction
from xv.util import toIntOrFloat
from xv.km.ui import SvgHelper 


from ..helper import BaseHelper
from xv.km.managers import BaseManager
from xv.km.kelements import KProblemTemplate


# _laplace_plotter_manager
# LaplacePlotterManager

class LaplacePlotterManager(BaseManager):
    '''
    LaplacePlotterManager
    Example:
        ke = LaplacePlotterManager()
        ke.getRandomProblem()
        ke.getRandomProblem(problem_type = 0)
        ...
        
        ke.printProblem()
        ke.printAnswer()
        ke.printSolution()
        
        ke.printProblemTypes()
        
    '''    
    
    def __new__(cls, *args, **kwargs):
        instance = super(LaplacePlotterManager, cls).__new__(cls, *args, **kwargs)
        for base_class in LaplacePlotterManager.__bases__:
            base_class.__init__(instance, *args, **kwargs)
        return instance
    
    def __init__(self, *args, **kwargs):
        self.base_helper = BaseHelper(**kwargs)
        pass
    
    def _get_function_code(self, function):
        text = inspect.getsource(function)
        lines = text.rstrip().split("\n")
        lines = '\n'.join(lines[1:])
        return lines.strip()
        pass

    def _plot_laplace_function(self):
        import matplotlib.pyplot as plt
        import numpy as np
        
        def f(s, t):
            laplace = np.exp(s*t)
            return laplace.real, laplace.imag
        
        steps = 100
        theta = np.linspace(0, 10*np.pi, steps)
        s = 0.1 + 5j
        t = np.linspace(0, 10, steps)
        x, y = f(s, t)

        origin = np.array([0, 0, 0])

        def colormap(angle):
            color = np.array([0, 0, 0])
            if 0 < angle < np.pi/2:
                ratio = np.abs(angle/(np.pi/2))
                red = np.array([1, 0, 0])
                yellow = np.array([1, 1, 0]) * (ratio)
                blue = np.array([0, 0, 1]) * (1-ratio)
                color = np.mean([red, yellow, blue], axis=0)

            elif np.pi/2 < angle < np.pi:
                ratio = np.abs((angle - np.pi/2)/(np.pi/2))
                red = np.array([1, 0, 0]) * (1-ratio)
                yellow = np.array([1, 1, 0]) 
                green = np.array([0, 1, 0]) * (ratio)
                color = np.mean([red, yellow, green], axis=0)

            elif -np.pi < angle < -np.pi/2:
                ratio = np.abs((angle + np.pi/2)/(np.pi/2))
                yellow = np.array([1, 1, 0]) * (ratio)
                blue = np.array([0, 0, 1]) * (1-ratio)
                green = np.array([0, 1, 0])
                color = np.mean([yellow, blue, green], axis=0)
            
            elif -np.pi/2 < angle < 0:
                ratio = np.abs((angle)/(np.pi/2))
                red = np.array([1, 0, 0]) * (1-ratio)
                green = np.array([0, 1, 0]) * (ratio)
                blue = np.array([0, 0, 1])
                color = np.mean([red, green, blue], axis=0)
            
            else:
                color = 'k'
            return color    


        fig = plt.figure(figsize = (6, 6))
        ax = plt.axes(projection='3d')

        ax.plot(*origin, color = 'blue', marker = 'o')
        ax.text(origin[0] + 0.25, origin[1] + 0.25, origin[2] + 0.25, s = 'O', fontsize = 12)

        # ax.plot_surface(X, Y, Z, color= 'red', alpha = 0.5)

        for i in range(len(x)):
            angle = np.angle(complex(x[i], y[i]))
            color = colormap(angle)
            ax.plot(x[i], y[i], t[i], color = color, marker = 'o')

        ax.set_xlim(-10, 10)
        ax.set_ylim(-10, 10)
        ax.set_zlim(0, 10)
        plt.show()



    
    def _init_problem_types(self, *args, **kwargs):
        self._append_problem_types([
            '_problem_plot_laplace_function',
            
            ])
        pass

    def _problem_plot_laplace_function(self, *args, **kwargs):
        problemTemplate = KProblemTemplate()
        problemTemplate.problem_type = kwargs.get('problem_type')

        problem_s = f"""
        Plot a laplace function.
        """
        problemTemplate.add_problem(problem_s.strip(), 'text/html')
        plot_code = self._get_function_code(self._plot_laplace_function)
        
        soln_s = f"""
        <pre>
        {plot_code}
        </pre>
        """
        problemTemplate.add_solution(soln_s.strip(), 'text/html', 
        add_html_line_break = True)
        
        answer_s = f"""
        """
        problemTemplate.add_answer(answer_s.strip(), 'text/html')
        
        return problemTemplate
        
    
