from manim import *
import math

def norm(vec):
    return math.sqrt(sum([xi ** 2 for xi in vec]))

def normalize(vec):
    n = norm(vec)
    return [xi / n for xi in vec]

def plus(v1, v2):
    return [x1 + x2 for x1, x2 in zip(v1, v2)]

def minus(v1, v2):
    return [x1 - x2 for x1, x2 in zip(v1, v2)]

def mult(v, c):
    return [xi * c for xi in v]

def SegmentedLine(points):
    assert(len(points) >= 2)
    ret = []
    for p1, p2 in zip( points[:-1], points[1:] ):
        vec = minus(p2, p1)
        length = norm(vec)
        normalized = normalize(vec)
        rotated = [-normalized[1], normalized[0], 0]

        ret.append(p1)
        ret.append(plus(p1, mult(normalized, length * 1/3)))
        ret.append(plus(p1, plus(mult(normalized, length/2), mult(rotated, math.sqrt(3) * length / 6))))
        ret.append(plus(p1, mult(normalized, length * 2/3)))
    ret.append(points[-1])
    return ret

class snowflake(Scene):
    def construct(self):

        points = [[-7, -1, 0],
                  [ 7, -1, 0]]

        segment = VMobject(stroke_width=1)
        segment.set_points_as_corners(points)

        self.add(*[Dot(p) for p in points])

        self.play(Create(segment))
        self.wait(1)

        vg = VGroup()
        self.add(vg)
        for i in range(5):
            points = SegmentedLine(points)

            new_vg = VGroup(*[Dot(p, radius=0.1/(i+1)) for p in points])
            self.play(Create(new_vg))

            self.play(segment.animate.set_points_as_corners(points))

            self.wait(1)

            vg.add(*[e for e in new_vg])
            self.remove(new_vg)

        self.play(Uncreate(vg))
        self.play(*[FadeOut(mob) for mob in self.mobjects])


if __name__ == "__main__":
    sc = snowflake()
    sc.render(True)