from big_ol_pile_of_manim_imports import *


def p(x=0, y=0):
    if type(x) == list or type(x) == tuple or type(x) == np.ndarray:
        return np.array(x)
    return np.array((x, y, 0))


class MMScene(Scene):

    m1 = np.array([[1, 2], [2, 9]])
    m2 = np.array([[2, 4, 3], [6, 8, 3]])
    res_data = None
    x, y, res = None, None, None

    def construct(self):
        if self.m1.shape[1] != self.m2.shape[0]:
            raise ValueError("Invalid matrices")
        self.x = x = IntegerMatrix(self.m1)
        x.move_to(p(-0.1), RIGHT)
        self.y = y = IntegerMatrix(self.m2)
        y.move_to(p(0.1), LEFT)
        self.play(FadeIn(x), FadeIn(y))
        y = self.move_mobj(y, p(0.1, x.get_height() / 2 + 0.1), LEFT + DOWN)
        self.res_data = np.zeros((self.m1.shape[0], self.m2.shape[1]), dtype=int)
        self.res = res = IntegerMatrix(self.res_data)
        res.move_to(p(y.get_center()[0], x.get_center()[1]))
        self.play(FadeIn(res))
        self.circle_matrix(res, 0, 'green')
        self.wait(2)
        for i in range(self.res_data.shape[0]):
            for j in range(self.res_data.shape[1]):
                self.calc_res_elem(i, j)
        res = self.res
        self.wait(2)
        rc = self.res.copy()
        rc.move_to(ORIGIN)
        self.play(FadeOut(x), FadeOut(y), Transform(res, rc))
        self.remove(x, y, res)
        self.add(rc)
        self.res = res = rc
        self.wait(3)

    def circle_matrix(self, m: Matrix, elem: int, color='red'):
        e = m.get_entries()[elem]
        c = Circle(radius=max(e.get_width(), e.get_height()) / 2 + 0.2, color=color)
        c.move_to(e)
        return c

    def calc_res_elem(self, x, y):
        elem = self.res_data.shape[1] * x + y
        res_tex = TexMobject(f"{self.res_data[x, y]}")
        tex_pos = p(y=self.x.get_critical_point(DOWN)[1] - 0.1)
        res_tex.set_color(GREEN)
        res_tex.move_to(tex_pos, UP)
        c = self.circle_matrix(self.res, elem, 'green')
        self.play(ShowCreation(c))
        self.play(Write(res_tex))
        xcp, ycp = None, None
        for step in range(self.m1.shape[1]):
            x_elem = self.m1.shape[1] * x + step
            y_elem = self.m2.shape[1] * step + y
            xcp, ycp = self.calc_step(x_elem, y_elem, xcp, ycp)
            self.wait(0.5)
            a, b = self.m1[x, step], self.m2[step, y]
            dr = self.m1[x, step] * self.m2[step, y]
            self.res_data[x, y] += dr
            cur_res = res_tex.get_tex_string()
            rtm = TexMobject(f"_{{}}{cur_res}+{a}_{{}}\\times{b}^{{}}", tex_to_color_map={
                f'{a}_{{}}': RED,
                f'{b}^{{}}': YELLOW,
                f'_{{}}{cur_res}': GREEN
            })
            rtm.move_to(tex_pos, UP)
            rta = TexMobject(f"{cur_res}^{{}}+{dr}_{{}}", tex_to_color_map={
                f"{cur_res}^{{}}": GREEN,
                f"{dr}_{{}}": ORANGE
            })
            rta.move_to(tex_pos, UP)
            rt = TexMobject(f'{self.res_data[x, y]}')
            rt.set_color(GREEN)
            rt.move_to(tex_pos, UP)
            self.chain_transorm([res_tex, rtm, rta, rt], 0.5)
            res_tex = rt

            rm = IntegerMatrix(self.res_data)
            rm.move_to(self.res)
            cc = self.circle_matrix(rm, elem, 'green')
            self.play(Transform(self.res, rm), Transform(c, cc))
            self.remove(self.res, c)
            self.add(rm, cc)
            self.res = rm
            c = cc

        self.play(*[FadeOut(g) for g in [c, xcp, ycp, res_tex]])
        self.remove(c, xcp, ycp, res_tex)

    def chain_transorm(self, mobjs, wait_time, **kwargs):
        for a, b in zip(mobjs, mobjs[1:]):
            self.play(Transform(a, b, **kwargs))
            self.remove(a)
            self.add(b)
            self.wait(wait_time)

    def calc_step(self, x_elem, y_elem, xcp, ycp):

        if xcp is None or ycp is None:
            xc = self.circle_matrix(self.x, x_elem)
            yc = self.circle_matrix(self.y, y_elem, 'yellow')
            self.play(ShowCreation(xc), ShowCreation(yc))
        else:
            xc = xcp.copy()
            yc = ycp.copy()
            xc.move_to(self.x.get_entries()[x_elem])
            yc.move_to(self.y.get_entries()[y_elem])
            self.play(Transform(xcp, xc), Transform(ycp, yc))
            self.remove(xcp, ycp)
            self.add(xc, yc)
        return xc, yc

    def move_mobj(self, m, point, edge=ORIGIN, **kwargs):
        mc = m.copy()
        mc.move_to(point, edge)
        self.play(Transform(m, mc, lag_ratio=0, **kwargs))
        self.remove(m)
        self.add(mc)
        return mc