{-# 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..])