-- A function that given:
-- 1) NF, an integer specifying the number of fermions in the system
-- 2) Opp, a list of Operators (Opp.name) whereby name needs to be one of the 
-- standard operators defined in Quanty
-- 3,4) The indices of where to find to shell to act upon
-- returns:
-- nothing
-- modifies:
-- The list of operators is modified such that the operator acting on the given indices is added
function AddOperator(NF,Opp,jindexmin,jindexplus)
  for key, OppName in pairs(Opp) do
    Opp[key] = Opp[key] + NewOperator(key,NF,jindexmin,jindexplus)
  end
end

-- Given the orbitalLabels, orbitalIndices (same order as labels) and the number of Fermions
-- returns a list of operators 
function CreateAngularMomentumOperators(orbitalLabels, orbitalIndices, NF)
  local Opp={}
  Opp.Sx=0
  Opp.Sy=0
  Opp.Sz=0
  Opp.Lx=0
  Opp.Ly=0
  Opp.Lz=0
  Opp.Jx=0
  Opp.Jy=0
  Opp.Jz=0
  
  for key, orbname in pairs(orbitalLabels) do
    if(string.find(orbname,"s1/2")) then
      AddOperator(NF,Opp,{},orbitalIndices[key])
    end
    if(string.find(orbname,"p1/2")) then
      otherjstr = string.gsub(orbname,"p1/2","p3/2")
      for key2, orbname2 in pairs(orbitalLabels) do
        if(string.find(orbname2,otherjstr)) then
          AddOperator(NF,Opp,orbitalIndices[key],orbitalIndices[key2])
          break
        end
      end
    end
    if(string.find(orbname,"d3/2")) then
      otherjstr = string.gsub(orbname,"d3/2","d5/2")
      for key2, orbname2 in pairs(orbitalLabels) do
        if(string.find(orbname2,otherjstr)) then
          AddOperator(NF,Opp,orbitalIndices[key],orbitalIndices[key2])
          break
        end
      end
    end
    if(string.find(orbname,"f5/2")) then
      otherjstr = string.gsub(orbname,"f5/2","f7/2")
      for key2, orbname2 in pairs(orbitalLabels) do
        if(string.find(orbname2,otherjstr)) then
          AddOperator(NF,Opp,orbitalIndices[key],orbitalIndices[key2])
          break
        end
      end
    end
    if(string.find(orbname,"g7/2")) then
      otherjstr = string.gsub(orbname,"g7/2","g9/2")
      for key2, orbname2 in pairs(orbitalLabels) do
        if(string.find(orbname2,otherjstr)) then
          AddOperator(NF,Opp,orbitalIndices[key],orbitalIndices[key2])
          break
        end
      end
    end
    if(string.find(orbname,"h9/2")) then
      otherjstr = string.gsub(orbname,"h9/2","h11/2")
      for key2, orbname2 in pairs(orbitalLabels) do
        if(string.find(orbname2,otherjstr)) then
          AddOperator(NF,Opp,orbitalIndices[key],orbitalIndices[key2])
          break
        end
      end
    end
  end
  Opp.Ssqr=Chop(Opp.Sx*Opp.Sx+Opp.Sy*Opp.Sy+Opp.Sz*Opp.Sz)
  Opp.Ssqr.Name="S^2"
  Opp.Lsqr=Chop(Opp.Lx*Opp.Lx+Opp.Ly*Opp.Ly+Opp.Lz*Opp.Lz)
  Opp.Lsqr.Name="L^2"
  Opp.Jsqr=Chop(Opp.Jx*Opp.Jx+Opp.Jy*Opp.Jy+Opp.Jz*Opp.Jz)
  Opp.Jsqr.Name="J^2"
  return Opp
end

-- A function that returns a list of numberoperators counting the occupation of each shell
-- The function density matrix actually might be faster, but returns more information 
function CreateNumberOperators(orbitalLabels, orbitalIndices, NF)
  local Opp={}
  Opp["Ntot"]=0
  for key, orbname in pairs(orbitalLabels) do
    local ones={}
    for i = 1,#(orbitalIndices[key]) do
      ones[i]=1
    end
    Opp["N"..orbname] = NewOperator("Number",NF,orbitalIndices[key],orbitalIndices[key],ones)
    Opp["N"..orbname].Name="N"..orbname
    Opp["Ntot"]=Opp["Ntot"]+Opp["N"..orbname]
  end
  Opp["Ntot"].Name="Ntot"
  return Opp  
end

-- A function that generates the indices of many single particle orbitals
function GenerateIndices(LabelsClosedOrbitals, LabelsOpenOrbitals,LabelsExcitedOrbitals,LabelsEmptyOrbitals, Z, Valence)
    local Indices = {}
    Indices.Labels = Simon.Flatten({LabelsClosedOrbitals, LabelsOpenOrbitals,LabelsExcitedOrbitals,LabelsEmptyOrbitals})
    Indices.All, Indices.NF = CreateAtomicIndicesList(Indices.Labels,{{"Kappas",true}})
    Indices.Dict = CreateAtomicIndicesDict(Indices.Labels)
    Indices.Closed = Take(Indices.All,#LabelsClosedOrbitals)
    Indices.Open   = Take(Indices.All,#LabelsClosedOrbitals+1,#LabelsClosedOrbitals+#LabelsOpenOrbitals)
    Indices.Excited= Take(Indices.All,#LabelsClosedOrbitals+#LabelsOpenOrbitals+1,#LabelsClosedOrbitals+#LabelsOpenOrbitals+#LabelsExcitedOrbitals)
    Indices.Empty   = Take(Indices.All,#LabelsClosedOrbitals+#LabelsOpenOrbitals+#LabelsExcitedOrbitals+1,#LabelsClosedOrbitals+#LabelsOpenOrbitals+#LabelsExcitedOrbitals+#LabelsEmptyOrbitals)
    Indices.Z = Z
    Indices.Valence = Valence
    Indices.NElectrons = Z - Valence
    Indices.NClosedElectrons = #Simon.Flatten(Indices.Closed)
    Indices.NOpenElectrons   = Indices.NElectrons - Indices.NClosedElectrons
    return Indices
end

function CreateHamiltonian(Indices, dir)
    G, F = ReadFPLOBasisFunctions(Indices.Labels,dir.."/+fval.001.1", dir.."/+fval.001.2")--,{{"Method","BSpline"}})                                      
    local H = NewOperator("Dirac",Indices.NF,Indices.Z, Indices.All,G,F) + NewOperator("AtomicU",Indices.NF, Indices.All,G,F)
    H = (H + ConjugateTranspose(H))/2 * EnergyUnits.Ha.value
    return H
end

function StartDeterminantString(NF,NE)
  local str=""
  for i=1,NE do
    str=str.."1"
  end
  for i=NE+1,NF do
    str=str.."0"
  end
  return str
end

-- A function that does precisely what you thunk it does
function Sum(x)
    local sum = 0
    for i = 1, #x do
        sum = sum + x[i]
    end
    return sum
end

function Poppulation(M1x, M1y, M1z)
    local n = #M1x
    local pop = {}
    for i = 1, n do
        pop[i] = 1/n
    end
    for i = n, 2, -1 do
        local Rates = CalcRates(M1x,M1y,M1z,i)
        local totR  = Sum(Rates)
        if totR > 1E-16 then 
           for j = 1, i -1 do
            pop[j] = pop[j] + pop[i] * Rates[j] / totR
           end
           pop[i] = 0
        end
    end
    return pop
end

function CalcRates(M1x,M1y,M1z,j)
    local rates = {}
    for i = 1, j-1 do
        rates[i] = Complex.Re(M1x[j][i]*Conjugate(M1x[j][i]))
        rates[i] = Complex.Re(rates[i] + M1y[j][i]*Conjugate(M1y[j][i]))
        rates[i] = Complex.Re(rates[i] + M1z[j][i]*Conjugate(M1z[j][i]))
    end
    return rates
end

function ExpectationValues(psi,Hglobal,MagicNumber,Opp,OppN,orbitalLabels)
  print("\n")
  local E0 = Braket(psi[1],Hglobal,psi[1])
  local Hlocal = Hglobal-E0
  print("Ground-state energy = ", E0+MagicNumber .."\n")
  local OppNlist={}
  for i=1,#orbitalLabels do
    OppNlist[i]=OppN["N"..orbitalLabels[i]]
  end
  OppNlist[#orbitalLabels+1]=OppN["Ntot"]
  -- We assume that the states decay due to M1 transitions
  -- M1 transitions scale with the magnetic dipole moment between two states
  -- We thus need to calculate the moment between all states
  local M1x = Braket(psi,2*Opp.Sx+Opp.Lx,psi)
  local M1y = Braket(psi,2*Opp.Sy+Opp.Ly,psi)
  local M1z = Braket(psi,2*Opp.Sz+Opp.Lz,psi)
  -- The lifetime of a state is the norm square of the magnetic dipole strength to all states
  -- at lower energy
  local EinsteinLifetime={}
  for i=1,#psi do
    EinsteinLifetime[i]=0
    for j=1,i-1 do
      EinsteinLifetime[i]=EinsteinLifetime[i]+Conjugate(M1x[i][j])*M1x[i][j]
      EinsteinLifetime[i]=EinsteinLifetime[i]+Conjugate(M1y[i][j])*M1y[i][j]
      EinsteinLifetime[i]=EinsteinLifetime[i]+Conjugate(M1z[i][j])*M1z[i][j]
    end
    EinsteinLifetime[i]=Chop(EinsteinLifetime[i])
  end
  local NumberOfDeterminants={}
  for i=1,#psi do
    NumberOfDeterminants[i]=psi[i].N
  end
  pop = Poppulation(M1x,M1y,M1z)
  print("sum pop:",Sum(pop))
  PrintExpectationValues(psi,{Opp.Jsqr,Opp.Lsqr,Opp.Ssqr,Opp.Jz},Hlocal,{{"Lifetime",EinsteinLifetime},{"#det",NumberOfDeterminants},{"pop",pop}})
  PrintExpectationValues(psi,OppNlist,Hlocal)
  
  --print("============ M1 x")
  --print(M1x)
  --print("============ M1 y")
  --print(M1y)
  --print("============ M1 z")
  --print(M1z)
  
  --print(psi)
  
  return E0+MagicNumber

end
