#!/usr/bin/python3

import spintrum
import math
import matplotlib.pyplot as plt
import numpy as np
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D


def run(dirs):
    gammas = [10] #Hz/T
    jCouplings = [[0]]

    BThermal = 10e0
    T2 = 10
    sampleRate = 1e4
    T= 1/sampleRate
    B0=10.
    B1=0.3e0
    points = 1650
    f = B0*gammas[0]
    bx = np.array([np.sin(2*math.pi*f*t/sampleRate)*B1 for t in range(points)])
    by = np.array([np.cos(2*math.pi*f*t/sampleRate)*B1 for t in range(points)])

    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__SET_HAMILTONIAN,
                         {'Bx': 0, 'By': 0, 'Bz': B0})
    spinOp.add_operation(spintrum.SpinOperations.OPERATION__EVOLVE_TIME_DEPENDENT,
                         {'threads': 4, 'samplingRate': sampleRate, 'store': dirs,
                          'input': {'Bx': bx, 'By': by}})

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

#this will result in points in the form xyzxyzxyz (order doesn't make a difference)
data = run('xyz')

mpl.rcParams['legend.fontsize'] = 10

fig = plt.figure()
ax = fig.gca(projection='3d')
datanp = np.transpose(np.split(np.array(data),len(data)/3))

ax.plot(datanp[0], datanp[1], datanp[2], label='parametric curve')
ax.legend()

plt.show()