{-# LANGUAGE NoImplicitPrelude #-}
module TBit.Plots (bzPlot, paramPlot, bandPlot) where

import TBit.Types
import TBit.Framework
import TBit.Parameterization
import TBit.Hamiltonian.Eigenstates (eigenenergies)

import Prelude.Listless
import Data.List.Stream
import Data.Map (adjust, elems)
import Control.Monad.State
import Control.Parallel.Strategies
import Numeric.LinearAlgebra.HMatrix
import qualified Data.Traversable as T

{-|
    Given a hamiltonian and a set of parameters, plot a real-valued function
    over the Brillouin zone. For a Berry curvature plot, /e.g./,
    
    > bzPlot (kagomeAF, defaultParams) (bandCurvature 0)

    The output string is suitable for gnuplot, unless that string is
    reporting an error from the calculation.
-}
bzPlot :: (Parameterized Hamiltonian, Parameters) 
       -> (Hamiltonian -> Wavevector -> Parameterized Double) 
       -> String
bzPlot (h, ps) f = case crunch (h >>= computeBZPlot f) ps
                     of Left  err -> show err
                        Right dat -> unlines
                                   . map (\(k, d) -> (unwords . map show . toList $ k)
                                                      ++ " " ++ show d)
                                   $ dat

computeBZPlot :: (Hamiltonian -> Wavevector -> Parameterized Double)
              -> Hamiltonian 
              -> Parameterized [(Wavevector, Double)]
computeBZPlot f h = meshBZ 
                >>= T.mapM (\k -> T.sequence (k, f h k))
                >>= (return . elems)

{-|
    Given a hamiltonian and a set of parameters, plot a real-valued function
    over some range of a parameter. For a Chern number plot over the hopping
    parameter, /e.g./,
    
    > paramPlot (kagomeAF, defaultParams) (chern 1) ("t", [1.0, 1.1 .. 5.0])

    The output string is suitable for gnuplot, unless that string is
    reporting an error from the calculation.
-}
paramPlot :: (Parameterized Hamiltonian, Parameters)
          -> (Hamiltonian -> Parameterized Double)
          -> (String, [Double])
          -> String
paramPlot (h, ps) f (s, ss) = case crunch (computePMPlot f s ss h) ps
                                of Left  err -> show err
                                   Right dat -> unlines
                                              . map (\(p, d) -> show p 
                                                             ++ " " 
                                                             ++ show d)
                                              $ dat
                                        

computePMPlot :: (Hamiltonian -> Parameterized Double)
              -> String
              -> [Double]
              -> Parameterized Hamiltonian
              -> Parameterized [(Double, Double)]
computePMPlot f s ss h = liftM (zip ss) 
                       $ mapM (\v -> modify (setScalar s v) >> (h >>= f)) $ ss
    where setScalar str val ps = ps { scalarParams = adjust (\_ -> val :+ 0.0) str 
                                                   $ scalarParams ps }


multiPlot :: (Parameterized Hamiltonian, Parameters)
          -> [(Hamiltonian -> Parameterized Double)]
          -> (String, [Double])
          -> String
multiPlot sys fs params = unlines $ map (\f -> paramPlot sys f params) fs

{-|
    Given a hamiltonian and a set of parameters, plot the bands over
    a path anchored by the provided wavevectors. For instance,
    
    > gamma  = vector [0.0 ,  0.0]
    > point1 = vector [0.5 ,  1.0]
    > point2 = vector [0.5 , -1.0]
    > bandPlot (kagomeAF, defaultParams) [gamma, point1, point2, gamma]

    The output string is suitable for gnuplot, unless that string is
    reporting an error from the calculation.
-}
bandPlot :: (Parameterized Hamiltonian, Parameters)
         -> [Wavevector]
         -> String
bandPlot (h,ps) ks = either show id $ crunch (bandData h ks) ps

bandData :: Parameterized Hamiltonian
         -> [Wavevector]
         -> Parameterized String
bandData h ks = do mtx <- h
                   kPath ks >>= mapM (eigenenergies mtx)
                            >>= (return . unlines
                                        . map (\(x, e) -> show x ++ " " ++ show e)
                                        . concatMap T.sequence 
                                        . zip [0..])