# -*- coding: utf-8 -*-
"""
Created on Tue Aug 31 22:04:01 2021

@author: Julien

https://en.wikipedia.org/wiki/Sierpi%C5%84ski_triangle
"""

#  _____                            _       
# |_   _|                          | |      
#   | |  _ __ ___  _ __   ___  _ __| |_ ___ 
#   | | | '_ ` _ \| '_ \ / _ \| '__| __/ __|
#  _| |_| | | | | | |_) | (_) | |  | |_\__ \
# |_____|_| |_| |_| .__/ \___/|_|   \__|___/
#                 | |                       
#                 |_|

import cmath
import numpy as np

#   _____                _              _       
#  / ____|              | |            | |      
# | |     ___  _ __  ___| |_ __ _ _ __ | |_ ___ 
# | |    / _ \| '_ \/ __| __/ _` | '_ \| __/ __|
# | |___| (_) | | | \__ \ || (_| | | | | |_\__ \
#  \_____\___/|_| |_|___/\__\__,_|_| |_|\__|___/

X_LIMITS = (0, 1)
Y_LIMITS = (-0.1, np.sin(np.pi/3)+0.1)

#  ______                _   _                 
# |  ____|              | | (_)                
# | |__ _   _ _ __   ___| |_ _  ___  _ __  ___ 
# |  __| | | | '_ \ / __| __| |/ _ \| '_ \/ __|
# | |  | |_| | | | | (__| |_| | (_) | | | \__ \
# |_|   \__,_|_| |_|\___|\__|_|\___/|_| |_|___/

def compute_new(z1, z2, up):
    '''
    uses complex numbers to compute the next points to add for each edge.

    Parameters
    ----------
    z1 : complex
        first vertice.
    z2 : complex
        second vertice.
    up : boolean
        wether or not add the point in a clockwise or counterclockwise fashion.

    Returns
    -------
    a, b : complex, complex
        the new point(s) generated from the given edge.
    up : boolean
        returning "up" value for computing purposes
    '''
    v = (z2-z1)/2 #v is the 'vector' from z1 to the middle of the segment [z1, z2]
    
    #rotating the vectors by 60 degrees clokwise or counterclockwise depending on the up value
    #and on the starting point
    if up :
        a = z1 + v*1j**(2/3)
        b = z2 + v*1j**(4/3)
    else :
        a = z1 - v*1j**(4/3)
        b = z2 - v*1j**(2/3)
    
    up = not(up) #flipping the up value
    
    return a, b, up
    

def partial_sierpinski_triangle(T):
    '''
    Generates the next generation of the fractal curve algorithm. For each pair of consecutive points of T
    that make a vertice, we add between those two points the new point(s) generated by the "compute_new"
    method.    

    Parameters
    ----------
    T : list[complex]
        the current state of the algorithm.

    Returns
    -------
    res : list[complex]
        The next step of the algorithm.
    '''
    
    up = len(T)%10<3
    res = []
    for i in range(len(T)-1):
        res.append(T[i])
        a, b, up = compute_new(T[i], T[i+1], up)
        res.append(a)
        res.append(b)
    res.append(T[-1])
    return res

def sierpinski_triangle(n):
    '''
    Recursive method to build the final fractal curve of degree n.

    Parameters
    ----------
    n : int
        The degree of the fractal curve to build.

    Returns
    -------
    list[complex]
        The list of complex numbers that represent the n-th degree fractal curve.
    '''
    
    if n==0 :
        return [complex(0,0), complex(1,0)]
    else :
        return partial_sierpinski_triangle(sierpinski_triangle(n-1))

def pts_num(n):
    '''
    Optional function that returns the amount of points in the n-th generation of the algorithm.

    Parameters
    ----------
    n : int
        The degree of the fractal curve.

    Returns
    -------
    int
        The number of vertices needed to build the n-th degree fractal curve.
    '''
    
    n0 = 2
    return (3**n)*(n0-1)+1

#  _            _       
# | |          | |      
# | |_ ___  ___| |_ ___ 
# | __/ _ \/ __| __/ __|
# | ||  __/\__ \ |_\__ \
#  \__\___||___/\__|___/

if __name__ == '__main__' :
    import matplotlib.pyplot as plt
    
    N = 12
    T = sierpinski_triangle(N)
    Tr = [el.real for el in T]
    Ti = [el.imag for el in T]
    plt.plot(Tr, Ti, lw=np.exp(-N/10))
    plt.axis('equal')
    plt.show()
