'''
Created on 12-Jul-2021

@author: rishi_priyansh
'''
import random
import inspect
import traceback
import re
import operator
import datetime
import calendar
from time import time, ctime

from sympy import (symbols, latex, simplify, nsimplify, 
                   Rational,
                   Eq,
                   pi,
                   sin, cos, tan,
                   csc, sec, cot
                   )
from sympy.abc import theta
from sympy.physics import units

from xv.km.util import latex_fraction, latex_exponent, wrap_bracket
from xv.math.util import LatexEquation
from xv.km.ui import SvgHelper         
import numpy as np
import matplotlib.pyplot as plt

class TheoryOfRelativityHelper:
    def __init__(self, b = None, e = None, count = 2, **kwargs):
        self.verbose = kwargs.get('verbose', False)
        self.implies = symbols(u'\U000021D2')
        
        self._colors = ['indigo', 'blue', 'green', 'red', 'cyan', 'magenta', 
                        'violet', 'black', 'orange', 'black', 'grey']  
        self._markers = ['o', 'v', '^', '<', '>', '1', '2', '3', '4', 
                         '8', 's', 'p', 'P', '*', 'h', 'H', 'X', 'D', 'd', 'o']
        self._p_labels = 'ABCDEFGHIJKLMNPQRSTUVWXYZ'

        self._center_label = 'O'
        self._center = (0, 0)

        self._default_font_size = 10
        self._default_fig_size = 6
        self._axis_multiplier = 1.2
        pass

    def show_method_source_code(self, method_name):
        try:
            self_method = getattr(self, method_name)
            text = inspect.getsource(self_method)
            lines = text.split("\n")
            spaces = len(lines[0]) - len(lines[0].lstrip())
            
            lines = [line[spaces:] for line in lines]
            lines[0] = f'def {method_name}():'
            return "\n".join(lines)
            
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}:
            {e}"""
            print(sError)
            print(traceback.format_exc())
        
        return None            
        
    def get_mirror_center(self, center_of_mirrors, d_val):
        center_of_mirror_up = np.array([center_of_mirrors[0], center_of_mirrors[1] + (d_val/2)])
        center_of_mirror_down = np.array([center_of_mirrors[0], center_of_mirrors[1] - (d_val/2)])
        return center_of_mirror_up, center_of_mirror_down
        
    def get_mirrors(self, center_of_mirrors, length, d_val):
        center_of_mirror_up, center_of_mirror_down = self.get_mirror_center(center_of_mirrors, d_val)
        mirror_up = np.array([[center_of_mirror_up[0] - (length/2), center_of_mirror_up[1]],
                    [center_of_mirror_up[0] + (length/2), center_of_mirror_up[1]]])
    
        mirror_down = np.array([[center_of_mirror_down[0] - (length/2), center_of_mirror_down[1]],
                      [center_of_mirror_down[0] + (length/2), center_of_mirror_down[1]]])
        
        return mirror_up, mirror_down
            
    def plot_static_mirror(self, length = None, d_val = None, center_of_mirrors = None):
        
        if not length:
            length = 4
        if not d_val:
            d_val = 5
        if not center_of_mirrors:
            center_of_mirrors = np.array([0, 0])
        
        mirror_up, mirror_down = self.get_mirrors(center_of_mirrors, length, d_val)
        center_of_mirror_up, center_of_mirror_down = self.get_mirror_center(center_of_mirrors, d_val)

        d_arrow1 = np.array([[center_of_mirror_up[0] + 0.1, center_of_mirror_up[1]],
                        [0, -d_val]])

        d_arrow2 = np.array([[center_of_mirror_down[0] - 0.1, center_of_mirror_down[1]],
                        [0, d_val]])

        fig = plt.figure(figsize = (6, 6))
        plt.plot(mirror_up[:, 0], mirror_up[:, 1], color = 'black')
        plt.plot(mirror_down[:, 0], mirror_down[:, 1], color = 'black')

        plt.arrow(*d_arrow1.flatten(), length_includes_head = True, width = 0.05, color = 'red')
        plt.arrow(*d_arrow2.flatten(), length_includes_head = True, width = 0.05, color = 'red')

        plt.text(center_of_mirrors[0] + 0.2, center_of_mirrors[1], s = 'd', fontsize = 16)
        plt.gca().set_aspect('equal')
        return fig
        
    def plot_moving_mirror(self, length = None, d_val = None, center_of_mirrors1 = None, center_of_mirrors2 = None, center_of_mirrors3 = None):
        if not length:
            length = 4
        if not d_val:
            d_val = 5
        if not center_of_mirrors1:
            center_of_mirrors1 = np.array([0, 0])
        mirror1_up, mirror1_down = self.get_mirrors(center_of_mirrors1, length, d_val)
        center_of_mirror1_up, center_of_mirror1_down = self.get_mirror_center(center_of_mirrors1, d_val)
        
        if not center_of_mirrors2:
            center_of_mirrors2 = np.array([10, 0])
        mirror2_up, mirror2_down = self.get_mirrors(center_of_mirrors2, length, d_val)
        center_of_mirror2_up, center_of_mirror2_down = self.get_mirror_center(center_of_mirrors2, d_val)
        
        if not center_of_mirrors3:
            center_of_mirrors3 = np.array([20, 0])
        mirror3_up, mirror3_down = self.get_mirrors(center_of_mirrors3, length, d_val)
        center_of_mirror3_up, center_of_mirror3_down = self.get_mirror_center(center_of_mirrors3, d_val)
        
        L_arrow1 = np.array([center_of_mirror1_down, center_of_mirror2_up - center_of_mirror1_down ])
        L_arrow2 = np.array([center_of_mirror2_up, center_of_mirror3_down - center_of_mirror2_up ])
        
        mid_point_mirror_1_3_down = (center_of_mirror1_down + center_of_mirror3_down)/2
        u_t_arrow1 = np.array([mid_point_mirror_1_3_down - np.array([0, 1.5]), center_of_mirror1_down - mid_point_mirror_1_3_down])
        u_t_arrow2 = np.array([mid_point_mirror_1_3_down - np.array([0, 1.5]), center_of_mirror3_down - mid_point_mirror_1_3_down])
        
        fig  = plt.figure(figsize = (10, 10))
        plt.plot(mirror1_up[:, 0], mirror1_up[:, 1], color = 'black')
        plt.plot(mirror1_down[:, 0], mirror1_down[:, 1], color = 'black')
        
        plt.plot(mirror2_up[:, 0], mirror2_up[:, 1], color = 'black')
        plt.plot(mirror2_down[:, 0], mirror2_down[:, 1], color = 'black')
        
        plt.plot(mirror3_up[:, 0], mirror3_up[:, 1], color = 'black')
        plt.plot(mirror3_down[:, 0], mirror3_down[:, 1], color = 'black')
        
        plt.arrow(*L_arrow1.flatten(), length_includes_head = True, width = 0.1, color = 'red')
        plt.arrow(*L_arrow2.flatten(), length_includes_head = True, width = 0.1, color = 'red')
        
        plt.text((L_arrow1[0] + (L_arrow1[1]/2))[0] - 1, (L_arrow1[0] + (L_arrow1[1]/2))[1], s = 'L', fontsize = 16, ha = 'center', va = 'center')
        plt.text((L_arrow2[0] + (L_arrow2[1]/2))[0] + 1, (L_arrow2[0] + (L_arrow2[1]/2))[1], s = 'L', fontsize = 16, ha = 'center', va = 'center')
        
        plt.arrow(*u_t_arrow1.flatten(), length_includes_head = True, width = 0.1, color = 'blue')
        plt.arrow(*u_t_arrow2.flatten(), length_includes_head = True, width = 0.1, color = 'blue')
        
        plt.text(mid_point_mirror_1_3_down[0], mid_point_mirror_1_3_down[1] - 1, s = fr'$u t$', fontsize = 16, ha = 'center', va = 'center', color = 'blue')
        
        plt.plot([center_of_mirror2_up[0], center_of_mirror2_down[0]], [center_of_mirror2_up[1], center_of_mirror2_down[1]], color = 'green')
        plt.text(center_of_mirrors2[0] + 0.5, center_of_mirrors2[1], s = 'd', fontsize = 16, ha = 'center', va = 'center')
        
        plt.gca().set_aspect('equal')

        return fig
        
    def plot_speed_and_time(self, t_0_val = None, c_val = None, u_val_dots = None):
        if not t_0_val:
            t_0_val = 10
        if not c_val:
            c_val = 1
        u_val = np.linspace(0, c_val - 10**(-2), 100)
        gamma_val = 1 / np.sqrt(1 - (u_val/c_val)**2)
        t_val = gamma_val*t_0_val 
        
        if not u_val_dots:
            u_val_dots = np.array([0, 0.25, 0.5, 0.6, 0.87, 0.95])
        gamma_val_dots = 1 / np.sqrt(1 - (u_val_dots/c_val)**2)
        t_val_dots = gamma_val_dots*t_0_val
        fig = plt.figure(figsize = (6, 6))
        plt.plot(u_val, t_val, color = 'blue')
        plt.scatter(u_val_dots, t_val_dots, color = 'red')
        
        for i in range(len(u_val_dots)):
            x_offset = 0
            if i == 3:
                x_offset = 0.1
                
            plt.text(u_val_dots[i] + x_offset, t_val_dots[i] + 3, s= f'({u_val_dots[i]}, {round(t_val_dots[i], 2)})', ha = 'right', va = 'center')
        
        plt.xlabel('speed of object(expressed in terms of c)', fontsize = 14)
        plt.ylabel('time delay', fontsize = 14) 
        return fig
    
    def plot_speed_and_gamma(self, c_val = None, u_val_dots = None):
        if not c_val:
            c_val = 1
        u_val = np.linspace(0, c_val - 10**(-2), 100)
        gamma_val = 1 / np.sqrt(1 - (u_val/c_val)**2)
        
        if not u_val_dots:
            u_val_dots = np.array([0, 0.25, 0.5, 0.6, 0.87, 0.95])
        gamma_val_dots = 1 / np.sqrt(1 - (u_val_dots/c_val)**2)
        fig = plt.figure(figsize = (6, 6))
        plt.plot(u_val, gamma_val, color = 'blue')
        plt.scatter(u_val_dots, gamma_val_dots, color = 'red')
        
        for i in range(len(u_val_dots)):
            x_offset = 0
            if i == 3:
                x_offset = 0.1
                
            plt.text(u_val_dots[i] + x_offset, gamma_val_dots[i] + 0.3, s= f'({u_val_dots[i]}, {round(gamma_val_dots[i], 2)})', ha = 'right', va = 'center')
        
        plt.xlabel('speed of object(expressed in terms of c)', fontsize = 14)
        plt.ylabel('gamma factor', fontsize = 14)
        return fig
    
    def plot_gamma_and_time(self, t_0_val = None, c_val = None, u_val_dots = None):
        if not t_0_val:
            t_0_val = 10
        if not c_val:
            c_val = 1
        u_val = np.linspace(0, c_val - 10**(-2), 100)
        gamma_val = 1 / np.sqrt(1 - (u_val/c_val)**2)
        t_val = gamma_val*t_0_val 

        if not u_val_dots:
            u_val_dots = np.array([0, 0.25, 0.5, 0.6, 0.87, 0.95])
        gamma_val_dots = 1 / np.sqrt(1 - (u_val_dots/c_val)**2)
        t_val_dots = gamma_val_dots*t_0_val 
        fig = plt.figure(figsize = (6, 6))
        plt.plot(gamma_val, t_val, color = 'blue')
        plt.scatter(gamma_val_dots, t_val_dots, color = 'red')

        for i in range(len(gamma_val_dots)):
            x_offset = 0
            y_offset = 1.5
            if i == 3 or i == 1:
                x_offset = 0.8
                y_offset = 0
            if i == 0 or i == 2:
                x_offset = -0.1
                y_offset = 0
    
            plt.text(gamma_val_dots[i] + x_offset, t_val_dots[i] + y_offset, s= f'({round(gamma_val_dots[i], 2)}, {round(t_val_dots[i], 2)})', ha = 'right', va = 'center')

        plt.xlabel('gamma factor', fontsize = 14)
        plt.ylabel('time delay', fontsize = 14)

        plt.xlim(0, 4)
        plt.ylim(9, 40)
        return fig 

    def plot_speed_and_length(self, l_0_val = None, c_val = None, u_val_dots = None):
        if not l_0_val:
            l_0_val = 10
        if not c_val:
            c_val = 1
        l_0_val = 10
        c_val = 1
        u_val = np.linspace(0, c_val - 10**(-2), 100)
        gamma_val = 1 / np.sqrt(1 - (u_val/c_val)**2)
        l_val = l_0_val/gamma_val

        if not u_val_dots:
            u_val_dots = np.array([0, 0.25, 0.5, 0.6, 0.87, 0.95])
        gamma_val_dots = 1 / np.sqrt(1 - (u_val_dots/c_val)**2)
        l_val_dots = l_0_val / gamma_val_dots
        fig = plt.figure(figsize = (6, 6))
        plt.plot(u_val, l_val, color = 'blue')
        plt.scatter(u_val_dots, l_val_dots, color = 'red')

        for i in range(len(u_val_dots)):
            x_offset = 0.05
            y_offset = 0
            if i == 0:
                x_offset = 0
                y_offset = -0.5
                
            plt.text(u_val_dots[i] + x_offset, l_val_dots[i] + y_offset, s= f'({round(u_val_dots[i], 2)}, {round(l_val_dots[i], 2)})', ha = 'left', va = 'center')

        plt.xlabel('speed of object(expressed in terms of c)', fontsize = 14)
        plt.ylabel('length contraction', fontsize = 14)
        return fig
    
    def plot_gamma_and_length(self, l_0_val = None, c_val = None, u_val_dots = None):
        if not l_0_val:
            l_0_val = 10
        if not c_val:
            c_val = 1
        u_val = np.linspace(0, c_val - 10**(-2), 100)
        gamma_val = 1 / np.sqrt(1 - (u_val/c_val)**2)
        l_val = l_0_val/gamma_val

        if not u_val_dots:
            u_val_dots = np.array([0, 0.25, 0.5, 0.6, 0.87, 0.95])
        gamma_val_dots = 1 / np.sqrt(1 - (u_val_dots/c_val)**2)
        l_val_dots = l_0_val / gamma_val_dots

        fig = plt.figure(figsize = (6, 6))
        plt.plot(gamma_val, l_val, color = 'blue')
        plt.scatter(gamma_val_dots, l_val_dots, color = 'red')

        for i in range(len(gamma_val_dots)):
            x_offset = 0.2
            y_offset = 0
            if i == 1:
                y_offset = -0.3
                
            plt.text(gamma_val_dots[i] + x_offset, l_val_dots[i] + y_offset, s= f'({round(gamma_val_dots[i], 2)}, {round(l_val_dots[i], 2)})', ha = 'left', va = 'center')

        plt.xlabel('gamma factor', fontsize = 14)
        plt.ylabel('length contraction', fontsize = 14)
        # plt.gca().set_aspect('equal')

        return fig 


    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 TheoryOfRelativityHelper:
