{-# LANGUAGE FlexibleInstances #-}

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

import Control.Monad (void)
import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Backend.Cairo
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Grid
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