-- 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

function Drop(list,n)
    local result = {}
    for i = n +1, #list do
        result[i-n] = list[i]
    end
    return result
end

function Take(list,n1,n2)
  local result={}
  if n2== nil then
    if n1 > 0 then
      for i = 1, n1 do
        result[i]=list[i]
      end
    else
      for i=1,-n1 do
        result[i]=list[#list+n1+i]
      end
    end
   else
     for i = n1, n2 do
       result[i+1-n1] = list[i]
     end
   end
  return result
end

function Join(...)
  local result = {}
  local n=1
  for i = 1, select("#",...) do
    for j=1,#select(i,...) do
      result[n]=select(i,...)[j]
      n=n+1
    end
  end
  return result
end

function Flatten(list)
  local result = {}
  local k=1
  for i=1,#list do
    for j=1,#list[i] do
      result[k] = list[i][j]
      k=k+1
    end
  end
  return result
end


function TakeFlatten(list,n)
    local result = {}
    local k = 1
    if n > 0 then
        for i = 1, n do
            for j = 1, #list[i] do
                result[k] = list[i][j]
                k = k + 1
            end
        end
    else
        for i = -n, 1, -1 do
            for j = 1, #list[#list-i+1] do
                result[k] = list[#list-i+1][j]
                k = k + 1
            end
        end
    end
    return result
end

function RemoveFromClosedOrbitals(Indices,ClosedLabels,patterns)
    local result = {}
    local k = 1
    for i = 1, #ClosedLabels do
        local found = false
        for j = 1, #patterns do
            if string.find(ClosedLabels[i],patterns[j]) then
                found = true
            end
        end
        if not found then 
            result[k] = Indices.Closed[i]
            k = k+1
        end
    end
    return result
end

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

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 Sum(x)
    local sum = 0
    for i = 1, #x do
        sum = sum + x[i]
    end
    return sum
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},Hlocal,{{"#det",NumberOfDeterminants}})
  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

function IsElement(x, M)
    for i = 1, #M do
        if M[i] == x then return true end
    end
    return false
end

function PrintRelativeVariance(H,psi)
    local E = {}
    local EtotSqr = 0
    local var = {}
    for i = 1, #psi do
        local Hpsi = H*psi[i]
        E[i] = Braket(Hpsi,psi[i])
        var[i] = Braket(Hpsi,Hpsi) - E[i]*E[i]
        EtotSqr = EtotSqr + E[i]*E[i]
    end
    for i = 1, #E do
       print(var[i]/ E[i]/E[i],var[i]/EtotSqr*#E)
    end
end

function FindEnergyLowering(psi, H, Restrictions, NTri)
  local psiloc = Copy(psi)
  psiloc.NF = H.NF
  spec, spectri = CreateSpectra(H,NewOperator("Identity",H.NF),psiloc,{{"NTri",NTri},{"Restrictions",Restrictions},{"E0",0},{"DenseBorder",0}})
  print(spectri)
  local val0 = Eigensystem(spectri[1])
  print("The energy of the lowest state including upto order n is:")
  for n=NTri,2,-1 do
    val,fun = Eigensystem(spectri[1])
    print(n,val[1])
    spectri[1][2][n]=0
  end
  return val0[1]
end

function InterpolateED(ed, x)
    local gval = {}
    local fval = {}
    for i = 1, #x do
        gval[i] = ed[i]
        fval[i] = ed[i+#x]
    end
    local gint = InterpolatingFunction.Spline(x,gval)
    local fint = InterpolatingFunction.Spline(x,fval)
    return gint, fint
end

function InterpolateEDList(eds, x)
    local glist = {}
    local flist = {}
    for i = 1, #eds do
        glist[i], flist[i] = InterpolateED(eds[i],x)
    end
    return glist, flist
end

-- Convert decimal number into fraction
local function to_frac(num)
   local W = math.floor(num)
   local F = num - W
   local pn, n, N = 0, 1
   local pd, d, D = 1, 0
   local x, err, q, Q
   repeat
      x = x and 1 / (x - q) or F
      q, Q = math.floor(x), math.floor(x + 0.5)
      pn, n, N = n, q*n + pn, Q*n + pn
      pd, d, D = d, q*d + pd, Q*d + pd
      err = F - N/D
   until math.abs(err) < 1e-15
   out = {}
   out[1] = N + D*W
   out[2] = D
   out[3] = err
   return out
end


-- Convert angular momentum into the corresponding letter (s,p,d...)
local function LtoLetter(l)
	if l == 0 then
		letter = "s"
	elseif l == 1 then
		letter = "p"
	elseif l == 2 then
		letter = "d" 
	elseif l == 3 then
		letter = "f" 
	elseif l == 4 then
		letter = "g"
	elseif l == 5 then
		letter = "h"
	else
		print("For angular momenta > 5, the assignment is not implemented yet.")
	end
	return letter
end


function AssignIndices(nb,lb,jb,mb,na,la,ja,ma,basis)

	ind_loc, NF_loc = CreateAtomicIndicesList(basis,{{"Kappas",true}})

	-- tau_b
	fraction_b = to_frac(jb)
	letter_b   = LtoLetter(lb)
	
	b_string = tostring(nb)..letter_b..tostring(fraction_b[1]).."/"..tostring(fraction_b[2])
	
	for key, orbname in pairs(basis) do
		if (string.find(orbname,b_string)) then
			b_indexlist = ind_loc[key]
			b_index 	= b_indexlist[jb + mb + 1]
		end
	end
	
	-- tau_a	
	fraction_a = to_frac(ja)
	letter_a   = LtoLetter(la)
	
	a_string = tostring(na)..letter_a..tostring(fraction_a[1]).."/"..tostring(fraction_a[2])
	
	for key, orbname in pairs(basis) do
		if (string.find(orbname,a_string)) then
			a_indexlist = ind_loc[key]
			a_index 	= a_indexlist[ja + ma + 1]
		end
	end
	
	return b_index, a_index
	
end



