import numpy as np
import matplotlib.pyplot as plt


def nextStep(A_list, length):
    """
    The function describing the curve.  It is this function that 
    must be changed if a different curve is needed (for example for 
    Koch's snowflake we would have
    
    for a in A_list:
        R.append(a)
        R.append(a - 60)
        R.append(a + 60)
        R.append(a)
    return R, length/3

    ).
    """
    R = []
    for k in range(len(A_list)):
        if k%2==0:
            a = A_list[k] + 45
            a = a%360
            R.append(a)
            R.append(a - 90)
        else:
            a = A_list[k] - 45
            a = a%360
            R.append(a)
            R.append(a + 90)
    return R, length/np.sqrt(2)


def anglesTo2dPoints(A_list, length):
    """
    Returns the poligonal line (a list of 2d_points) constructed 
    using the list of angles A_list and the length as length of 
    each segment.  The construction stars at the origin and each 
    angle (in degrees) is with respect to the Ox axis.
    """
    PL = []
    angle = 0
    point = np.array([0, 0])
    PL.append(point)
    for a in A_list:
        a = a*np.pi/180
        v = np.array([np.cos(a), np.sin(a)])
        point = point + length * v
        PL.append(point)
    return PL
    

    
def forPlotPL(PL_list):
    """
    PL_list is a list of 2d-points (vectors).  It returns the x and y 
    lists of the corresponding segments.
    """
    x = [p[0] for p in PL_list]
    y = [p[1] for p in PL_list]
    return x, y


plt.ion()
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, aspect="equal")
ax.set_xlim(-.5, 1.3)
ax.set_ylim(-.5, .8)

nbSteps = 11
A_list = [0]
length = 1

delta_t = 1
# proportionnel a l'intervalle de temps de vue de chaque image

### first image
PL = anglesTo2dPoints(A_list, length)
x, y = forPlotPL(PL)
dragon, = ax.plot(x, y, lw=.75)


### animation
for k in range(nbSteps):
    ax.set_title(f"dragon curve after {k+1} steps")
    A_list, length = nextStep(A_list, length)
    PL = anglesTo2dPoints(A_list, length)
    x, y = forPlotPL(PL)
    dragon.set_xdata(x)
    dragon.set_ydata(y)
    plt.pause(delta_t)
    plt.show()