module TBit.Hamiltonian.Eigenstates ( eigenstates
                                    , eigenbras
                                    , eigenkets
                                    , eigenenergies
                                    , eigensystem ) where

import TBit.Types
import Data.List (sortBy, zip, sort, map) -- .Stream
import Prelude hiding (zip, map)
import Numeric.LinearAlgebra.HMatrix

-- | Returns a list of eigenvectors sorted by eigenvalue. The lowest
--   energy state is the first element of the returned list.
eigenstates :: Hamiltonian -> Wavevector -> Parameterized [Eigenstate]
eigenstates h k = return . map snd $ eigensystem' (h k)

-- |As 'eigenstates', but converts each vector to a column matrix
--  for convenience in certain caclulations.
eigenkets :: Hamiltonian -> Wavevector -> Parameterized [Eigenket]
eigenkets h k = eigenstates h k >>= (return . map asColumn)

-- |As 'eigenkets', takes the conjugate transpose of each ket.
eigenbras :: Hamiltonian -> Wavevector -> Parameterized [Eigenbra]
eigenbras h k = eigenkets h k >>= (return . map tr)

-- |Returns a list of eigenvalues, sorted in ascending order.
eigenenergies :: Hamiltonian -> Wavevector -> Parameterized [Energy]
eigenenergies h = return . energies . h

energies :: Matrix ℂ -> [Energy]
energies = sort . toList . eigenvaluesSH'

eigensystem' :: Matrix ℂ -> [(Energy,Eigenstate)]
eigensystem' = sortBy (\t1 t2 -> compare (fst t1) (fst t2))
             . (\(vals,vecs) -> zip (toList vals) (toColumns vecs))
             . eigSH'

-- |Returns the full eigensystem, sorted by energy. Equivalent to zipping
--  the results of 'eigenenergies' and 'eigenstates'.
eigensystem :: Hamiltonian -> Wavevector -> Parameterized [(Energy,Eigenstate)]
eigensystem h k = return $ eigensystem' (h k)