'''
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


# _plotter3d_manager
# Plotter3DManager

class Plotter3DManager(BaseManager):
    '''
    Plotter3DManager
    Example:
        ke = Plotter3DManager()
        ke.getRandomProblem()
        ke.getRandomProblem(problem_type = 0)
        ...

        ke.printProblem()
        ke.printAnswer()
        ke.printSolution()

        ke.printProblemTypes()

    '''

    def __new__(cls, *args, **kwargs):
        instance = super(Plotter3DManager, cls).__new__(cls, *args, **kwargs)
        for base_class in Plotter3DManager.__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_cylinder(self):
        import matplotlib.pyplot as plt
        import numpy as np

        radius = 3
        length = 5
        sides = 20
        z = np.linspace(-length/2, length/2, sides + 1)
        theta = np.linspace(0, 2*np.pi, sides + 1)
        theta_grid, Zc = np.meshgrid(theta, z)

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

        Xc = radius*np.cos(theta_grid) + center[0]
        Yc = radius*np.sin(theta_grid) + center[1]
        Zc = Zc + center[2]

        def plot_circle(center, theta):
            for i in range(len(theta) - 1):
                x1 = np.array([
                    [0, radius*np.cos(theta[i])],
                    [0, radius*np.cos(theta[i + 1])]
                ]) + center[0]
                y1 = np.array([
                    [0, radius*np.sin(theta[i])],
                    [0, radius*np.sin(theta[i + 1])]
                ]) + center[1]
                z1 = np.zeros(y1.shape) + center[2]

                ax.plot_surface(x1, y1, z1, color = 'red', alpha = 0.5)

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

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

        plot_circle([0, 0, length/2], theta)
        plot_circle([0, 0, -length/2], theta)

        ax.plot_surface(Xc, Yc, Zc, color = 'blue', alpha = 0.5)

        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_zlim(-5, 5)

        plt.show()

    def _plot_cone(self):
        import matplotlib.pyplot as plt
        import numpy as np

        radius = 3
        length = 5
        sides = 20
        theta = np.linspace(0, 2*np.pi, sides + 1)
        center = np.array([0, 0, 0])

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

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

        for i in range(len(theta) - 1):
            x1 = np.array([
                [0, radius*np.cos(theta[i])],
                [0, radius*np.cos(theta[i + 1])]
            ])
            y1 = np.array([
                [0, radius*np.sin(theta[i])],
                [0, radius*np.sin(theta[i + 1])]
            ])
            z1 = np.zeros(y1.shape)  + length
            zc = np.array([
                [0, length],
                [0, length]
            ])
            ax.plot_surface(x1, y1, zc, color = 'blue', alpha = 0.5)
            ax.plot_surface(x1, y1, z1, color = 'red', alpha = 0.5)

        # ax.plot_surface(Xc, Yc, Zc, color = 'blue', alpha = 0.5)

        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_zlim(0, 5)

        plt.show()

    def _plot_sphere(self):
        import matplotlib.pyplot as plt
        import numpy as np

        radius = 3
        theta = np.linspace(0, 2*np.pi, 100)
        alpha = np.linspace(0, np.pi, 100)

        x = radius*np.outer(np.cos(theta), np.sin(alpha))
        y = radius*np.outer(np.sin(theta), np.sin(alpha))
        z = radius*np.outer(np.ones(theta.shape), np.cos(alpha))

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

        ax.plot_surface(x, y, z, color = 'blue')
        # ax.set_xlim([-1.2, 1.2])
        # ax.set_ylim([-1.2, 1.2])
        # ax.set_zlim([-1.2, 1.2])

        plt.show()

    def _plot_rotated_sphere(self):
        import matplotlib.pyplot as plt
        import numpy as np

        radius = 3
        sides = 3
        theta = np.linspace(0, 2*np.pi, 50)
        center = np.array([0, 0, 0])
        alpha = np.linspace(0, np.pi, sides + 1)

        def rotate_circles(x1, y1, z1, rotation_matrix_x):
            for j in range(y1.shape[0]):
                for k in range(y1.shape[1]):
                    vec = np.array([x1[j, k], y1[j, k], z1[j, k]])
                    vec_r = rotation_matrix_x.dot(vec)
                    x1[j, k], y1[j, k], z1[j, k] = vec_r[0], vec_r[1], vec_r[2]
            return x1, y1, z1

        def plot_circle(theta, rotation_matrix_x):
            for i in range(len(theta) - 1):
                x1 = np.array([
                    [0, radius*np.cos(theta[i])],
                    [0, radius*np.cos(theta[i + 1])]
                ])
                y1 = np.array([
                    [0, radius*np.sin(theta[i])],
                    [0, radius*np.sin(theta[i + 1])]
                ])
                z1 = np.zeros(y1.shape)

                rotate_circles(x1, y1, z1, rotation_matrix_x)

                ax.plot_surface(x1, y1, z1, color = 'red', alpha = 0.5)

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

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

        for i in range(len(alpha) - 1):
            rotation_matrix_x = np.array([
                [1, 0, 0],
                [0, np.cos(alpha[i]), -np.sin(alpha[i])],
                [0, np.sin(alpha[i]), np.cos(alpha[i])],
            ])
            plot_circle(theta, rotation_matrix_x)

        # ax.plot_surface(Xc, Yc, Zc, color = 'blue', alpha = 0.5)

        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_zlim(-5, 5)

        plt.show()

    def _plot_sliced_sphere(self):
        import matplotlib.pyplot as plt
        import numpy as np

        radius = 3
        shapes = 10
        zs = np.linspace(-radius, radius, shapes + 2)
        theta = np.linspace(0, 2*np.pi, 50)
        center = np.array([0, 0, 0])
        # alpha = np.linspace(0, np.pi, sides + 1)


        def plot_circle(theta):
            for z in zs:
                for i in range(len(theta) - 1):
                    x1 = np.array([
                        [0, np.sqrt(radius**2 - z**2)*np.cos(theta[i])],
                        [0, np.sqrt(radius**2 - z**2)*np.cos(theta[i + 1])]
                    ])
                    y1 = np.array([
                        [0, np.sqrt(radius**2 - z**2)*np.sin(theta[i])],
                        [0, np.sqrt(radius**2 - z**2)*np.sin(theta[i + 1])]
                    ])

                    z1 = np.zeros(y1.shape) + z

                    ax.plot_surface(x1, y1, z1, color = 'red', alpha = 0.5)

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

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

        plot_circle(theta)

        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_zlim(-5, 5)


    def _init_problem_types(self, *args, **kwargs):
        self._append_problem_types([
            '_problem_plot_cylinder',
            '_problem_plot_cone',
            '_problem_plot_sphere',
            '_problem_plot_rotated_sphere',
            '_problem_plot_sliced_sphere',
            ])
        pass


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

        problem_s = f"""
        Plot a cylinder.
        """
        problemTemplate.add_problem(problem_s.strip(), 'text/html')
        plot_code = self._get_function_code(self._plot_cylinder)

        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

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

        problem_s = f"""
        Plot a cone.
        """
        problemTemplate.add_problem(problem_s.strip(), 'text/html')

        plot_code = self._get_function_code(self._plot_cone)

        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

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

        problem_s = f"""
        Plot a sphere.
        """
        problemTemplate.add_problem(problem_s.strip(), 'text/html')

        plot_code = self._get_function_code(self._plot_sphere)

        soln_s = f"""
        <pre>
        {plot_code}
        </pre>
        """
        problemTemplate.add_solution(soln_s.strip(), 'text/html',
        add_html_line_break = True)

        soln_s = f"""
        import matplotlib.pyplot as plt
        import numpy as np

        theta = np.linspace(0, 2*np.pi, 100)
        alpha = np.linspace(0, np.pi, 100)

        x = np.outer(np.cos(theta), np.sin(alpha))
        y = np.outer(np.sin(theta), np.sin(alpha))
        z = np.outer(np.ones(theta.shape), np.cos(alpha))

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

        ax.plot_surface(x, y, z, color = 'blue')
        ax.set_xlim([-1.2, 1.2])
        ax.set_ylim([-1.2, 1.2])
        ax.set_zlim([-1.2, 1.2])

        """
        problemTemplate.add_solution(soln_s.strip(), 'text/html')

        answer_s = f"""
        """
        problemTemplate.add_answer(answer_s.strip(), 'text/html')

        return problemTemplate

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

        problem_s = f"""
        Plot a rotated sphere.
        """
        problemTemplate.add_problem(problem_s.strip(), 'text/html')

        plot_code = self._get_function_code(self._plot_rotated_sphere)

        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

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

        problem_s = f"""
        Plot a sliced sphere.
        """
        problemTemplate.add_problem(problem_s.strip(), 'text/html')

        plot_code = self._get_function_code(self._plot_sliced_sphere)

        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



