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

from mpmath import mp, mpf
mp.dps = 2000

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

noise_ratio = (0.0*1e-38) #1e8

total_periods=10#DO NOT CHANGE to fit T2s_scale autoscale
rf_samples_per_period = 10
# 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
time_steps = np.array(mp.linspace((0), (period*total_periods), rf_samples_per_period*total_periods))
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))
    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 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 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


mag_sin, mag_cos = gen_rf_signal(time_steps)

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()


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

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.")

lambdas = GetLambda(mag_sin)
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()