import numpy,os
from diatom import Hamiltonian

fpath = os.path.dirname(os.path.abspath(__file__))
I1 = Hamiltonian.RbCs['I1']
I2 = Hamiltonian.RbCs['I2']
Nmax = 6

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

print("Calculating mF")
Nvec,I1vec,I2vec = Hamiltonian.Generate_vecs(Nmax,I1,I2)
N2 = Hamiltonian.vector_dot(Nvec,Nvec)
F = Nvec+I1vec+I2vec
Fz=F[2]
mF = numpy.round(numpy.einsum('ik,ij,jk->k',
            Hyperfine_States[:,:,-1],Fz,Hyperfine_States[:,:,-1]).real)

Nsq = numpy.round(numpy.einsum('ik,ij,jk->k',
            Hyperfine_States[:,:,-1],N2,Hyperfine_States[:,:,-1]).real)

numpy.savetxt(fpath+"\\N6 MF.csv",mF,delimiter=',')
print("Saved mF to:"+fpath+"\\N6 MF.csv")

numpy.savetxt(fpath+"\\N6 N2.csv",Nsq,delimiter=',')
