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

noise_ratio = 0.0 #1e8

total_periods = 50
rf_samples_per_period = 1

# 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 = [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]
#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)
#print(voxel_phases)
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 assumed_signal(time, values):
    amplitudes = values[:voxel_num]
    T2s_decay = values[voxel_num:]*T2s_scale
    mag_sin = 0.0
    mag_cos = 0.0
    for i in range(voxel_num):
        amp = amplitudes[i] * (
                np.exp(-time/T2s_decay[i])
                ) 
        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


mag_sin, mag_cos = gen_rf_signal(time_steps)

def fitness(amplitudes, *args):
    assumed_sin, assumed_cos = assumed_signal(time_steps, amplitudes)
    diff = 0
    #amp = np.sqrt(mag_sin**2 + mag_cos**2)
    # diff += np.sqrt(np.mean(np.square((mag_sin - assumed_sin)/amp)))
    # diff += np.sqrt(np.mean(np.square((mag_cos - assumed_cos)/amp)))
    diff += np.sqrt(np.mean(np.square((mag_sin - assumed_sin))))
    diff += np.sqrt(np.mean(np.square((mag_cos - assumed_cos))))
    diff += np.sqrt(np.mean(np.square(mag_sin/mag_cos - assumed_sin/assumed_cos)))
    return diff

def hyper_fit(amplitudes, *args):
    result = minimize(fitness, amplitudes, bounds = bounds, method='L-BFGS-B', options={'ftol': 1e-16,'gtol': 1e-16, 'disp': False})
    print (result.fun)
    return result.fun
#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))

bounds = []
amplitude_minmax = (0,1)
T2s_minmax = (T2s_min/T2s_scale,1)
for i in range(voxel_num):
    bounds.append(amplitude_minmax)
for i in range(voxel_num):
    bounds.append(T2s_minmax)

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 = minimize(hyper_fit, x0, bounds = bounds, method='L-BFGS-B', options={'ftol': 1e-16,'gtol': 1e-16, 'disp': True})
result = minimize(fitness, x0, bounds = bounds, method='L-BFGS-B', options={'ftol': 1e-16,'gtol': 1e-16, 'disp': True})

# result = differential_evolution(hyper_fit, bounds, polish=True
#                                     , maxiter = voxel_num*2*500
#                                 , disp = True
# )

# bestresult = minimize(fitness, x0, bounds = bounds, method='L-BFGS-B', options={'ftol': 1e-16,'gtol': 1e-16, 'disp': True})
# for x in range(50):
#     print("Try to solve: ", x)
#     x0 = np.random.rand(2*voxel_num)
#     result = minimize(fitness, x0, bounds = bounds, method='L-BFGS-B', options={'ftol': 1e-16,'gtol': 1e-16, 'disp': False})
#     if result.fun < bestresult.fun:
#         bestresult = result        
#         print("new best: ", bestresult.fun)
# result = bestresult

#result = minimize(fitness, x0, bounds = bounds, method='COBYLA', tol=1e-16, options={ 'maxiter':20000,'disp': True})
#result = minimize(fitness, result.x, bounds = bounds, method='L-BFGS-B', options={'gtol': 1e-16, 'disp': True})
#result = minimize(hyper_fit, x0, bounds = bounds, method='Nelder-Mead', tol=1e-16,  options={'maxiter':100,'disp': True})

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