#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
#from scipy.special import gamma, binom
import numpy as np
import matplotlib.pyplot as plt
import scipy.sparse as sparse
from scipy.sparse.linalg import eigsh
from numpy.linalg import eig

from mpmath import mp, mpf
mp.dps = 2000

voxel_num = 5
phase_range = np.pi/2 *0.0
phase_init = np.pi/2  #(0.0)
U_points = voxel_num * 3000

noise_ratio = 0.0*(1.2e-16) #(1.2e-16) #1e8

total_periods=10#DO NOT CHANGE to fit T2s_scale autoscale
rf_samples_per_period = 20

# max polynomial order equals  rf_samples_per_period * total_periods 

# B0=1.5T freq=64Mhz, period = 15.6 ns
period = (10.0) #ms
omega = 2.0*np.pi/period
#T2s_scale = 0.01 #ms # need to be 10ms
T2s_scale = total_periods*period/rf_samples_per_period/voxel_num*8 #ms # need to be 10ms
T2s_min = T2s_scale/20.0 
#ms
total_steps = rf_samples_per_period*total_periods
if total_steps % 2 != 0: total_steps -= 1
time_steps = np.array(mp.linspace((0), (period*total_periods), total_steps))
tmp = [np.random.rand()*0.0+(1.0) for n in range(voxel_num)]
voxel_amplitudes = np.array(tmp)
tmp = [np.random.rand() for n in range(voxel_num)]
voxel_T2s_decay = np.sort(np.array(tmp))*(T2s_scale-2*T2s_min) + T2s_min

voxel_all = np.append(voxel_amplitudes,voxel_T2s_decay/T2s_scale)
if voxel_num == 5:
#    voxel_all = np.array([mpf(0.2),(0.6),(0.02),(0.1)])
    voxel_all = np.array([(0.822628),(0.691376),(0.282906),(0.226013),(0.90703),(0.144985),(0.228563),(0.340353),(0.462462),(0.720518)])
    #voxel_all = np.array([(0.592606),(0.135168),(0.365712),(0.667536),(0.437378),(0.918822),(0.943879),(0.590338),(0.685997),(0.658839)])
    voxel_amplitudes = voxel_all[:voxel_num]
    voxel_T2s_decay = voxel_all[voxel_num:]*T2s_scale

# a_i = np.array([(0.3),(0.1),(0.15),(0.1)])
# d_i = np.array([(0.7),(0.9),(0.2),(0.67)])
# voxel_num = len(a_i)


voxel_phases = np.array(np.linspace(0,phase_range, voxel_num))
# if len(voxel_amplitudes) != len(voxel_phases):
#     print("ERROR! Size of amplitude and phase arrays do not match!")
#     raise

ampl = []
def gen_rf_signal(time):
    '''Generates demodulated signal at radio frequence using voxels
    amplitudes, T2s decays, and phases.

    '''
    tmp = [mpf(0.0) for n in range(len(time))]
    mag_sin = np.array(tmp)
    mag_cos = np.array(tmp)    
    for t in range(len(time)):
        for i in range(voxel_num):
            # mag_sin[t] += a_i[i]*(
            #     (d_i[i] + np.random.rand()*noise_ratio)**time[t]
            # )
            amp = voxel_amplitudes[i] * (
                mp.exp(-time[t]/voxel_T2s_decay[i])
                ) + ( 0.0 
                     #     + np.random.rand()*noise_ratio
                        )
            # if t == 0:
                #print("a_{:d}".format(i),float(voxel_amplitudes[i]* np.sin(voxel_phases[i] + phase_init)))
                #print("d_{:d}".format(i),float( np.exp(-(period/rf_samples_per_period)/voxel_T2s_decay[i]) ))
            mag_sin[t] += amp * np.sin((voxel_phases[i] + phase_init))
            mag_cos[t] += amp * np.cos((voxel_phases[i] + phase_init))
        mag_sin[t] *= (1 + (np.random.rand()-0.5)*noise_ratio)
        #mag_cos[t] += np.random.rand()*noise_ratio
    return mag_sin, mag_cos

def factorial(n):
    return mp.gamma(n+1)

def binom(n,k):
    return factorial(n)/(factorial(k)*factorial(n-k))


def K ( i, j):
    polyL = L[i] #precomputed shiftedLegendre(i)
    return polyL.coeffs[-j-1]

def GetU (lambdas):
    n_max = len(lambdas)
    P = np.ones((n_max+1,U_points))
    x = np.linspace(0,1, U_points)
    P[1] = 2*x-1
    for n in range(1,n_max):
        P[n+1] = ((2*n+1)/(n+1)) * (2*x-1) * P[n]  - (n/(n+1))*P[n-1]
    U = np.zeros(U_points)
    for i in range (len(lambdas)):
        # if i%10 == 0: print(i)
        U += lambdas[i]*P[i]*np.sqrt(2*i+1)
        #polyL = L[i] #shiftedLegendre(i)
        #U += lambdas[i]*polyL(x)
    return U

def GetLambda(mag_rf):
    M_cutoff = len(mag_rf)
    all_lambda = []
    for i in range(M_cutoff):
        lambd = mpf(0.0)
        for j in range(i+1):
            lambd += K(i,j)*mpf(mag_rf[j])
#            print("K({:d},{:d}) =".format(i,j), float(K(i,j)))
        all_lambda.append(float(lambd))
    # tmp = [mpf(0.0) for n in range(M_cutoff)]
    # all_lambda = np.array(tmp)
    # all_lambda[10] = mpf(1.0)
    return all_lambda

def GetAbsLastLambda(mag_rf):
    M_cutoff = len(mag_rf)
    all_lambda = []
    i = M_cutoff - 1
    lambd = mpf(0.0)
    for j in range(i+1):
        lambd += K(i,j)*mpf(mag_rf[j])
    #    print("K({:d},{:d}) =".format(i,j), float(K(i,j)))
    all_lambda.append(float(lambd))
        # tmp = [mpf(0.0) for n in range(M_cutoff)]
    # all_lambda = np.array(tmp)
    # all_lambda[10] = mpf(1.0)
    return np.abs(all_lambda[0])


def shiftedLegendre(n):
    coeffs = []
    for k in range(n+1):
        val = mpf(-1)**n * binom(mpf(n),mpf(k)) * binom(n+k,k) * (-1)**k
        coeffs.insert(0,val*mp.sqrt(2*n+1))
    return np.poly1d(coeffs)

def GetGnk(c,n,k):
    Gnk = mpf(0.0)
    for m in range(k+1):
        Gnk+=(-1)**m * binom(mpf(k),mpf(m)) * c[n+m]
    return Gnk

def GetGnkSum(c):
    total = 0
    N = len(c)
    for n in range(2):
        for k in range (N-n-2,N-n):
            Gnk = float(GetGnk(c , n, k))
            if Gnk <0:
                total += Gnk
    return total



def FixGnk(mag_sin):
    total = GetGnkSum(mag_sin)
    print(total)

    
def mat1mp (mag_sin):
    N = len(mag_sin)
    NN = int((N-1)/2)
    tmp = np.array([mpf(0.0) for n in range((NN+1)**2)])
    tmp2 = np.reshape(tmp,( NN+1 , NN+1 ))
    mat = mp.matrix(tmp2)
    # mat = np.zeros(( NN+1 , NN+1 ))
    for i in range(NN+1):
        for j in range(NN+1):
            mat[i,j] = mag_sin[i+j+1]
    # print(mat)
    return mat
def mat1 (mag_sin):
    N = len(mag_sin)
    NN = int((N-1)/2)
    mat = np.zeros(( NN+1 , NN+1 ))
    for i in range(NN+1):
        for j in range(NN+1):
            mat[i,j] = mag_sin[i+j+1]
    # print(mat)
    return mat

def mat2 (mag_sin):
    N = len(mag_sin)
    NN = int((N-1)/2)
    mat = np.zeros(( NN+1 , NN+1 ))
    for i in range(NN+1):
        for j in range(NN+1):
            mat[i,j] = mag_sin[i+j] - mag_sin[i+j+1]
    return mat


def fitness(corrector):
    mag_sin = mag_sin_0*(1 + mpf(1.1)*(corrector-0.5)*noise_ratio)
    res = GetAbsLastLambda(mag_sin)
    global globi
    globi = globi + 1    
    if globi % 100 == 0: print("REsult", res)
    return res


mag_sin_0, mag_cos_0 = gen_rf_signal(time_steps)
mag_sin = np.copy(mag_sin_0)
mag_cos = np.copy(mag_cos_0)

L = [] # Shifted Legendre polinomials
for i in range(len(mag_sin)):
    polyL = shiftedLegendre(i)
    L += [polyL]

globi = 1

# import time
# time1 = time.time()
# corr = np.zeros(len(mag_sin))
# fitness(corr)
# from scipy.optimize import minimize
# # = minimize(fitness,corr,method='L-BFGS-B')
# v = minimize(fitness,corr,method='TNC',options={'disp': False, 'minfev': 0, 'scale': None, 'rescale': -1, 'offset': None, 'gtol': -1, 'eps': 1e-08, 'eta': -1, 'maxiter': None, 'maxCGit': -1, 'mesg_num': None, 'ftol': -1, 'xtol': -1, 'stepmx': 0, 'accuracy': 0})

# # from scipy.optimize import differential_evolution
# # bounds = []
# # for i in range(len(mag_sin)):
# #     bounds.append((0,1))
# # v = differential_evolution(fitness,bounds)

# fitness(v.x)
# print(v.fun, "at", v.x)
# time2 = time.time()
# print( 'it took %0.5f s' % ((time2-time1)))

print("N = ",len(mag_sin))
print("Noise = ", noise_ratio )
# print("c_n =",["%0.8f"%(x) for x in mag_sin])
# print("mat1 =",mat1(mag_sin))
# print("mat2 =",mat2(mag_sin))

#print( 'eigsh took %0.5f s' % ((time2-time1)))
#print("eigsh1 = ",v1,"eigsh2 = ",v2)
# v1f = eig(mat1(mag_sin))
# v2f = eig(mat2(mag_sin))
# print("eig1 = ",np.sort(v1f[0]))
# print("eig2 = ",np.sort(v2f[0]))

# v = eigsh(mat2(mag_sin), k=6, return_eigenvectors=False, which='LA')
# print(v)

# N = len(mag_sin)
# print("N = ", N)
# for n in range(2):
#     for k in range (N-n-2,N-n):
#         Gnk = float(GetGnk(mag_sin, n, k))
#         if Gnk <0:
#             print("G",n,k," = ", Gnk)
        
#FixGnk(mag_sin)

sign = ""
# for i in range(voxel_num):
#     if i%5 == 0:
#         sign+="\n"
#     sign = sign + '{:3.2g}'.format(float(a_i[i]))+"/"+'{:3.2g}'.format(float(d_i[i]))+", "
# 
# #    print(mp.exp(-1.0/voxel_T2s_decay[i]))
               
 
plt.plot(mag_sin, ls=' ', marker='o')
plt.title("Signal to restore amp/decay_T:"+sign)
plt.savefig("signal.pdf")
plt.clf()



recL = []

    
x = np.linspace(0,1, U_points)
# polyL_val = np.array([float(L[-1](x[n])) for n in range(U_points)])
# plt.plot(x,polyL_val)
# plt.title("Legendre polynom of order "+str(len(L)))
# plt.savefig("polyL.pdf")
# plt.clf()
#print("Before lambda.")

import time
time1 = time.time()
lambdas = GetLambda(mag_sin)
time2 = time.time()
print( 'GetLambda took %0.5f s' % ((time2-time1)))
print((lambdas))
U = GetU(lambdas)

sign ="" 
for i in range(voxel_num):
    sign+="{:4.3g}, ".format(float( mp.exp(-(period/rf_samples_per_period)/voxel_T2s_decay[i]) ))
#print(sign)
x = np.linspace(0,1, U_points)
mag_x = np.linspace(0,1, len(mag_sin))
import matplotlib as matplotlib
matplotlib.rcParams.update({'font.size': 28})
plt.figure(figsize=(30, 20))
plt.plot(x,U, lw=0.2)
plt.title(r"$\sum^M \lambda_i L_i(x)$", y=1.02)
plt.xlim(0,1)
plt.xlabel(sign, fontsize=28)
plt.savefig("plt.pdf")
plt.clf()

# rms = np.sqrt(np.mean(U**2))/np.max(U)
# #print(rms)
# status = "OK"
# if rms < 0.2: status = "BAD"
# print("rms =", rms, status)