from manim import *
import numpy as np
import random

class newVectorScene(ThreeDScene):
    def construct(self):

        self.camera.set_focal_distance(50)
        amplitudeAxes = ThreeDAxes(
                                    x_range = [-1,1,1],
                                    y_range = [-1,1,1],
                                    z_range = [-1,1,1],
                                    x_length = 3.5,
                                    y_length = 3.5,
                                    z_length = 3.5,
                                    tips=False,
                                    ).shift(LEFT*5)

        amplitudeLabels = VGroup(
                                    MathTex(r'\uparrow').scale(1.2).next_to(amplitudeAxes[1], UP),
                                    MathTex(r'\downarrow').scale(1.2).next_to(amplitudeAxes[0], RIGHT),
                                 )

        shadeIn3D = False
        amplitudeVectorUp = Arrow(amplitudeAxes.c2p(0, 0, 0), amplitudeAxes.c2p(0, 1, 0), buff=0,
                                  color=BLUE, shade_in_3d=shadeIn3D)
        amplitudeVectorDown = Arrow(amplitudeAxes.c2p(0, 0, 0), amplitudeAxes.c2p(1, 0, 0), buff=0,
                                    color=ORANGE, shade_in_3d=shadeIn3D)

        amplitudeGroup = VGroup(amplitudeAxes,amplitudeLabels,amplitudeVectorUp,amplitudeVectorDown)
        amplitudeGroup.rotate(-TAU / 8, axis=np.array([0, np.cos(TAU / 12), np.sin(TAU / 12)])).scale(1.3)

        conceptualArrow = Arrow(start=LEFT, end=RIGHT).next_to(amplitudeGroup).next_to(amplitudeGroup).shift(DOWN*0.7)

        upText = MathTex(r'Up =').next_to(conceptualArrow).shift(UP*1.5).set_color(BLUE)

        upColumn = Matrix(
                            [[1],[0]],
                            left_bracket='(',
                            right_bracket=')',
                            bracket_h_buff=0.1
                          ).next_to(upText).shift(UP*0.095).set_color(BLUE)

        downText = MathTex(r'Down =').next_to(upColumn).shift(RIGHT*1.5).set_color(ORANGE)

        downColum = Matrix(
                            [[0],[1]],
                            left_bracket='(',
                            right_bracket=')',
                            bracket_h_buff=0.1
                           ).next_to(downText).shift(UP*0.095).set_color(ORANGE)

        newText = MathTex(r'New \; vector? =').next_to(conceptualArrow).shift(DOWN*1.5+RIGHT*1.2).set_color('#4E9843')

        newColumn = Matrix(
                            [['a'],['b']],
                            left_bracket='(',
                            right_bracket=')',
                            bracket_h_buff=0.1
                           ).next_to(newText).shift(UP*0.095).set_color('#4E9843')

        mathGroup = VGroup(upText,upColumn,downText,downColum,newText,newColumn)

        self.add(amplitudeGroup,conceptualArrow)
        self.wait()

        self.play(
                    ReplacementTransform(amplitudeVectorUp.copy(),mathGroup[0:2]),
                    ReplacementTransform(amplitudeVectorDown.copy(), mathGroup[2:4]),
                  )
        self.wait()

        amplitudeVectorI = Arrow(amplitudeAxes.c2p(0, 0, 0), amplitudeAxes.c2p(0, 0, 1), buff=0,
                                    color='#4E9843', shade_in_3d=shadeIn3D)

        self.play(
                    Indicate(mathGroup[4:6],color='#4E9843',scale_factor=1.5),
                    Indicate(amplitudeVectorI,color='#4E9843',scale_factor=1.5)
                  )
        self.wait()

        self.play( Rotate( mobject=amplitudeVectorI,angle=-TAU/4, about_point=amplitudeAxes.c2p(0, 0, 0),
                           axis=amplitudeAxes.c2p(1, 0, 0)-amplitudeAxes.c2p(0, 0, 0) ) )
        self.wait()

        upDownPlane = Polygon(amplitudeAxes.c2p(-1,-1,0), amplitudeAxes.c2p(-1,1,0), amplitudeAxes.c2p(1,1,0),
                            amplitudeAxes.c2p(1,-1,0),fill_color=PURPLE,fill_opacity=0.5,stroke_width=0)

        self.play(FadeIn(upDownPlane))
        self.wait()