from consts import *
from components import *

import numpy as np
from manim import *


class Introduction(TitledScene):
    def construct(self):
        self.add_title("Goldreich--Goldwasser--Halevi")
        self.wait()

        text_1 = Tex(
            r"""
            \begin{itemize} 
                \item Lattice-based cryptosystem.
                \item Devised in 1997 by Goldreich, Goldwasser, and Halevi.
                \item Broken in 1999 by Nguyen.
            \end{itemize}
            """,
            font_size=MEDIUM_FONT,
        )
        self.add(text_1)
        self.play(Write(text_1, run_time=4.0))
        self.wait()


class Premise(TitledScene):
    def construct(self):
        self.add_title("Lattices")

        # A lattice is a subspace of a vector space that is constructed by taking integer multiples of some basis
        # vectors.
        # For example, take the real plane R2.
        plane = NumberPlane(axis_config={"stroke_width": 0.0})
        plane.set_z_index(-10)
        plane.set_opacity(0.75)
        dot = Dot(ORIGIN)
        self.add(dot, plane)
        self.play(Create(dot), Create(plane))
        self.wait()

        # We can construct the 2D grid of integers with the elementary basis
        lattice_1 = VGroup()
        arrow_1 = Arrow(ORIGIN, [1, 0, 0], buff=0)
        arrow_2 = Arrow(ORIGIN, [0, 1, 0], buff=0)

        self.play(Create(arrow_1), Create(arrow_2))
        self.wait()

        for i in range(-50, 50):
            for j in range(-25, 25):
                lattice_1.add(Dot([i, j, 0]))

        self.play(Create(lattice_1))

        # By moving these basis vectors but maintaining their linear independency, other lattices can be formed
        self.play(
            Transform(arrow_1, Arrow(ORIGIN, [1.5, 0, 0], buff=0)),
            *[
                Transform(
                    dot,
                    Dot(
                        dot.get_center()
                        * np.matrix([[1.5, 0, 0], [0, 1, 0], [0, 0, 1]])
                    ),
                )
                for dot in lattice_1
            ]
        )
        self.wait()

        self.play(
            Transform(arrow_1, Arrow(ORIGIN, [1, -0.5, 0], buff=0)),
            *[
                Transform(
                    dot,
                    Dot(
                        dot.get_center()
                        * np.matrix([[2 / 3, 0, 0], [0, 1, 0], [0, 0, 1]])
                        * np.matrix([[1, -0.5, 0], [0, 1, 0], [0, 0, 1]])
                    ),
                )
                for dot in lattice_1
            ]
        )
        self.wait()

        self.play(
            Transform(arrow_1, Arrow(ORIGIN, [1, -1, 0], buff=0)),
            Transform(arrow_2, Arrow(ORIGIN, [1, 1, 0], buff=0)),
            *[
                Transform(
                    dot,
                    Dot(
                        dot.get_center()
                        * np.matrix([[1, 0.5, 0], [0, 1, 0], [0, 0, 1]])
                        * np.matrix([[1, -1, 0], [1, 1, 0], [0, 0, 1]])
                    ),
                )
                for dot in lattice_1
            ]
        )
        self.wait()

        self.play(
            Transform(arrow_1, Arrow(ORIGIN, [2, 0, 0], buff=0)),
            *[
                Transform(
                    dot,
                    Dot(
                        dot.get_center()
                        * np.matrix([[1 / 2, 1 / 2, 0], [-1 / 2, 1 / 2, 0], [0, 0, 1]])
                        * np.matrix([[2, 0, 0], [1, 1, 0], [0, 0, 1]])
                    ),
                )
                for dot in lattice_1
            ]
        )
        self.wait()


class MatrixRep(TitledScene):
    def construct(self):
        self.add_title("Lattices")

        matrix_comp = MathTex(
            r"\begin{bmatrix}1 & 1\\ 1 & -1\end{bmatrix} \sim \begin{bmatrix}1 & 2\\ 1 & 0\end{bmatrix}"
        )
        lattice_comp = MathTex(
            r"L\left(\begin{bmatrix}1 & 1\\ 1 & -1\end{bmatrix}\right) = L\left(\begin{bmatrix}1 & 2\\ 1 & 0\end{bmatrix}\right)"
        )

        self.play(Create(matrix_comp))
        self.wait()

        self.play(Transform(matrix_comp, lattice_comp))
        self.wait()

        self.play(ApplyMethod(matrix_comp.shift, 2 * UP))
        self.wait()

        l_def = MathTex(
            r"L(\begin{bmatrix}\mathbf{b_1} & \mathbf{b_2}\end{bmatrix}) = \{ k_1 \mathbf{b_1} + k_2 \mathbf{b_2} \mid k_i \in \mathbb{Z} \}"
        )

        l_def.shift(DOWN)

        self.play(Create(l_def))
        self.wait()


class LatticeProblems(Scene):
    def construct(self):
        text = Tex("How does this give us a trapdoor?", font_size=LARGE_FONT)

        self.play(Create(text))
        self.wait()

        self.play(Transform(text, Tex("Lattice problems", font_size=LARGE_FONT)))
        self.wait()

        self.play(ApplyMethod(text.to_corner, UP + LEFT))
        self.wait()


class OrthoDefect(Scene):
    def construct(self):
        text = Tex("How do we decide which bases are ``good''?", font_size=LARGE_FONT)

        self.play(Create(text))
        self.wait()

        self.play(Transform(text, Tex("Orthogonal defect", font_size=LARGE_FONT)))
        self.wait()

        self.play(ApplyMethod(text.to_corner, UP + LEFT))
        self.wait()

        diagram_1 = VGroup()
        diagram_1.add(Dot(ORIGIN))
        diagram_1.add(Arrow(ORIGIN, [1, 0, 0], buff=0))
        diagram_1.add(Arrow(ORIGIN, [0, 1, 0], buff=0))

        diagram_1.shift(2 * LEFT)
        diagram_1.set_color(GREEN)
        label_1 = Tex("Defect 1", font_size=MEDIUM_FONT)
        label_1.next_to(diagram_1, DOWN)

        diagram_2 = VGroup()
        diagram_2.add(Dot(ORIGIN))
        diagram_2.add(Arrow(ORIGIN, [0.25, 1, 0], buff=0))
        diagram_2.add(Arrow(ORIGIN, [0.5, 1, 0], buff=0))

        diagram_2.shift(2 * RIGHT)
        diagram_2.set_color(RED)
        label_2 = Tex(r"Defect $\approx$ 4.6", font_size=MEDIUM_FONT)
        label_2.next_to(diagram_2, DOWN)

        self.play(Create(diagram_1), Create(diagram_2))
        self.wait()

        self.play(Create(label_1), Create(label_2))
        self.wait()