#!/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 = (1.2e-16) #(1.2e-16) #1e8

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

# 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
n_mu = total_steps
#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),(0.2)])
d_i = np.array([(0.7),(0.9),(0.2),(0.67),(0.41)])
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)
    v_sum = 0
    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)**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))
            # if t == 0:
            #     v = voxel_amplitudes[i] * (
            #     mp.exp(-(time[1])/voxel_T2s_decay[i])
            #     )* np.sin((voxel_phases[i] + phase_init))
            #     v_sum += v
            #     print(float(v))
            #     #sign+="{:4.3g}, ".format(float( mp.exp(-(period/rf_samples_per_period)/voxel_T2s_decay[i]) ))

        mag_sin[t] *= (1 + (np.random.rand()-0.5)*noise_ratio)
        # if t == 0:
        #     print(v_sum, mag_sin[t])
        #mag_cos[t] += np.random.rand()*noise_ratio
    return mag_sin, mag_cos


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)

# mu = [np.random.rand() for i in range(n_mu)]
mu = mag_sin
#print(mu)

P = np.zeros((n_mu+1, n_mu+1),dtype=mpf)
P[0,0] = 1
for i in range(n_mu):
    P[i,1] = (-1)**i * mu[i]

for j in range(2,n_mu+1):
    for i in range(n_mu-1,-1,-1):
        if i+j > n_mu: continue
        # print("i =",i, "  j =",j, "   P = ",P[i,j])
        P[i,j] = P[0,j-1]*P[i+1,j-2] - P[0,j-2]*P[i+1,j-1]

# for i in range (n_mu+1):
#     line = [float("%5.3g"%(x)) for x in P[i,:]]
#     print(line)
#print(P)

alpha = np.zeros(n_mu,dtype=mpf)
for k in range(1,n_mu):
    alpha[k] = (P[0,k+1])/((P[0,k])*(P[0,k-1]))
# print("alpha:",len(alpha))
# line = [float("%5.5g"%(x)) for x in alpha]
# print(line)


M=np.zeros((voxel_num,voxel_num), dtype=mpf)
for i in range(voxel_num):
    for j in range(voxel_num):
        if i == j:
            M[i,j] = alpha[i*2]+alpha[i*2 +1]
        if abs(i-j) == 1:
            m = max(i,j)
            M[i,j] = -np.sqrt( alpha[2*m-1] * alpha[2*m] )
M1 = mp.matrix(M.tolist())
#print(M1)

spectra = mp.eig(M1)[0]
line = [float("%5.5g"%(x)) for x in spectra]
print("eig(M1) =",line)

# print("N = ",len(mag_sin))
# print("Noise = ", noise_ratio )
# print("c_n =",["%0.8f"%(x) for x in 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]))
           
    
x = np.linspace(0,1, U_points)

import time
time1 = time.time()

time2 = time.time()
#print( 'It took %0.5f s' % ((time2-time1)))

sign ="" 
for i in range(voxel_num):
    sign+="{:4.3g}, ".format(float( mp.exp(-(period/rf_samples_per_period)/voxel_T2s_decay[i]) ))
print(d_i)
# 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()

plt.plot(mag_sin, ls=' ', marker='o')
plt.title("Signal to restore amp/decay_T:"+sign)
plt.savefig("signal.pdf")
plt.clf()