#!/usr/bin/python3

import spintrum
import math
import matplotlib.pyplot as plt
import numpy as np
import time

gammas = [4257.7e4,4257.7e4,4257.7e4,4257.7e4,1070.8e4]
multips = [2,2,2,2,2]
jCouplings = \
[
    [0,-0.8,-0.8,-0.8,226.8],
    [0,0,0,0,4.0],
    [0,0,0,0,4.0],
    [0,0,0,0,4.0],
    [0,0,0,0,0],
]

BThermal = 1.8e0
T2 = 10
sampleRate = 2e3
T= 1/sampleRate
points = 80000

gammah = 2*math.pi*4257.7


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': 4})


start_time = time.time()

signal = spintrum.simulate(gyromagneticRatios=gammas,
                           jCouplings=jCouplings,
                           spinMultiplicities=multips,
                           spinOperations=spinOp)

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

print("Simulation lastet: " + repr(time.time()-start_time) + "s")

plt.plot(signal)
plt.show()

fft = spintrum.FFTSpectralDensity(signal, sampleRate)

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