#!/usr/bin/python3

import spintrum
import math
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize
import multiprocessing
import mpmath as mp


def filter_spectrum(spec, freq_lim):
    x_axis = np.array([])
    y_axis = np.array([])
    for i in range(len(freq_lim)):
        for j in range(len(spec["x"])):
            if freq_lim[i][0] <= spec["x"][j] <= freq_lim[i][1]:
                x_axis = np.append(x_axis,spec["x"][j])
                y_axis = np.append(y_axis,spec["y"][j])
    return {"x": x_axis, "y": y_axis}


#defining initial parameters

with open("../data/methylformateSignal.txt") as f:
    data = f.readlines()
    data = np.array(list(map(np.double,data)))
    data = data - np.mean(data)

gammas = [4257.7e4,4257.7e4,4257.7e4,4257.7e4,1070.8e4]
multips = [2,2,2,2,2]

gammah = 2*math.pi*4257.7e4
BThermal = 1.8e0
T2 = 1
points = len(data)
sampleRate = 2000
spectrumRange = [[70,119],[121,179],[181,239],[241,299]]
jCouplings = \
[
    [0,0,0,0,0],
    [0,0,0,0,0],
    [0,0,0,0,0],
    [0,0,0,0,0],
    [0,0,0,0,0],
]

experimental_spectrum = filter_spectrum(spintrum.FFTSpectralDensity(data, 2000), spectrumRange)


gen = spintrum.SpinSimulator(gyromagneticRatios=gammas,
                             jCouplings=jCouplings,
                             spinMultiplicities=multips,
                             doPrint=False)

spinOp = spintrum.SpinOperations()
spinOp.add_operation(spintrum.SpinOperations.OPERATION__THERMAL_POPULATE,
                     {'Bx': 0, 'By': 0, 'Bz': BThermal, 'T': 293.778})
spinOp.add_operation(spintrum.SpinOperations.OPERATION__TIP_SPINS,
                     {'direction': 'y', 'BVsTArea': 4*math.pi/gammah})
spinOp.add_operation(spintrum.SpinOperations.OPERATION__SET_HAMILTONIAN,
                     {'Bx': 0, 'By': 0, 'Bz': 0})
spinOp.add_operation(spintrum.SpinOperations.OPERATION__INIT_TIME_INDEPENDENT_EVOLUTION,
                     {'samplingRate': sampleRate, 'measurementDirection': 'z'})
spinOp.add_operation(spintrum.SpinOperations.OPERATION__EVOLVE_TIME_INDEPENDENT,
                     {'points': points, 'threads': multiprocessing.cpu_count()})

gen.update_parameters(spinOperations=spinOp, jCouplings=jCouplings, gyromagneticRatios=gammas)


def generate_spectrum(params):

    amplitude = params[3]
    T2 = params[4]
    jCouplings = \
        [
            [0, params[0], params[0], params[0],params[1]],
            [0, 0, 0, 0, params[2]],
            [0, 0, 0, 0, params[2]],
            [0, 0, 0, 0, params[2]],
            [0, 0, 0, 0, 0],

        ]


    spinOp = spintrum.SpinOperations()
    spinOp.add_operation(spintrum.SpinOperations.OPERATION__THERMAL_POPULATE,
                         {'Bx': 0, 'By': 0, 'Bz': BThermal, 'T': 293.778})
    spinOp.add_operation(spintrum.SpinOperations.OPERATION__TIP_SPINS,
                         {'direction': 'y', 'BVsTArea': 4 * math.pi / gammah})
    spinOp.add_operation(spintrum.SpinOperations.OPERATION__SET_HAMILTONIAN,
                         {'Bx': np.float(params[5]), 'By': 0, 'Bz': np.float(params[6])})
    spinOp.add_operation(spintrum.SpinOperations.OPERATION__INIT_TIME_INDEPENDENT_EVOLUTION,
                         {'samplingRate': sampleRate, 'measurementDirection': 'z'})
    spinOp.add_operation(spintrum.SpinOperations.OPERATION__EVOLVE_TIME_INDEPENDENT,
                         {'points': points, 'threads': multiprocessing.cpu_count()})

    gen.update_parameters(jCouplings=jCouplings, spinOperations=spinOp)

    signal = amplitude*gen.simulate()

    signal = signal - np.mean(signal)
    signal = [signal[i]*math.exp(-i/sampleRate/T2) for i in range(len(signal))]

    fft = filter_spectrum(spintrum.FFTSpectralDensity(signal, 2000), spectrumRange)

#    plt.plot(fft['x'],fft['y'],realSpectrum['x'],realSpectrum['y'])
#    plt.show()

    return fft



def objective_func(params):
    print(params)
    spect = generate_spectrum(params)
    funcValue = np.sqrt(np.sum((spect['y'] - experimental_spectrum['y']) ** 2))
    print('Objective function value: ',funcValue)
    return funcValue



pars = np.array([-1.06785235e+00,2.26791767e+02,4.26684847e+00,4.53620118e-05,2.73284204e+00,3.11541220e-09,2.80718054e-09])
#pars = np.array([-1.06785235e+00,2.26791767e+02,4.26684847e+00,4.53620118e-05,2.73284204e+00])

# get_objective_func(pars)

optimum_params = scipy.optimize.fmin(objective_func, pars, ftol = 5.0e-13, maxfun = 1)
optjCouplings = mp.matrix([[0, optimum_params[0], optimum_params[0], optimum_params[0], optimum_params[1]], [0, 0, 0, 0, optimum_params[2]], [0, 0, 0, 0, optimum_params[2]], [0, 0, 0, 0, optimum_params[2]], [0, 0, 0, 0, 0]])

CorrelationMatrix = spintrum.get_correlation_matrix(optimum_params, generate_spectrum, objective_func)
print('Correlation matrix: ','\n',CorrelationMatrix)

print('Parameters at minimum: ','\n', optimum_params)
errorbars = spintrum.get_errorbars(CorrelationMatrix)
print('Standard deviations of parameters: ','\n',errorbars)

#writout to file
np.set_printoptions(suppress=True,precision=20)
fr = open("../data/methylformateFitResult.txt","w")
fr.write('Minimum at J-couplings:\n')
fr.write(str(optjCouplings)+'\n')
fr.write('Minimum at paramters:\n')
fr.write(str(optimum_params) + '\n')
fr.write('Parameters standard deviations:\n')
fr.write(str(errorbars)+'\n')
fr.write('Parameters correlation matrix:\n')
fr.write(str(CorrelationMatrix)+'\n')
fr.write('Minimum objective function value: ' + str(objective_func(optimum_params)))
fr.close()

fft_final = generate_spectrum(optimum_params)

plt.plot(fft_final['x'], fft_final['y'], experimental_spectrum['x'], experimental_spectrum['y'])

plt.show()