import numpy
from matplotlib import pyplot
import os
from jqc import jqc_plot
from scipy import constants
from diatom import Hamiltonian
from matplotlib import gridspec

h = constants.h
muB = constants.physical_constants['Bohr magneton'][0]
muN = constants.physical_constants['nuclear magneton'][0]
cwd = os.path.dirname(os.path.abspath(__file__))

jqc_plot.plot_style('thesis')

colours = jqc_plot.colours

Hyperfine_States = numpy.load(cwd+"\\states.npy")

Hyperfine_energy = numpy.genfromtxt(cwd+"\\Energies.csv",
                                    delimiter=',')
RbCs = Hamiltonian.RbCs
I1 = RbCs['I1']
I2 = RbCs['I2']
Nmax = 1

indices = []
I_vector = []
Nvec,I1vec,I2vec = Hamiltonian.Generate_vecs(Nmax,I1,I2)
Nop = numpy.array([n for n in range(0,Nmax+1) for mi in range(0,32*(2*n+1))])
muZ = Hamiltonian.Zeeman_Ham(Nmax,I1,I2,RbCs)

F = (Nvec+I1vec+I2vec)
Fz= F[2]
order = numpy.argsort(Hyperfine_energy[1:,-1])

Hyperfine_energy[1:,:]=Hyperfine_energy[order+1,:]

for i in range(32):
    pyplot.plot(Hyperfine_energy[0,:],Hyperfine_energy[1+i,:]/h+500e3)
pyplot.show()

Hyperfine_States= Hyperfine_States[:,order,:]


print("Calculating mF")
mF = numpy.round(numpy.einsum('ik,ij,jk->k',
        numpy.conjugate(Hyperfine_States[:,:,-1]),Fz,
        Hyperfine_States[:,order,-1])).real


print("Calculating mu")
mu =numpy.einsum('ikl,ij,jkl->kl',
        numpy.conjugate(Hyperfine_States[:,:,:]),muZ,
        Hyperfine_States[:,order,:]).real
B = Hyperfine_energy[0,:]*1e4

mu = numpy.vstack((B,mu/muN))

numpy.savetxt("mu_output.csv",mu,delimiter=',')
print(mu)
mF = numpy.insert(mF,0,0.1)
Nop = numpy.insert(Nop,0,0.1)
print(mF,mF.shape)
mu= numpy.insert(mu,0,mF,axis=1)
mu= numpy.insert(mu,0,Nop,axis=1)
numpy.savetxt("mu_output2.csv",mu,delimiter=',',
header="0 N/A N/A, B(G) \n1:N, mF, mu(B) [units of MuN]...")
