{-# LANGUAGE FlexibleInstances #-}

import System.Random
import Data.Random.Normal
import Control.Monad.State
import Numeric.LinearAlgebra as Lin

import Control.Monad(void)
import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Grid
import Graphics.Rendering.Chart.Backend.Cairo
import Graphics.Rendering.Chart.Plot.Histogram(defaultNormedPlotHist)
instance Default (PlotHist x Double) where
    def = defaultNormedPlotHist

-- | Evaluates the Winger surmise.
ρWinger :: Floating t => t -> t
ρWinger s = pi * s / 2 * exp(- pi * s^2 / 4)

-- | Computes the splitting of the center eigenvalues of a symmetric matrix.
eigenSplitting :: Herm Double -> Double
eigenSplitting m = e2 - e1
                   where
                     e = sortVector $ eigenvaluesSH m
                     i = (Lin.size e) `quot` 2
                     e1 = e Lin.! (i - 1)
                     e2 = e Lin.! i

-- | Produces random GOE matrices.
goeGen :: RandomGen g => Int -> State g (Herm Double)
goeGen n = get >>= (\g -> let (g1, g2) = split g in 
  (put g1) >> (return $ sym $ matrix n $ take (n^2) $ normals' (0.0, 2.0) g2))

-- | Produces random ∓1 matrices.
pm1Gen :: RandomGen g => Int -> State g (Herm Double)
pm1Gen n = state (\g -> let (m, g2) = runState (goeGen n) g in (trustSym $ cmap signum (unSym m), g2))

-- | Produces n eigenvalue splittings from the matrix ensemble m.
eigenSplittings :: RandomGen g => Int -> State g (Herm Double) -> State g [Double]
eigenSplittings n m = state (\g -> let (mats, g2) = runState (sequence $ replicate n m) g in (eigenSplitting <$> mats, g2))

hsChart :: RandomGen g => g -> State g (Herm Double) -> Int -> Int -> String -> Layout Double Double
hsChart g gen num bins tit = execEC $ do
  plot $ fmap histToPlot $ liftEC $ do
    plot_hist_title .= tit
    plot_hist_bins .= bins
    plot_hist_values .= scaledData
    plot_hist_range .= Just (0.0, 5.0)
    plot_hist_norm_func .= const (\x -> fromIntegral bins * fromIntegral x / fromIntegral num / 5.0)
  plot (line "ρ_Winger" [ [(s,ρWinger s) | s <- [0,0.005..5]] ])
  where
    rawData = evalState (eigenSplittings num gen) g
    mean = sum rawData / fromIntegral (length rawData)
    scaledData = (\x -> x / mean) <$> rawData

mkgrid g num bins = title `wideAbove` aboveN [ besideN [ layoutToGrid (hsChart g ((fst gen) n) num bins (snd gen Prelude.<> show n)) | gen <- gens] | n <- ns]
  where
    gens = [(goeGen, "GOE, N = "), (pm1Gen, "±1, N = ")]
    ns = [2, 4, 10]
    title = setPickFn nullPickFn $ label ls HTA_Centre VTA_Centre "Winger surmise compared with eigenvalue splitting of random matrices"
    ls = def { _font_size = 15, _font_weight = FontWeightBold }
    

main = do
  g <- newStdGen
  void $ renderableToFile (def { _fo_format = PDF}) "01.06.pdf" $ fillBackground def $ gridToRenderable $ mkgrid g 1000 50