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

noise_ratio = 1e8

total_periods = 1000
rf_samples_per_period = 10

# 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_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
print("amp/decay", voxel_amplitudes,voxel_T2s_decay)
voxel_all = np.append(voxel_amplitudes,voxel_T2s_decay/T2s_scale)
print("all   ", voxel_all)
#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):
    signal = 0.0
    for i in range(voxel_num):
        signal += voxel_amplitudes[i]*np.sin(
            omega*time + voxel_phases[i] + phase_init
            ) * (
                np.exp(-time/voxel_T2s_decay[i])
                ) + ( 0.0 
                           + np.random.rand()/noise_ratio
                    )
    return signal

def assumed_signal(time, values):
    amplitudes = values[:voxel_num]
    T2s_decay = values[voxel_num:]*T2s_scale
    signal = 0.0
    for i in range(voxel_num):
        signal += amplitudes[i]*np.sin(
            omega*time + voxel_phases[i] + phase_init
            ) * (
                np.exp(-time/T2s_decay[i])
                )
    return signal

rf_signal_measured = gen_rf_signal(time_steps)

def fitness(amplitudes):
    diff = rf_signal_measured - assumed_signal(time_steps, amplitudes)
    return np.sqrt(np.mean(np.square(diff)))

#print(voxel_phases)
#print (voxel_amplitudes)

import matplotlib.pyplot as plt
plt.plot(time_steps, rf_signal_measured)
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)
result = differential_evolution(fitness, bounds, polish=True
                                    #, maxiter = voxel_num*2*500
                                    )
#result.x[voxel_num:] = result.x[voxel_num:]/T2s_scale
print("eval  ",result.x, 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")