#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
from scipy.optimize import differential_evolution
from scipy.optimize import minimize
from scipy.special import gamma, binom
import numpy as np
voxel_num = 2
phase_range = np.pi/2
phase_init = 0.0
U_points = voxel_num * 1000

noise_ratio = 0.0 #1e8

total_periods = 8
rf_samples_per_period = 16

# B0=1.5T freq=64Mhz, period = 15.6 ns
period = 15.6/1000/1000 #ms
omega = 2.0*np.pi/period
#T2s_scale = 0.01 #ms # need to be 10ms
T2s_scale = total_periods*period #ms # need to be 10ms
T2s_min = T2s_scale/1000.0
#print(period)
#ms
time_steps = np.linspace(0, period*total_periods, rf_samples_per_period*total_periods)
voxel_amplitudes = np.random.rand(voxel_num)
voxel_T2s_decay = np.random.rand(voxel_num)*(T2s_scale-T2s_min) + T2s_min
voxel_all = np.append(voxel_amplitudes,voxel_T2s_decay/T2s_scale)
if voxel_num == 5:
    voxel_all = np.array([0.822628,0.691376,0.282906,0.226013,0.90703,0.144985,0.328563,0.440353,0.662462,0.720518])
    #voxel_all = [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
#first estimate    0.551777 0.190833 0.271438 0.814036 0.347389 0.926153 0.908453 0.581414 0.666012 0.673226

#voxel_amplitudes = [0.4, 0.8, 0]
#voxel_amplitudes = [0.9, 0.092893218813452, 0.5]
#voxel_amplitudes = [0.6, 0.517157287525381, 0.4]

test_amplitudes = np.zeros(voxel_num)
test_amplitudes = voxel_amplitudes
voxel_phases = 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


def gen_rf_signal(time):
    mag_sin = 0.0
    mag_cos = 0.0
    for i in range(voxel_num):
        amp = voxel_amplitudes[i] * (
                np.exp(-time/voxel_T2s_decay[i])
                ) + ( 0.0 
                           # + np.random.rand()*noise_ratio
                    )
        mag_sin += amp * np.sin(
            voxel_phases[i] + phase_init
            )
        mag_cos += amp * np.cos(
            voxel_phases[i] + phase_init
            )
    return mag_sin, mag_cos

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

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

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

def GetLambda(mag_rf):
    M_cutoff = len(mag_rf)
    all_lambda = []
    for i in range(M_cutoff):
        lambd = 0
        for j in range(i+1):
            lambd += K(i,j)*mag_rf[i]
        all_lambda.append(lambd)
    all_lambda = np.zeros(M_cutoff)
    all_lambda[19] = 1
    return all_lambda

def GetU (lambdas):
    x = np.linspace(0,1, U_points)
    U = np.zeros(U_points)
    for i in range (len(lambdas)):
        polyL = L[i] #shiftedLegendre(i)        
        U += lambdas[i]*polyL(x)
    return U



mag_sin, mag_cos = gen_rf_signal(time_steps)
L = [] # Shifted Legendre polinomials
for i in range(len(mag_sin)):
    polyL = shiftedLegendre(i)        
    L += [polyL]
#print(len(L))
print(L[20])
lambdas = GetLambda(mag_sin)
U = GetU(lambdas)

import matplotlib.pyplot as plt
x = np.linspace(0,1, U_points)
mag_x = np.linspace(0,1, len(mag_sin))
# crop = 1
# plt.plot(x[:-crop],U[:-crop])
plt.plot(x,U)
# plt.plot(mag_x,mag_sin)
#plt.xlim(0.2,0.8)
#plt.ylim(0.0,2.8)
plt.savefig("plt.pdf")
#plt.show()

#print(voxel_phases)
#print (voxel_amplitudes)

# import matplotlib.pyplot as plt
# plt.plot(time_steps, mag_sin)
# plt.plot(time_steps, mag_cos)
# plt.show()

# #print(fitness(test_amplitudes))

amplitude_minmax = (0,1)
T2s_minmax = (T2s_min/T2s_scale,1)

x0 = np.full(2*voxel_num,0.5)
x0 = [0.551777,0.190833,0.271438,0.814036,0.347389,0.926153,0.908453,0.581414,0.666012,0.673226]


# #result.x[voxel_num:] = result.x[voxel_num:]/T2s_scale
# print("amp/decay", voxel_amplitudes,voxel_T2s_decay)
# print("all   ", voxel_all)
# print("eval  ",result.x, "\n=====> fun=",result.fun)



# # print("Diff")
# # print((voxel_amplitudes-result.x))
# # print("percent")
# print("percent",np.abs(voxel_all-result.x)*100)
# if np.max(np.abs(voxel_all[:voxel_num]-result.x[:voxel_num])*100)>0.5:
#     print ("==============  !!!LARGE!!! ===============")

# print("\n")