from big_ol_pile_of_manim_imports import *


class MainScene(Scene):
    SL = 3
    scale_matrix = np.identity(2) * SL
    matrix = np.array([
        [0, 1],
        [1, 0]
    ]).T

    def construct(self):
        plane = self.make_plane()
        x_label, t_label = self.make_labels(plane)

        box = self.make_box()
        self.wait(0.5)
        lines = self.make_lines()
        objects = [plane, *lines]
        x_label_t = x_label.copy().next_to((4, 0, 0), DL)
        self.play(*[ApplyMatrix(self.matrix, m, run_time=1) for m in objects] + [Transform(x_label, x_label_t)])
        self.wait(3)

    def make_plane(self):
        plane = NumberPlane(background_line_style={
            "stroke_opacity": 0.5
        })

        self.add(plane)
        self.play(ShowCreation(plane, run_time=0.5))
        return plane

    def make_labels(self, plane):
        x_label = plane.get_x_axis_label("x")
        t_label = plane.get_y_axis_label("t")
        self.add(x_label, t_label)
        self.play(FadeIn(x_label), FadeIn(t_label))
        return x_label, t_label

    def make_box(self):
        box = Square(side_length=1, color=GREY)
        box.scale(self.SL)
        box.move_to((self.SL / 2, self.SL / 2, 0))
        self.add(box)
        self.play(ShowCreation(box, run_time=0.5))
        return box

    def make_lines(self):
        lines_data = [
            {
                'color': GREEN,
                'coords': [(0, 0, 0), (1, 0.2, 0)]
            },
            {
                'color': GREEN,
                'coords': [(1, 0.2, 0), (0, 0.4, 0)]
            },
            {
                'color': GREEN,
                'coords': [(0, 0.4, 0), (1, 0.6, 0)]
            },
            {
                'color': GREEN,
                'coords': [(1, 0.6, 0), (0, 0.8, 0)]
            },
            {
                'color': GREEN,
                'coords': [(0, 0.8, 0), (1, 1, 0)]
            }
        ]
        lines = []
        for ld in lines_data:
            l = Line(*(np.array(ld['coords']) * self.SL), stroke_color=ld['color'])
            lines.append(l)
            self.add(l)
            self.play(ShowCreation(l, run_time=1 / len(lines_data)))

        return lines