from numpy import zeros 
from numpy import mean 
from numpy import column_stack
from numpy import vstack
from numpy import concatenate
from scipy.stats import norm    
from esig import stream2sig
from esig import sigkeys
from matplotlib.pylab import plot
from sklearn.linear_model import LogisticRegression

#ARMA simulation
def SimARMA(n,k,c,y0,phi,sigma=1.0):
    if n<=k:
        return "error: n is smaller than k"
    y=zeros(n)
    for i in range(0,k):
        y[i] = y0[i]+sigma*norm.rvs()
    for i in range(k,n):
        y[i] = c+sigma*norm.rvs()
        for j in range(0,k):
            y[i] += phi[j]*y[i-j-1]
    return y

def Leadtransform(m,x):
    y=zeros(2*m)
    y[0] = x[0]
    y[1] = x[1]
    y[2*m-1] = x[m-1]
    for i in range(1,m):
        y[2*i] = x[i]
        y[2*i-1] = x[i]
    return y

def Lagtransform(m,x):
    y=zeros(2*m)
    y[2*m-1] = x[m-1]
    for i in range(0,m):
        y[2*i] = x[i]
        y[2*i+1] = x[i]
    return y

# Routine implementing an ARMA simulation
n = 50 #dimension of the sample
m = 300 #length of the stream
k = 2 #degree
y0 = [1,2] #initial values of the series

depth = 3 #Degree for the signature
c = 0 #centering

#Model one
phi_1 = [0.1,0.1] #weights
sample_1 = SimARMA(m,k,c,y0,phi_1,sigma=1) #initialization of the stream
lead_1 = Leadtransform(m,sample_1)
lag_1 = Lagtransform(m,sample_1)
leadlag_1 = column_stack([lead_1,lag_1])
sig_1 = stream2sig(leadlag_1, depth) # compute the signature of the steam

#Model two
phi_2 = [0.3,0.4] #weights
sample_2 = SimARMA(m,k,c,y0,phi_2,sigma=1) #initialization of the stream
lead_2 = Leadtransform(m,sample_2)
lag_2 = Lagtransform(m,sample_2)
leadlag_2 = column_stack([lead_2,lag_2])
sig_2 = stream2sig(leadlag_2, depth) # compute the signature of the steam

# Plotting the time series
dom = list(range(m))
#plot(dom, sample_1, '-', color='blue')
#plot(dom, sample_2, '-', color='red')
#poner varias graficas y poner line thickness bajo

#iterated simulations of ARMA model 1 and extraction of the signature
for i in range(1,n):
    stream_1 = SimARMA(m,k,c,y0,phi_1,sigma=1)
    plot(dom, stream_1, '-', color='blue',linewidth=0.1)
    lead_1 = Leadtransform(m,stream_1)
    lag_1 = Lagtransform(m,stream_1)
    leadlag_1 = column_stack([lead_1,lag_1])
    aux = stream2sig(leadlag_1, depth)
    sig_1 = vstack([sig_1,aux])
    
    stream_2 = SimARMA(m,k,c,y0,phi_2,sigma=1)
    plot(dom, stream_2, '-', color='red',linewidth=0.1)
    lead_2 = Leadtransform(m,stream_2)
    lag_2 = Lagtransform(m,stream_2)
    leadlag_2 = column_stack([lead_2,lag_2])
    aux = stream2sig(leadlag_2, depth)
    sig_2 = vstack([sig_2,aux])

#labels of the signature
sig_keys = sigkeys(2, depth)
print(sig_keys)

# Matrix XX and YY
XX = zeros((2*n,12))
YY = concatenate(([0]*n,[1]*n))
for j in range(0,12-1):
        XX[2*n-1,j] = sig_2[n-1,j]
for i in range(0,n-1):
    for j in range(0,12-1):
        XX[i,j] = sig_1[i,j]
        XX[n+i,j] = sig_2[i,j]

# Logistic Regression
LogReg =  LogisticRegression(penalty='l1', solver='liblinear')
#LogReg =  LogisticRegression(penalty='l2', solver='liblinear')
LogReg.fit(XX,YY)


# Testing
#New iterated simulations of ARMA model 1 and extraction of the signature
for i in range(1,n):
    stream_1 = SimARMA(m,k,c,y0,phi_1,sigma=1)
    lead_1 = Leadtransform(m,stream_1)
    lag_1 = Lagtransform(m,stream_1)
    leadlag_1 = column_stack([lead_1,lag_1])
    aux = stream2sig(leadlag_1, depth)
    sig_1 = vstack([sig_1,aux])
    
    stream_2 = SimARMA(m,k,c,y0,phi_2,sigma=1)
    lead_2 = Leadtransform(m,stream_2)
    lag_2 = Lagtransform(m,stream_2)
    leadlag_2 = column_stack([lead_2,lag_2])
    aux = stream2sig(leadlag_2, depth)
    sig_2 = vstack([sig_2,aux])
    
# New Matrix XX and YY
XX = zeros((2*n,12))
YY = concatenate(([0]*n,[1]*n))
for j in range(0,12-1):
        XX[2*n-1,j] = sig_2[n-1,j]
for i in range(0,n-1):
    for j in range(0,12-1):
        XX[i,j] = sig_1[i,j]
        XX[n+i,j] = sig_2[i,j]

# Counting successes
test =  [-1]*(2*n)
for i in range(0,2*n-1):
    if LogReg.predict([XX[i,:]])==YY[i]:
        test[i] = 1
    if LogReg.predict([XX[i,:]])!=YY[i]:
            test[i] = 0
me = mean(test)

#comparing sampled signatures
#x_1 = sig_1[:,1]
#x_2 = sig_1[:,2]
#x_3 = sig_1[:,3]
#x_4 = sig_1[:,4]
#x_5 = sig_1[:,5]
#x_6 = sig_1[:,6]
#x_7 = sig_1[:,7]
#x_8 = sig_1[:,8]
#x_9 = sig_1[:,9]
#x_10 = sig_1[:,10]
#x_11 = sig_1[:,11]
#x_12 = sig_1[:,12]
#x_13 = sig_1[:,13]
#x_14 = sig_1[:,14]


#a_1 = sig_2[:,4]
#b_1 = sig_2[:,5]

#plot( x_1, y_1, '.', color='blue')
#plot( a_1, b_1, '.', color='red')