module TBit.Magnetic.OrbitalMagnetization ( orbMag
                                          , dichroism
                                          , dichroicIntegrand
                                          , intrinsicOM
                                          , integrandMk
                                          , bandIntrinsicOM
                                          ) where

import TBit.Types
import TBit.Framework
import TBit.Hamiltonian.Eigenstates
import TBit.Numerical.Derivative
import TBit.Parameterization

import Data.List.Stream (delete)
import Control.Monad.Except
import Prelude hiding (mapM, sequence)
import Numeric.LinearAlgebra.HMatrix

-- |Returns the total orbital magnetization due to filling the first
--  n bands.
orbMag :: Filling -> Hamiltonian -> Parameterized Magnetization
orbMag n h = bzIntegral (integrandOM n h)

-- |Returns the gauge-invariant self-rotational orbital magnetization, i.e.
--  the circular dichroism, of the occupied bands, accomplished by integrating
--  'dichroicIntegrand'.
dichroism :: Filling -> Hamiltonian -> Parameterized Magnetization
dichroism n h = bzIntegral (dichroicIntegrand n h)

-- |Returns the intrinsic orbital magnetization of the n/th/ band, namely
--  the integral of _m_(k) from (Xiao et al., 2005).
bandIntrinsicOM :: BandIndex -> Hamiltonian -> Parameterized Magnetization
bandIntrinsicOM n h = bzIntegral (integrandMk n h)

-- |Sums 'bandIntrinsicOM' over the first n bands.
intrinsicOM :: Filling -> Hamiltonian -> Parameterized Magnetization
intrinsicOM n h = liftM sum $ sequence [ bandIntrinsicOM j h | j <- [0 .. pred n] ]

-- |Essentially equation (12) from PRB _77_, 054438 (2008) with alpha, beta set
--  to x, y respectively so as to give the z-component of the (expectation
--  value of) the circular dichroism pseudovector. This form takes a double sum over
--  occupied and then unoccupied states, which improves over the implementation of
--  'integrandMk' by allowing for intra-(un)occupied band degeneracies as long as 
--  there is still an electronic band gap.
-- 
--  As is consistent with our API for these types of functions, the 
--  'TBit.Types.Filling' argument should be a positive integer counting the
--  number of filled bands.
dichroicIntegrand :: Filling -> Hamiltonian -> Wavevector -> Parameterized Magnetization
dichroicIntegrand b h k = do ket <- eigenkets h k
                             bra <- eigenbras h k
                             eng <- eigenenergies h k
                             (Spacing s) <- getMesh
                             return . negate . imagPart . sum
                                    $ zipWith (/)
                                      [ num $ (bra!!n) <> hx s <> (ket!!m)
                                           <> (bra!!m) <> hy s <> (ket!!n)
                                      | m <- occ , n <- unocc ]
                                      [ (eng!!m - eng!!n) :+ 0.0
                                      | m <- occ , n <- unocc ]
    where num m = m ! 0 ! 0
          occ = [0 .. pred b]
          unocc = [b .. pred dim]
          dim = fst $ size $ h k
          twice = (*) 2.0
          kx = k ! 0
          ky = k ! 1
          hx s = diff (s :+ 0.0) (\x -> h $ vector [realPart x, ky]) (kx :+ 0.0)
          hy s = diff (s :+ 0.0) (\y -> h $ vector [kx, realPart y]) (ky :+ 0.0)

                        

-- |Gives the gauge-invariant self-rotational orbital magnetism, which is
--  proportional to the circular dichroism for a particular band index.
--  This is m(k) in Xiao et al's semiclassical approach (hence the name).
integrandMk :: BandIndex -> Hamiltonian -> Wavevector -> Parameterized Magnetization
integrandMk n h k = do ket <- eigenkets h k
                       bra <- eigenbras h k
                       eng <- eigenenergies h k
                       (Spacing s) <- getMesh
                       if   (close eng n (pred n)) || (close eng n (succ n))
                       then throwError $ UndefinedError (degenErr eng)
                       else return 
                          $ imagPart . sum
                          $ zipWith (/)
                                    [ num $! (bra!!m) <> hx s <> (ket!!n)
                                          <> (bra!!n) <> hy s <> (ket!!m)
                                    | m <- delete n [0 .. pred dim]]
                                    [ (eng!!n - eng!!m) :+ 0.0
                                    | m <- delete n [0 .. pred dim]]
    where num m = m ! 0 ! 0
          dim = fst $ size $ h k
          kx = k ! 0
          ky = k ! 1
          eps = 0.000000001
          hx s = diff (s :+ 0.0) (\x -> h $ vector [realPart x, ky]) (kx :+ 0.0)
          hy s = diff (s :+ 0.0) (\y -> h $ vector [kx, realPart y]) (ky :+ 0.0)
          close _ (-1) _ = False
          close _ _ (-1) = False
          close eng i j = (abs $ eng!!i - eng!!j) < eps
          degenErr eng = case (close eng n (pred n), close eng n (succ n))
                           of (True, _) -> errMsg $ pred n
                              (_, True) -> errMsg $ succ n
                              otherwise -> "This error was in error."
          errMsg j     =  "Cannot integrate Berry curvature due to a gap "
                           ++"closing between bands "++(show $ succ j)++" and "
                           ++(show $ succ n)++"."

-- |The full orbital magnetization for a particular wavevector; integrating over
--  the Brillouin zone would give the actual OM due to a particular band filling.
integrandOM :: Filling -> Hamiltonian -> Wavevector -> Parameterized Magnetization
integrandOM n h k = do ket <- eigenkets h k
                       bra <- eigenbras h k
                       eng <- eigenenergies h k
                       (Spacing s) <- getMesh
                       if   any (< eps) (diffs eng)
                       then throwError $ UndefinedError (degenErr eng)
                       else return 
                          $ imagPart . sum
                          $ zipWith (/)
                                    [ (*) ((eng!!m + eng!!l) :+ 0.0) .
                                      num $! (bra!!m) <> hx s <> (ket!!l)
                                          <> (bra!!l) <> hy s <> (ket!!m)
                                    | m <- occ , l <- unocc ]
                                    [ (eng!!m - eng!!l)^2 :+ 0.0
                                    | m <- occ , l <- unocc ]
    where num m = m ! 0 ! 0
          occ = [0 .. pred n]
          unocc = [n .. pred dim]
          dim = fst $ size $ h k
          diffs eng = [ abs $ (eng!!m) - (eng!!l)
                      | m <- occ, l <- unocc ]
          kx = k ! 0
          ky = k ! 1
          eps = 0.000000001
          hx s = diff (s :+ 0.0) (\x -> h $ vector [realPart x, ky]) (kx :+ 0.0)
          hy s = diff (s :+ 0.0) (\y -> h $ vector [kx, realPart y]) (ky :+ 0.0)
          degenErr eng = let (j,l,_) = head $ filter (\(_,_,z) -> (z < eps))
                                            $ [ (m,p,abs $ eng!!m - eng!!p)
                                              | m <- occ, p <- unocc ]
                          in "Cannot integrate Berry curvature due to a gap "
                           ++"closing between bands "++(show $ succ j)++" and "
                           ++(show $ succ l)++"."