#_____________________________________________________________________
#
# Schramm–Loewner evolution
#
# Greg Zitelli, 2016
#_____________________________________________________________________

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib as mpl
import matplotlib.pyplot as plt

##################################################

# Push function
def push(z,delta,Delta):
    l=np.sqrt(np.power(z,2)-4*Delta+0*1j)
    if l.imag<0:
        l=-l
    return l+delta

def try_this(gamma,n,delta,Delta):
    b = 0
    for m in range(1,n+1):
        b = push(b,delta.ix[n-m+1],Delta.ix[n-m+1])
    return b

def create_delta(tk):
    Delta = pd.Series(index=tk.index[1:])
    for l in Delta.index:
        Delta.ix[l] = tk.ix[l]-tk.ix[l-1]
    return Delta

def fixing_function(gamma,delta,tk,fix):
    delta.ix[0] = 0
    delta = delta.sort_index().cumsum()
    for n in range(len(fix)):
        g = pd.concat([tk,delta],axis=1)
        gg=pd.DataFrame(g.ix[fix[n]-1:fix[n]].mean(),columns=[len(g.index)]).transpose()
        gg.ix[:,1]=gg.ix[:,1]+0.5*np.sqrt(k*g.ix[fix[n]-1:fix[n]].diff().tail(1)[0].tolist()[0])*np.random.normal(0,1)
        g = pd.concat([g,gg]).sort_values(0)
        g.index = range(len(g.index))
        tk = g.ix[:,0]
        delta = g.ix[:,1]
    delta = delta.diff()[1:]
    return [delta,tk]


# Set up
k = 6
T = 1
N = 1000
dt = T/N
tk = pd.Series(index=range(N+1)) # The clock
for n in tk.index:
    tk.ix[n] = np.power(T*n*dt,1)

# Define the little and big delta vectors initially
delta = pd.Series(index=tk.index[1:])
Delta = create_delta(tk)
for n in delta.index:
    delta.ix[n] = np.random.normal(0,1)*np.sqrt(k*Delta.ix[n])

# Run the experiment and save data
Delta = create_delta(tk)
gamma = pd.Series(index=tk.index)
gamma.ix[0] = 0
eps = 0.005
n = 1
# This loop can be interrupted and restarted in order to check progress
while n<=(len(tk)-1):
    ttry = 0
    while np.abs(try_this(gamma,n,delta,Delta)-gamma.ix[n-1])>eps:
        ttry = ttry + 1
        fix = [n]
        [delta,tk] = fixing_function(gamma,delta,tk,fix)
        Delta = create_delta(tk)
        print('trying... ' + str(ttry))
    gamma.ix[n] = try_this(gamma,n,delta,Delta)
    print('done! ' +str(n/len(tk))+', '+str(tk.ix[n]))
    n = n+1
    gamma.to_pickle('gamma6.pickle')


# Create the image
plt.clf()
gamma = gamma.dropna()
gamma2 = gamma.apply(lambda t:1j*(t-1j)/(t+1j))
plt.plot(gamma2.real,gamma2.imag,'r-')
plt.xlim((-1,1))
plt.ylim((-1,1))
plt.xticks(())
plt.yticks(())
#plt.show()
fig = plt.gcf()
plt.gca().set_aspect('equal', adjustable='box')
fig.savefig('gamma6.png',transparent=False,dpi=150,bbox_inches='tight')