-- | Reanimate Fuckery
module CalamityBot.Commands.Reanimate
    ( reanimateGroup,
    ) where

import Calamity.Commands as C
import Calamity.Commands.Context (FullContext)
import Calamity
import CalamityBot.Utils.Reanimate
import qualified Data.Text as T
import qualified Polysemy as P
import Reanimate
import Reanimate.Render (Format(RenderWebm, RenderGif), Raster(RasterAuto))

import           Data.Complex
import           Graphics.SvgTree
import qualified Graphics.SvgTree.Types
import           Linear.V2
import           Codec.Picture
import qualified Data.ByteString as B
import Network.Mime (defaultMimeLookup)
import Control.Lens
import Network.Wreq

protectText :: T.Text -> T.Text
protectText = T.concatMap (fromString . protectChar)

protectChar :: Char -> String
protectChar '#'  = "\\#"
protectChar '$'  = "\\$"
protectChar '%'  = "\\%"
protectChar '^'  = "\\^{}"
protectChar '&'  = "\\&"
protectChar '{'  = "\\{"
protectChar '}'  = "\\}"
protectChar '~'  = "\\~{}"
protectChar '\\' = "\\textbackslash{}"
protectChar '_'  = "\\_{}"
protectChar x = [x]

renderText :: T.Text -> Animation
renderText text =
  let rt = center . latex $ text
      rtRotated = mkAnimation 5 (\t -> rotateAroundCenter (t * 360) rt)
  in addStatic (mkBackground "white") rtRotated

findSVG :: [Attachment] -> Maybe Attachment
findSVG = listToMaybe . filter (\a -> "image/svg" `B.isPrefixOf` defaultMimeLookup (a ^. #filename))

usableTree :: Tree -> Bool
usableTree None = False
usableTree (UseTree _ _) = False
usableTree (DefinitionTree _) = False
usableTree _ = True

reanimateGroup :: BotC r => P.Sem (DSLState FullContext r) ()
reanimateGroup = void
  . help (const "Commands related to reanimate fuckery")
  . C.group "reanimate"
  $ do
    help (const "Render a message") $
      command @'[KleenePlusConcat T.Text] "render" \ctx msg -> do
        let anim = renderText msg
        r <- P.embed $ renderToMemory anim RasterAuto RenderWebm 400 400 30
        case r of
          Right s -> do
            void $ tell ctx (TFile "lol.webm" s)
          Left e ->
            putLBSLn $ "Failed with reason: " <> e

    help (const "Render a fourier thing of an svg") $
      command @'[] "renders" \ctx -> do
        case findSVG (ctx ^. #message . #attachments) of
          Just svg -> do
            r <- P.embed $ Network.Wreq.get (T.unpack $ svg ^. #url)
            let file = r ^. responseBody
            let Just doc = parseSvgFile "a.svg" (decodeUtf8 file)
            let tree = head . fromList . filter usableTree $ doc ^. Graphics.SvgTree.Types.documentElements
            let f = mkSVGLatex $ flipYAxis tree
            let anim = setDuration 20 $ scene $ do
                  _ <- newSpriteSVG $ mkBackgroundPixel (PixelRGBA8 252 252 252 0xFF)
                  play $ fourierA f (fromToS 0 5)      -- Rotate 15 times
                    & setDuration 30
                    & signalA (reverseS . powerS 2 . reverseS) -- Start fast, end slow
                    & pauseAtEnd 2
                  play $ fourierA f (constantS 0)       -- Don't rotate at all
                    & setDuration 10
                    & reverseA
                    & signalA (powerS 2)                       -- Start slow, end fast
                    & pauseAtEnd 2
            r <- P.embed $ renderToMemory anim RasterAuto RenderGif 480 360 15
            case r of
              Right s -> do
                void $ tell ctx (TFile "lol.gif" s)
              Left e ->
                putLBSLn $ "Failed with reason: " <> e
          Nothing ->
            void $ tell @T.Text ctx "Couldn't find an svg"

    help (const "Render a fourier thing") $
      command @'[KleenePlusConcat T.Text] "renderf" \ctx msg -> do
        let f = mkFourierLatex msg
        let anim = setDuration 20 $ scene $ do
              _ <- newSpriteSVG $ mkBackgroundPixel (PixelRGBA8 252 252 252 0xFF)
              play $ fourierA f (fromToS 0 5)      -- Rotate 15 times
                & setDuration 30
                & signalA (reverseS . powerS 2 . reverseS) -- Start fast, end slow
                & pauseAtEnd 2
              play $ fourierA f (constantS 0)       -- Don't rotate at all
                & setDuration 10
                & reverseA
                & signalA (powerS 2)                       -- Start slow, end fast
                & pauseAtEnd 2
        r <- P.embed $ renderToMemory anim RasterAuto RenderGif 480 360 15
        case r of
          Right s -> do
            void $ tell ctx (TFile "lol.gif" s)
          Left e ->
            putLBSLn $ "Failed with reason: " <> e

scaleToFit :: Double -> Double -> Tree -> Tree
scaleToFit w h t =
  let oh = svgHeight t
      ow = svgWidth t
  in scale (min (w / ow) (h / oh)) t

-- layer 2
fourierA :: Fourier -> (Double -> Double) -> Animation
fourierA f genPhi = animate $ \t ->
    let circles = setFourierLength (t* fourierLen f) f
        coeffs = fourierCoefficients $ rotateFourier (genPhi t) circles
    in mkGroup
    [ drawCircles coeffs
    , withStrokeColor "green" $
      withStrokeLineJoin JoinRound $
      withFillOpacity 0 $
      withStrokeWidth (defaultStrokeWidth*2) $
      mkLinePath $ mkFourierOutline circles
    , let x :+ y = sum coeffs in
      translate x y $ withFillColor "red" $ mkCircle (defaultStrokeWidth*3)
    ]

drawCircles :: [Complex Double] -> SVG
drawCircles [] = mkGroup []
drawCircles ( x :+ y : xs) =
  translate x y $ drawCircles' xs

drawCircles' :: [Complex Double] -> SVG
drawCircles' circles = mkGroup
    [ worker circles
    , withStrokeColor "black" $
      withStrokeLineJoin JoinRound $
      withFillOpacity 0 $
      mkLinePath [ (x, y) | x :+ y <- scanl (+) 0 circles ]]
  where
    worker [] = None
    worker (x :+ y : rest) =
      let radius = sqrt(x*x+y*y) in
      mkGroup
      [ withStrokeColor "dimgrey" $
        withFillOpacity 0 $
        mkCircle radius
      , translate x y $ worker rest ]

-- layer 1
newtype Fourier = Fourier {fourierCoefficients :: [Complex Double]}

mkSVGLatex :: Tree -> Fourier
mkSVGLatex t = mkFourier $ lineToPoints 500 $
  toLineCommands $ extractPath $ center $ scaleToFit 9.6 7.2 $ t

mkFourierLatex :: T.Text -> Fourier
mkFourierLatex t = mkFourier $ lineToPoints 500 $
  toLineCommands $ extractPath $ center $ scaleToFit 9.6 7.2 $ latex t

fourierLen :: Fourier -> Double
fourierLen f = sum $ map magnitude $ drop 1 $ take 500 $ fourierCoefficients f

pointAtFourier :: Fourier -> Complex Double
pointAtFourier = sum . fourierCoefficients

mkFourier :: [RPoint] -> Fourier
mkFourier points = Fourier $ findCoefficient 0 :
    concat [ [findCoefficient n, findCoefficient (-n)] | n <- [1..] ]
  where
    findCoefficient :: Int -> Complex Double
    findCoefficient n =
        sum [ toComplex point * exp (negate (fromIntegral n) * 2 *pi * i*t) * deltaT
            | (idx, point) <- zip [0::Int ..] points, let t = fromIntegral idx/nPoints ]
    i = 0 :+ 1
    toComplex (V2 x y) = x :+ y
    deltaT = recip nPoints
    nPoints = fromIntegral (length points)

setFourierLength :: Double -> Fourier -> Fourier
setFourierLength _ (Fourier []) = Fourier []
setFourierLength len0 (Fourier (first:lst)) = Fourier $ first : worker len0 lst
  where
    worker _len [] = []
    worker len (c:cs) =
      if magnitude c < len
        then c : worker (len - magnitude c) cs
        else [c * realToFrac (len / magnitude c)]

rotateFourier :: Double -> Fourier -> Fourier
rotateFourier phi (Fourier coeffs) =
    Fourier $ worker coeffs (0::Integer)
  where
    worker [] _ = []
    worker (x:rest) 0 = x : worker rest 1
    worker [left] n = worker [left,0] n
    worker (left:right:rest) n =
      let n' = fromIntegral n in
      left * exp (negate n' * 2 * pi * i * phi') :
      right * exp (n' * 2 * pi * i * phi') :
      worker rest (n+1)
    i = 0 :+ 1
    phi' = realToFrac phi

mkFourierOutline :: Fourier -> [(Double, Double)]
mkFourierOutline fourier =
    [ (x, y)
    | idx <- [0 .. granularity]
    , let x :+ y = pointAtFourier $ rotateFourier (idx/granularity) fourier
    ]
  where
    granularity = 500


{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell   #-}

module Bench where

import           Control.Lens                    (set)
import           Control.Lens.Operators
import           Control.Lens.TH
import           Criterion
import           Criterion.Types                 (Regression (..), Report (..),
                                                  SampleAnalysis (..))
import qualified Data.Binary                     as Bin
import qualified Data.ByteString.Lazy            as B
import           Data.Csv
import qualified Data.Map.Lazy                   as Map
import           Data.Maybe                      (fromJust)
import           Data.Monoid                     ((<>))
import qualified Data.Text                       as T
import qualified Data.Text.Encoding              as T
import           GHC.Generics                    (Generic)
import           Statistics.Resampling.Bootstrap (Estimate (..))

data Bench a = Bench {
      _benchDescription :: T.Text
    , _benchData        :: [(Double, a)]
    }
    deriving (Functor, Generic, Show, Foldable, Traversable)

makeLenses ''Bench

data BenchGroup a = BenchGroup {
      _bgId          :: T.Text
    , _bgDescription :: T.Text
    , _bgBenches     :: [Bench a]
    , _bgXAxisName   :: T.Text
    }
    deriving (Functor, Generic, Show, Foldable, Traversable)

makeLenses ''BenchGroup

instance Bin.Binary a => Bin.Binary (Bench a)
instance Bin.Binary a => Bin.Binary (BenchGroup a)

runBenchGroup :: BenchGroup Benchmarkable -> IO (BenchGroup Report)
runBenchGroup group = do
    putStrLn $ "Benchmarking group: " <> T.unpack (group ^. bgDescription) <>
               " (" <> T.unpack (group ^. bgId) <> ")"
    results <- mapM describeAndRun (group ^. bgBenches)
    return $ set bgBenches results group

describeAndRun :: Bench Benchmarkable -> IO (Bench Report)
describeAndRun bench = do
    putStrLn $ "Running benchmark: " <> T.unpack (bench ^. benchDescription)
    mapM benchAndRetry bench
  where
    benchAndRetry singleBench = do
        result <- benchmark' singleBench
        if getRSquared result < 0.99
            then putStrLn "R^2 < 0.99 => RETRY!" >> benchAndRetry singleBench
            else return result

getOLS :: Report -> Double
getOLS = estPoint . fromJust . Map.lookup "iters" . regCoeffs . head . anRegress . reportAnalysis

getRSquared :: Report -> Double
getRSquared = estPoint . regRSquare . head . anRegress . reportAnalysis

getMean :: Report -> Double
getMean = estPoint . anMean . reportAnalysis

getHeader :: BenchGroup a -> [T.Text]
getHeader group = ["framework", group ^. bgXAxisName, "time"]

getMeasurements :: Bench a -> [Measurement a]
getMeasurements bench =
    map (\(x, y) -> Measurement (bench ^. benchDescription) x y)
        (bench ^. benchData)

data Measurement a = Measurement {
      _mLabel  :: T.Text
    , _mXValue :: Double
    , _mYValue :: a
    }
    deriving (Functor, Generic, Show)

makeLenses ''Measurement

instance ToField a => ToRecord (Measurement a)

toCsv :: ToField a => BenchGroup a -> B.ByteString
toCsv group =
       B.intercalate "," (map (B.fromStrict . T.encodeUtf8) (getHeader group))
    <> "\n"
    <> encode (concatMap getMeasurements (group ^. bgBenches))


{-# LANGUAGE GADTs #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UnboxedTuples #-}

-- |
-- Module      : PSO.Neural
-- Description : Particle Swarm Optimisation
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Neural
  (
  -- mutableMcmcLoop
  -- , energyHH1DOpen
  -- , energyHH1DOpenC
  -- , energyHH1DOpenMKLC
  -- , energyHH1DOpenZ
  -- , energyVarHH1DOpen
  -- , RBM(..)
    Rbm
  , mkRbm
  , uniformRbm
  , IsRbm(..)
  , fromRbm
  , toRbm
  , unsafeFreezeRbm
  , unsafeThawRbm
  , mcmcHeisenberg1D
  -- , MCMC(..)
  -- , unsafeFreezeMcmc
  -- , Mcmc
  -- , HH1DOpen(..)
  -- , listAll
  , Measurement(..)
  , HasMean(..)
  , HasVar(..)
  ) where

import Prelude hiding (map, mapM, zipWith, zipWithM)
import qualified Prelude as Prelude

import Control.Exception(assert)
import Control.Monad (Monad(..), (>=>), (<=<))
import Control.Monad.Primitive
import Control.Monad.ST

import Data.Bits
import Data.Coerce
import Data.Complex
import Data.Functor.Identity
import Data.NumInstances.Tuple
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..), SPEC(..))
import qualified Data.Vector.Fusion.Stream.Monadic as Stream
import qualified Data.Vector.Fusion.Bundle.Monadic as Bundle

import Foreign.C.Types
import Foreign.Storable
import Foreign.ForeignPtr

import GHC.Generics (Generic)

import Lens.Micro
import Lens.Micro.TH

import System.IO.Unsafe

import PSO.Internal.Neural
import PSO.Random
import PSO.Swarm

type family RealOf x

type instance RealOf Float = Float

type instance RealOf Double = Double

type instance RealOf CFloat = CFloat

type instance RealOf CDouble = CDouble

type instance RealOf (Complex a) = a

type family RbmCore a

type instance RbmCore Float = RbmC

-- type instance RbmCore (Complex Double) = RbmZ

newtype Rbm a = R (ForeignPtr (RbmCore a)) deriving (Generic)

newtype MRbm s a = MR (ForeignPtr (RbmCore a)) deriving (Generic)

class Storable a => IsRbm a where
  newRbm :: PrimMonad m => Int -> Int -> m (MRbm (PrimState m) a)
  cloneRbm :: PrimMonad m => Rbm a -> m (MRbm (PrimState m) a)
  getWeights :: PrimMonad m => MRbm (PrimState m) a -> m (V.MVector (PrimState m) a, V.MVector (PrimState m) a)
  getVisible :: PrimMonad m => MRbm (PrimState m) a -> m (V.MVector (PrimState m) a, V.MVector (PrimState m) a)
  getHidden  :: PrimMonad m => MRbm (PrimState m) a -> m (V.MVector (PrimState m) a, V.MVector (PrimState m) a)
  sizeWeights :: Rbm a -> Int
  sizeVisible :: Rbm a -> Int
  sizeHidden  :: Rbm a -> Int
  -- debugPrintRbm :: Rbm a -> IO ()

unsafeThawRbm :: PrimMonad m => Rbm a -> m (MRbm (PrimState m) a)
unsafeThawRbm (R p) = return (MR p)

unsafeFreezeRbm :: PrimMonad m => MRbm (PrimState m) a -> m (Rbm a)
unsafeFreezeRbm (MR p) = return (R p)

-- setWeights :: PrimMonad m => MRbm (PrimState m) a -> V.Vector a -> V.Vector a -> m ()
-- setWeights mrbm wR wI = do
--   (mwR, mwI) <- getWeights mrbm
--   V.copy mwR wR
--   V.copy mwI wI
--
-- setVisible :: PrimMonad m => MRbm (PrimState m) a -> V.Vector a -> V.Vector a -> m ()
-- setVisible mrbm aR aI = do
--   (maR, maI) <- getVisible mrbm
--   V.copy maR aR
--   V.copy maI aI
--
-- setHidden :: PrimMonad m => MRbm (PrimState m) a -> V.Vector a -> V.Vector a -> m ()
-- setHidden mrbm bR bI = do
--   (mbR, mbI) <- getHidden mrbm
--   V.copy mbR bR
--   V.copy mbI bI

fromRbm :: (Monad m, IsRbm a) => Rbm a -> Stream.Stream m a
fromRbm rbm = Bundle.elements $ Bundle.fromVectors (runST $ unsafeThawRbm rbm >>= go)
  where
    go :: IsRbm a => MRbm s a -> ST s [V.Vector a]
    go mrbm = do
      (wR, wI) <- getWeights mrbm
      (aR, aI) <- getVisible mrbm
      (bR, bI) <- getHidden mrbm
      Prelude.mapM V.freeze [wR, wI, aR, aI, bR, bI]

toRbm :: (PrimMonad m, IsRbm a) => Int -> Int -> Stream.Stream m a -> m (Rbm a)
toRbm n m (Stream step t) = do
  mrbm <- newRbm n m
  (wR, wI) <- getWeights mrbm
  (aR, aI) <- getVisible mrbm
  (bR, bI) <- getHidden mrbm
  copy SPEC wR (n * m) 0 t >>=
    copy SPEC wI (n * m) 0 >>=
      copy SPEC aR n 0 >>=
        copy SPEC aI n 0 >>=
          copy SPEC bR m 0 >>=
            copy SPEC bI m 0
  unsafeFreezeRbm mrbm
  where
    copy !_ !v !size !i s
      | i < size = step s >>= \r ->
        case r of
          Yield x s' -> MV.write v i x >> copy SPEC v size (i + 1) s'
          Skip    s' -> copy SPEC v size i s'
          Done       -> error "No, we're not done yet!"
      | otherwise = return s

-- | Creates a new Restricted Boltzmann Machine with given weights.
mkRbm :: IsRbm a
      => V.Vector (Complex a) -- ^ Visible bias \(a^N\).
      -> V.Vector (Complex a) -- ^ Hidden bias \(a^M\).
      -> V.Vector (Complex a) -- ^ Weights \(a^{M\times N}\). *Column-major layout is assumed!*.
      -> Rbm a
mkRbm a@(V.length -> nx1) b@(V.length -> mx1) w@(V.length -> mxn)
  | mxn == nx1 * mx1 = runST $ toRbm nx1 mx1 $ Bundle.elements $
    Bundle.fromVectors [ realPart `V.map` w, imagPart `V.map` w
                       , realPart `V.map` a, imagPart `V.map` a
                       , realPart `V.map` b, imagPart `V.map` b
                       ]
  | otherwise = error "mkRbm: Dimensions mismatch."

-- | Creates an RBM with weights distributed uniformly within given intervals.
uniformRbm ::
     (PrimMonad m, IsRbm a, Ord a, UniformDist m (Complex a))
  => Int -- ^ Number of visible nodes.
  -> Int -- ^ Number of hidden nodes.
  -> (a, a) -- ^ Range for visible bias.
  -> (a, a) -- ^ Range for hidden bias.
  -> (a, a) -- ^ Range for weights.
  -> m (Rbm a)
uniformRbm !n !m !(amin, amax) !(bmin, bmax) !(wmin, wmax)
  | n < 0 || m < 0 = error "uniformRbm: Negative dimensions??"
  | amin > amax || bmin > bmax || wmin > wmax =
    error "uniformRbm: Lower bound is per definition not \
          \greater than the upper bound."
  | otherwise = do
    a <- uniformVector n (amin :+ amin, amax :+ amin)
    b <- uniformVector m (bmin :+ bmin, bmin :+ bmax)
    w <- uniformVector (n * m) (wmin :+ wmin, wmax :+ wmax)
    return $ mkRbm a b w

instance IsRbm Float where
  {-# INLINE newRbm #-}
  newRbm n m = unsafeIOToPrim $ MR <$> _RbmC'construct n m
  {-# INLINE cloneRbm #-}
  cloneRbm (R p) = unsafeIOToPrim $ MR <$> _RbmC'clone p
  {-# INLINE getWeights #-}
  getWeights (MR p) = _RbmC'getWeights p
  {-# INLINE getVisible #-}
  getVisible (MR p) = _RbmC'getVisible p
  {-# INLINE getHidden #-}
  getHidden (MR p) = _RbmC'getHidden p
  {-# INLINE sizeVisible #-}
  sizeVisible (R fp) = unsafePerformIO $
    withForeignPtr fp (peek >=> return . _rbmC'sizeVisible)
  {-# INLINE sizeHidden #-}
  sizeHidden (R fp) = unsafePerformIO $
    withForeignPtr fp (peek >=> return . _rbmC'sizeHidden)
  {-# INLINE sizeWeights #-}
  sizeWeights x = sizeHidden x * sizeVisible x

getDims :: IsRbm a => Rbm a -> (Int, Int)
getDims x = (sizeVisible x, sizeHidden x)
{-# INLINE getDims #-}

zipWithM :: (IsRbm a, IsRbm b, IsRbm c, PrimMonad m)
         => (a -> b -> m c) -> Rbm a -> Rbm b -> m (Rbm c)
zipWithM f x y = assert (getDims x == getDims y) $
  toRbm (sizeVisible x) (sizeHidden x) $ Stream.zipWithM f (fromRbm x) (fromRbm y)

zipWith :: forall a b c. (IsRbm a, IsRbm b, IsRbm c)
        => (a -> b -> c) -> Rbm a -> Rbm b -> Rbm c
zipWith f x y = runST $ zipWithM f' x y
  where f' :: a -> b -> ST s c
        f' a b = return $ f a b

mapM :: (IsRbm a, IsRbm b, PrimMonad m)
     => (a -> m b) -> Rbm a -> m (Rbm b)
mapM f x = toRbm (sizeVisible x) (sizeHidden x) $ Stream.mapM f $ fromRbm x

map :: (IsRbm a, IsRbm b)
     => (a -> b) -> Rbm a -> Rbm b
map f x = runST $ mapM (return . f) x

instance Num (Rbm Float) where
  (+) x y = zipWith (+) x y
  (-) x y = zipWith (-) x y
  (*) x y = zipWith (*) x y
  abs x = map abs x
  signum x = map signum x
  negate x = map negate x
  fromInteger = undefined

instance Scalable Float (Rbm Float) where
  scale λ = map (*λ)

instance (PrimMonad m, DeltaWell m Float Float)
  => DeltaWell m Float (Rbm Float) where
    upDeltaWell κ p x = zipWithM (upDeltaWell κ) p x

{-
instance RBM (Complex Double) where
  newRbm nrVis nrHid = unsafeIOToPrim $ liftM MR (_newRbmZ nrVis nrHid)
  cloneRbm (R p) = unsafeIOToPrim $ liftM MR (_cloneRbmZ p)
  setWeights (MR p) w = unsafeIOToPrim $ _setWeightsZ p (coerce w)
  setVisible (MR p) a = unsafeIOToPrim $ _setVisibleZ p (coerce a)
  setHidden (MR p) b = unsafeIOToPrim $ _setHiddenZ p (coerce b)
  sizeVisible (R p) = _sizeVisibleZ p
  sizeHidden (R p) = _sizeHiddenZ p
  debugPrintRbm (R p) = _printRbmZ p

instance AXPBY (Complex Float) (Complex Float) where
  axpby a (R x) b (MR y) = unsafeIOToPrim $ _caxpbyRbmC (coerce a) x (coerce b) y

instance AXPBY (Complex Double) (Complex Double) where
  axpby a (R x) b (MR y) = unsafeIOToPrim $ _zaxpbyRbmZ (coerce a) x (coerce b) y

instance SCAL (Complex Float) (Complex Float) where
  scal a (MR x) = unsafeIOToPrim $ _cscaleRbmC (coerce a) x

instance SCAL Float (Complex Float) where
  scal a (MR x) = unsafeIOToPrim $ _cscaleRbmC (coerce a :+ 0) x

instance SCAL (Complex Double) (Complex Double) where
  scal a (MR x) = unsafeIOToPrim $ _zscaleRbmZ (coerce a) x

instance SCAL Double (Complex Double) where
  scal a (MR x) = unsafeIOToPrim $ _zscaleRbmZ (coerce a :+ 0) x

instance (Num a, AXPBY a a, SCAL a a) => Num (Rbm a) where
  (+) x y = runST $ cloneRbm x >>= \z ->
    axpby (1 :: a) y (1 :: a) z >> unsafeFreezeRbm z
  (-) x y = runST $ cloneRbm x >>= \z ->
    axpby ((-1) :: a) y (1 :: a) z >> unsafeFreezeRbm z
  (*) x y = undefined
  negate x = runST $ cloneRbm x >>= \z -> scal ((-1) :: a) z >> unsafeFreezeRbm z
  fromInteger = undefined
  abs = undefined
  signum = undefined

instance (SCAL λ a) => Scalable λ (Rbm a) where
  scale k x = runST $ cloneRbm x >>= \z -> scal k z >> unsafeFreezeRbm z

instance MCMC (Complex Float) where
  newMcmc (R rbm) spin = unsafeIOToPrim $
    _newMcmcC rbm (coerce spin) >>= \mcmc -> return (MM mcmc rbm)
  logWF (M p _) = coerce $ _logWFC p
  logQuotient1 (M p _) flip1 = coerce $ _logQuotientWF1C p flip1
  logQuotient2 (M p _) flip1 flip2 = coerce $ _logQuotientWF2C p flip1 flip2
  propose1 m@(M p _) = (^2) . magnitude . exp . logQuotient1 m -- _proposeC1 p flip1
  propose2 (M p _) flip1 flip2 = _propose2C p flip1 flip2
  update1 (MM p _) flip1 = unsafeIOToPrim $ _accept1C p flip1
  update2 (MM p _) flip1 flip2 = unsafeIOToPrim $ _accept2C p flip1 flip2
  debugPrintMcmc (M p _) = _printMcmcC p

instance MCMC (Complex Double) where
  newMcmc (R rbm) spin = unsafeIOToPrim $
    _newMcmcZ rbm (coerce spin) >>= \mcmc -> return (MM mcmc rbm)
  logWF (M p _) = coerce $ _logWFZ p
  logQuotient1 (M p _) flip1 = coerce $ _logQuotientWF1Z p flip1
  logQuotient2 (M p _) flip1 flip2 = coerce $ _logQuotientWF2Z p flip1 flip2
  propose1 m@(M p _) = (^2) . magnitude . exp . logQuotient1 m -- _proposeZ1 p flip1
  propose2 (M p _) flip1 flip2 = _propose2Z p flip1 flip2
  update1 (MM p _) flip1 = unsafeIOToPrim $ _accept1Z p flip1
  update2 (MM p _) flip1 flip2 = unsafeIOToPrim $ _accept2Z p flip1 flip2
  debugPrintMcmc (M p _) = _printMcmcZ p

mutableMcmcLoop ::
     (PrimMonad m)
  => (α -> Int -> m β)
  -> (α -> β -> m Bool)
  -> (α -> β -> m ())
  -> α
  -> Int
  -> Producer α m ()
mutableMcmcLoop propose accept update ψ n = do
  proposal <- lift $ propose ψ n
  isAcceptable <- lift $ accept ψ proposal
  yield ψ
  when isAcceptable (lift (update ψ proposal))
  mutableMcmcLoop propose accept update ψ (n + 1)

class MCMC a => HH1DOpen a where
  locEnergyHH1DOpen :: Mcmc a -> a

instance HH1DOpen (Complex Float) where
  locEnergyHH1DOpen (M p _) = coerce (_locEHH1DOpenC p)

instance HH1DOpen (Complex Double) where
  locEnergyHH1DOpen (M p _) = coerce (_locEHH1DOpenZ p)

unpack2Bools :: (Num α, Bits α) => Int -> α -> V.Vector Bool
unpack2Bools n x = V.fromList $ doUnpack n x []
  where doUnpack 0 _ []   = [False]
        doUnpack 0 _ bits = bits
        doUnpack n x bits = doUnpack (n - 1) (x `shiftR` 1) $
                              (testBit x 0) : bits

listAll :: forall m a.
     ( PrimMonad m
     , HH1DOpen a
     , Num a
     )
  => Rbm a -> m (V.Vector a)
listAll rbm =
  let n = sizeVisible rbm
      toSpin True  = 1
      toSpin False = -1
      σs :: [V.Vector a]
      σs = V.map toSpin <$> unpack2Bools n <$> [(0 :: Int) .. (2 ^ n - 1)]
   in V.fromList <$> (newMcmc rbm >=> unsafeFreezeMcmc >=> return . logWF) `mapM` σs

energyHH1DOpen ::
     forall m a λ.
     ( PrimMonad m
     , Ord (RealOf a)
     , Fractional a
     , HH1DOpen a
     , Randomisable m Bool
     , Randomisable m (RealOf a)
     , UniformDist m Int
     -- , MonadIO m
     -- , a ~ Complex λ
     -- , Fractional λ
     -- , Real λ
     -- , Show λ
     )
  => Rbm a -> Int -> Int -> m a
energyHH1DOpen rbm offset steps =
  let n = sizeVisible rbm
      proposer :: MMcmc (PrimState m) a -> Int -> m Int
      proposer _ i = {- return (i `mod` n) -} uniform (0, n - 1)
      acceptor :: MMcmc (PrimState m) a -> Int -> m Bool
      acceptor mx flip = unsafeFreezeMcmc mx >>= \x ->
                            ((<= propose1 x flip) <$> random)
      updater :: MMcmc (PrimState m) a -> Int -> m ()
      updater = update1
      func :: MMcmc (PrimState m) a -> m a
      func = unsafeFreezeMcmc >=> return . locEnergyHH1DOpen
      mmcmc :: m (MMcmc (PrimState m) a)
      mmcmc = randomSpin n >>= newMcmc rbm
      states = mutableMcmcLoop proposer acceptor updater
   in mmcmc >>= \mmcmc' ->
                  (/ fromIntegral steps)
                    <$> P.sum
                  --     P.fold debugAdd
                  --            (toRational 0, toRational 0)
                  --            (\(x, y) -> fromRational (x / toRational steps)
                  --                         :+ fromRational (y / toRational steps))
                             (states mmcmc' 0 >-> P.drop offset
                                              >-> P.mapM func
                                              >-> P.take steps)

energyHH1DOpenC :: forall m.
     ( PrimMonad m
     , Randomisable m Bool
     , Randomisable m Float
     , UniformDist m Int
     )
  => Rbm (Complex Float) -> Int -> Int -> m (Complex Float)
energyHH1DOpenC rbm offset steps =
  let n = sizeVisible rbm
      mmcmc = randomSpin n >>= newMcmc rbm
      ints :: m (V.Vector Int)
      ints = uniformVector (offset + steps) (0, n - 1)
      floats :: m (V.Vector Float)
      floats = V.replicateM (offset + steps) random
   in ints >>= \ is ->
        floats >>= \ fs ->
          mmcmc >>= \ (MM p (RbmC m)) ->
            unsafeIOToPrim $ (_runMcmcC p offset steps is fs) >>= \ x ->
              touchForeignPtr m >> return (coerce x)

energyHH1DOpenZ :: forall m.
     ( PrimMonad m
     , Randomisable m Bool
     , Randomisable m Double
     , UniformDist m Int
     )
  => Rbm (Complex Double) -> Int -> Int -> m (Complex Double)
energyHH1DOpenZ rbm offset steps =
  let n = sizeVisible rbm
      mmcmc = randomSpin n >>= newMcmc rbm
      ints :: m (V.Vector Int)
      ints = uniformVector (offset + steps) (0, n - 1)
      floats :: m (V.Vector Double)
      floats = V.replicateM (offset + steps) random
   in ints >>= \ is ->
        floats >>= \ fs ->
          mmcmc >>= \ (MM p (RbmZ m)) ->
            unsafeIOToPrim $ (_runMcmcZ p offset steps is fs) >>= \ x ->
              touchForeignPtr m >> return (coerce x)
-}

mcmcHeisenberg1D ::
     PrimMonad m => Rbm Float -> Int -> Int -> m (Measurement Float)
mcmcHeisenberg1D (R rbm) offset steps = unsafeIOToPrim $ _RbmC'heisenberg1D rbm offset steps

instance Eq a => Eq (Measurement a) where
  (==) (Measurement _ x) (Measurement _ y) = x == y

instance Ord a => Ord (Measurement a) where
  -- (<=) (Measurement a x) (Measurement b y) = x <= y
  (<=) (Measurement a x) (Measurement b y) = a <= b

instance Num a => Num (Measurement a) where
  (+) (Measurement a b) (Measurement c d) = Measurement (a + c) (b + d)
  (-) (Measurement a b) (Measurement c d) = Measurement (a - c) (b - d)
  (*) (Measurement a b) (Measurement c d) = Measurement (a * c) (b * d)
  negate (Measurement a b) = Measurement (negate a) (negate b)
  fromInteger n = Measurement (fromInteger n) (fromInteger n)
  abs (Measurement a b) = Measurement (abs a) (abs b)
  signum (Measurement a b) = Measurement (signum a) (signum b)

instance Fractional a => Fractional (Measurement a) where
  (/) (Measurement a b) (Measurement c d) = Measurement (a / c) (b / d)
  recip (Measurement a b) = Measurement (recip a) (recip b)

makeFields ''Measurement

{-

debugAdd :: forall λ. (Real λ, Fractional λ, Show λ)
         => (λ, λ) -> (λ, λ) -> (λ, λ)
debugAdd (x, y) (a, b) = unsafePerformIO $ do
  putStrLn $ show x ++ "\t" ++ show y ++ "\t" ++ show a ++ "\t" ++ show b
  return (x + a, y + b)

energyVarHH1DOpen ::
     forall m a.
     ( PrimMonad m
     , Ord a
     , RealFloat a
     , HH1DOpen (Complex a)
     , Randomisable m Bool
     , Randomisable m a
     , UniformDist m Int
     -- , MonadIO m
     , Show a
     )
  => Rbm (Complex a) -> Int -> Int -> m (Measurement a)
energyVarHH1DOpen rbm offset steps =
  let n = sizeVisible rbm
      proposer :: MMcmc (PrimState m) (Complex a) -> Int -> m Int
      proposer _ i = return (i `mod` n) {- uniform (0, n - 1) -}
      acceptor :: MMcmc (PrimState m) (Complex a) -> Int -> m Bool
      acceptor mx flip = unsafeFreezeMcmc mx >>= \x ->
                            ((<= propose1 x flip) <$> random)
      updater :: MMcmc (PrimState m) (Complex a) -> Int -> m ()
      updater = update1
      func :: MMcmc (PrimState m) (Complex a) -> m (a, a)
      func = unsafeFreezeMcmc >=> return . locEnergyHH1DOpen
                              >=> return . (\ x -> (realPart x, (realPart x)^2))
      mmcmc :: m (MMcmc (PrimState m) (Complex a))
      mmcmc = randomSpin n >>= newMcmc rbm
      states = mutableMcmcLoop proposer acceptor updater
   in mmcmc >>= \mmcmc' ->
                  P.fold debugAdd
                         0
                         (\(h, h2) -> let mean = (h  / fromIntegral steps)
                                          var  = (h2 / fromIntegral steps) - mean^2
                                       in MV mean var)
                         (states mmcmc' 0 >-> P.drop offset
                                          >-> P.mapM func
                                          >-> P.take steps)
-}


module Main where

import           Control.Concurrent
import           Control.Concurrent.Chan
import           Control.Monad.Writer
import           Data.Array
import           Data.Bits
import qualified Data.ByteString           as B
import           Data.Complex
import           Data.Function
import           Data.List
import           Data.Word
import           Network.Socket            hiding (sendTo)
import           Network.Socket.ByteString
import           System.IO

artWorkLeft :: Rythm
artWorkLeft = Rest 1 :+: Hit 1 1 :+: incrAmp  (incrDur artWorkLeft)

artWorkRight :: Rythm
artWorkRight = Hit 1 1 :+: Rest 1 :+: incrAmp (incrDur artWorkRight)

artWorkMiddle :: Rythm
artWorkMiddle = Hit 1 1 :+: Hit 1 1 :+: go 1 True
  where
    go :: Int -> Bool -> Rythm
    go dur up =
        Rest 1 :+: Hit 119 dur :+: go dur' up'
      where
        dur' =
            if up
                then dur + 1
                else dur - 1
        up' = up && dur < 119 || (not up && dur < 2)

decrAmpN 0 r = r
decrAmpN n r = r' :+: decrAmpN (n - 1) r'
  where
    r' = decrAmp r

incrAmpN 0 r = r
incrAmpN n r = r' :+: incrAmpN (n-1) r'
  where
    r' = incrAmp r

nTimes r 1 = r
nTimes r n = r :+: nTimes r (n-1)

ssht :: Rythm
ssht = Hit 12 1

incrDur (Hit a d) = Hit a (cap 1 120 (d + 1))
incrDur (r1 :+: r2) = incrDur r1 :+: incrDur r2
incrDur ( = r

decrAmp :: Rythm -> Rythm
decrAmp (Hit a d) | a > 1 = Hit (a - 1) d
                  | otherwise = Hit a d
decrAmp (r1 :+: r2) = decrAmp r1 :+: decrAmp r2
decrAmp r = r

incrAmp :: Rythm -> Rythm
incrAmp (Hit a d) | a < 120 = Hit (a + 1) d
                  | otherwise = Hit a d
incrAmp (r1 :+: r2) = incrAmp r1 :+: incrAmp r2
incrAmp r = r

data Action
    = Render Track
    | Stop
    deriving (Show)

ipLeft = "2001:67c:20a1:1095:ba27:ebff:feb9:db12"
ipMiddle = "2001:67c:20a1:1095:ba27:ebff:fe23:60d7"
ipRight = "2001:67c:20a1:1095:ba27:ebff:fe71:dd32"

-- small panel: 16x20
-- panel: 3 x 6 small: 48 x 120

data Panel = P Track (Chan Action)

main :: IO ()
main = do
    l <- startPanel 'L' artWorkLeft (mandelBrot 0 newImage) ipLeft
    m <- startPanel 'M' artWorkMiddle (mandelBrot 49 newImage) ipMiddle
    r <- startPanel 'R' artWorkRight (mandelBrot 97 newImage) ipRight
    hSetBuffering stdin NoBuffering
    hSetBuffering stdout NoBuffering
    xxxLoop l m r
    hSetBuffering stdin LineBuffering
    hSetBuffering stdout LineBuffering
    stopPanel l
    stopPanel m
    stopPanel r

xxxLoop :: Panel -> Panel -> Panel -> IO ()
xxxLoop l m r = do
    i <- readInput
    processInput i l m r

processInput :: Input -> Panel -> Panel -> Panel -> IO ()
processInput Quit _ _ _ = return ()
processInput Tick l m r = do
    let l' = updatePanel l
        m' = updatePanel m
        r' = updatePanel r
    renderToPanel l
    renderToPanel m
    renderToPanel r
    xxxLoop l' m' r'

updatePanel :: Panel -> Panel
updatePanel (P t c) = P (dropCurrent t) c
  where
    dropCurrent (_ :>>: t') = t'
    dropCurrent _ = Done

data Input
    = Quit
    | Tick

readInput :: IO Input
readInput = do
    c <- getChar
    return
        (case c of
             ' ' -> Tick
             _ -> Quit)

startPanel :: Char -> Rythm -> Image -> String -> IO Panel
startPanel name artWork initialImage host = do
    c <- newChan
    ai <- head <$> getAddrInfo Nothing (Just host) (Just "2323")
    s <- socket (addrFamily ai) Datagram defaultProtocol
    void $ forkIO (panelLoop name (addrAddress ai) c s)
    return (P (renderRythm artWork initialImage) c)

stopPanel :: Panel -> IO ()
stopPanel = sendAction Stop

renderToPanel :: Panel -> IO ()
renderToPanel p@(P t _) = sendAction (Render t) p

sendAction :: Action -> Panel -> IO ()
sendAction a (P _ c) = writeChan c a

panelLoop :: Char -> SockAddr -> Chan Action -> Socket -> IO ()
panelLoop name a c s = do
    putStr [name, '?']
    action <- readChan c
    case action of
        Stop -> do
            putStr [name, '$']
            return ()
        Render t -> do
            panelRenderNext t
            panelLoop name a c s
  where
    panelRenderNext t =
        case t of
            Done -> do
                putStr [name, 'X']
                return ()
            Sleep d -> do
                putStr [name, '<']
                threadDelay (durationToMicros d)
                putStr [name, '>']
            Draw i -> do
                putStr [name, '#']
                void $ sendTo s (convertToDisplay i) a
            t' :>>: _ -> panelRenderNext t'

-- * MUSIC API

type Amplitude = Int
type Duration = Int

durationToMicros :: Int -> Int
durationToMicros = (microsPerCol *)
  where
    microsPerCol = 1500000 `div` panelW

data Rythm = Hit Amplitude Duration | Rest Duration | Rythm :+: Rythm
    deriving (Show)

type Micros = Int
data Track = Done | Sleep Micros | Draw Image | Track :>>: Track
    deriving (Show)

instance Monoid Track where
    mempty = Done
    mappend Done t = t
    mappend t Done = t
    mappend (ll :>>: lr) r = ll :>>: mappend lr r
    mappend l r = l :>>: r

renderRythm :: Rythm -> Image -> Track
renderRythm r orig = execWriter (go r orig)
  where
    go :: Rythm -> Image -> Writer Track Image
    go (r1 :+: r2) i = do
        i' <- go r1 i
        go r2 i'
    go (Rest d) i = do
        tell (Sleep (cap 1 120 d))
        return i
    go (Hit a d) i = do
        let i' = flipBlock d a i
        tell (Draw i')
        tell (Sleep (cap 1 120 d))
        return i'

cap l u t | l <= t && t <= u = t
          | l <= t         = u
          | otherwise     = l

-- * Image API
type Image = Array (Int,Int) Bool

panelW, panelH :: Int
panelW = 48
panelH = 120

capCoords (x, y) = (cap 1 48 x, cap 1 120 y)

mandelBrot :: Int -> Image -> Image
mandelBrot offset img = foldr mandelPixel img (indices img)
  where
    mandelPixel (x,y) = setPixel x y (inMandelSet (0.0 :+ 0.0) 0)
      where
        maxIter = 100
        c :: Complex Double
        c =
            (fromIntegral (offset + x) * (4.0 / 144.0) - 2.0) :+
            (fromIntegral y * (4.0 / 120.0) - 2.0)
        inMandelSet z i
          | ((magnitude z) <= 2.0) && i < maxIter =
              inMandelSet (z * z + c) (i + 1)
          | otherwise = rem i 2 == 1

newImage :: Image
newImage =
    array
        ((1, 1), (panelW, panelH))
        [((x, y), False) | x <- [1 .. panelW]
                         , y <- [1 .. panelH]]

flipBlock width height img = foldr (flipCol height) img [1 .. width]

flipCol height col img = foldr flipPixel img [(col, y) | y <- [1..height]]

flipPixel (x,y) img = setPixel x y (not (getPixel x y img)) img

setPixel :: Int -> Int -> Bool -> Image -> Image
setPixel x y v img = img // [((x, y), v)]

getPixel :: Int -> Int -> Image -> Bool
getPixel x y img = img ! (x,y)

convertToDisplay :: Image -> B.ByteString
convertToDisplay img = pixels
  where
    pixels = B.pack $ toFlipDotFormat $ toBit <$> elems img
    toFlipDotFormat :: [Word8] -> [Word8]
    toFlipDotFormat x = go x []
      where
        go x acc =
            case take 8 x of
              [] -> reverse acc
              sx -> go (drop 8 x) (toByte (reverse sx) 1 0 : acc)
    toByte [] _ acc =  acc
    toByte (b:rest) p acc = toByte rest (p*2) (acc .|. b * p)
    toBit :: Bool -> Word8
    toBit = fromIntegral . fromEnum

    (_,(w,h)) = bounds img

toPixel :: Int -> Char
toPixel = toEnum


{-# LANGUAGE FlexibleContexts #-}
-- {-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE LambdaCase #-}
module Main where

import           Common                         ( Analysis(..)
                                                , Path(..)
                                                )
import           Control.Monad                  ( foldM
                                                , forM
                                                , forM_
                                                , replicateM
                                                , zipWithM_
                                                )
import           Control.Monad.Except           ( runExceptT )
import           Data.Either                    ( rights )
import           Data.List                      ( unzip4 )
import qualified Data.List                     as L
import           Data.Maybe                     ( catMaybes
                                                , listToMaybe
                                                , mapMaybe
                                                )
import qualified Data.Vector                   as V
import           GHC.Float                      ( int2Double )
import qualified Graphics.Matplotlib           as Plt
import qualified GreedyParser                  as Greedy
import           Inference.Conjugate            ( Hyper
                                                , HyperRep
                                                , Prior(sampleProbs)
                                                , Probs
                                                , ProbsRep
                                                , Trace
                                                , evalTraceLogP
                                                , evalTracePredLogP
                                                , getPosterior
                                                , jeffreysPrior
                                                , traceTrace
                                                , uniformPrior
                                                )
import           Musicology.Pitch               ( SPitch )
import qualified Numeric.Log                   as Log
import           PVGrammar                      ( Edge
                                                , PVAnalysis
                                                , loadAnalysis
                                                , loadInput
                                                , topEdges
                                                )
import           PVGrammar.Parse                ( protoVoiceEvaluator
                                                , pvCountUnrestricted
                                                )
import           PVGrammar.Prob.Simple          ( PVParams(PVParams)
                                                , observeDerivation'
                                                , sampleDerivation'
                                                )
import qualified Parser
import qualified Statistics.Sample             as Stats
import           System.FilePath                ( (<.>)
                                                , (</>)
                                                )
import qualified System.FilePattern            as FP
import qualified System.FilePattern.Directory  as FP
import qualified System.Random.MWC.Probability as MWC
import           System.Random.Stateful         ( StatefulGen
                                                , initStdGen
                                                , newIOGenM
                                                )
import           Text.Pretty.Simple             ( CheckColorTty(..)
                                                , OutputOptions(..)
                                                , defaultOutputOptionsNoColor
                                                , pPrintOpt
                                                )
import           Text.Printf                    ( printf )

prettyPrint :: (Show a) => a -> IO ()
prettyPrint = pPrintOpt
  NoCheckColorTty
  (defaultOutputOptionsNoColor { outputOptionsCompact = True })

dataDir :: FilePath
dataDir = "data/"

main = mainLearn

nBaselines :: Int
nBaselines = 100 -- default: 100

nSamples :: Int
nSamples = 250 -- default: 250

mainLearn :: IO ()
mainLearn = do
  -- initialize
  genPure <- initStdGen
  gen     <- newIOGenM genPure
  genMWC  <- MWC.create -- uses a fixed seed
  let prior = uniformPrior @PVParams
  -- load data
  articleExamples <- loadDir
    (dataDir </> "theory-article")
    ["05b_cello_prelude_1-4", "09a_hinunter", "03_bwv784_pattern"]
  Just bwv939 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0939"
  Just bwv940 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0940"
  let dataset = bwv939 : bwv940 : articleExamples
  -- let dataset = take 3 articleExamples
  putStrLn "list of pieces:"
  forM_ dataset $ \(name, _ana, _trace, _surface) -> do
    putStrLn $ "  " <> name
  -- compute overall posterior
  posteriorTotal <- learn prior dataset
  prettyPrint posteriorTotal
  -- cross validation
  let splits = leaveOneOut dataset
  crossPerps <- forM splits (comparePerNote gen genMWC prior)
  let (logpPriors, logpPosts, counts, basemeans, baselines) =
        L.unzip5 crossPerps
      count           = sum counts
      logppnPrior     = sum logpPriors / count
      logppnTrained   = sum logpPosts / count
      logppnBaseline  = sum basemeans / count
      -- look at each split
      logppns         = zipWith (/) logpPosts counts -- fmap (\(_, logp, count, _, _) -> logp / count) crossPerps
      baselinelogppns = zipWith (\bs n -> (/ n) <$> bs) baselines counts
      testpieces      = fst <$> splits
      testnames       = (\(name, _, _, _) -> name) <$> testpieces
  showSplits logppns baselinelogppns counts testnames
      -- nsplits         = int2Double $ length splits
      -- meanlogppn      = sum logppns / nsplits
  putStrLn $ "prior logppn: " <> show logppnPrior
  putStrLn $ "prior perppn: " <> show (exp $ negate logppnPrior)
  putStrLn $ "overall trained logppn: " <> show logppnTrained
  putStrLn $ "overall trained perppn: " <> show (exp $ negate logppnTrained)
  putStrLn $ "baseline logppn: " <> show logppnBaseline
  putStrLn $ "baseline perppn: " <> show (exp $ negate logppnBaseline)
  putStrLn $ "prior-posterior Δlogppn:" <> show (logppnPrior - logppnTrained)
  putStrLn $ printf "logppn (nats) & %.3f & %.3f & %.3f\\\\"
                    logppnPrior
                    logppnTrained
                    logppnBaseline
  let natsbits = log 2
  putStrLn $ printf "logppn (bits) & %.3f & %.3f & %.3f\\\\"
                    (logppnPrior / natsbits)
                    (logppnTrained / natsbits)
                    (logppnBaseline / natsbits)
  putStrLn $ printf "perppn & %.2f & %.2f & %.2f\\\\"
                    (exp $ negate logppnPrior)
                    (exp $ negate logppnTrained)
                    (exp $ negate logppnBaseline)

testEstimator = do
  (long, short)   <- loadExamples
  articleExamples <- loadDir
    (dataDir </> "theory-article")
    ["05b_cello_prelude_1-4", "09a_hinunter", "03_bwv784_pattern"]
  Just bwv939 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0939"
  Just bwv940 <- loadItem (dataDir </> "bach" </> "fünf-kleine-präludien")
                          "BWV_0940"
  let dataset = bwv939 : bwv940 : articleExamples
  genMWC <- MWC.createSystemRandom
  let prior = uniformPrior @PVParams
  posteriorTotal <- learn prior dataset
  estimates      <- derivationLogProb genMWC posteriorTotal long
  -- let means = meanOfLogs <$> tail (L.inits estimates)
  -- -- print means
  -- Plt.onscreen $ Plt.line [1 .. length means] means
  print $ meanOfLogs estimates

countRareIntDerivs = do
  Just (_, _, _, surface) <- loadItem (dataDir </> "theory-article")
                                      "10c_rare_int"
  count <- Parser.parseSize pvCountUnrestricted surface
  print count

-- loading data
-- ------------

type Piece
  = (String, PVAnalysis SPitch, Trace PVParams, Path [SPitch] [Edge SPitch])

loadItem :: FilePath -> FilePath -> IO (Maybe Piece)
loadItem dir name = do
  ana <- loadAnalysis (dir </> name <.> "analysis.json")
  case ana of
    Left  _err -> pure Nothing
    Right a    -> if anaTop a == PathEnd topEdges
      then do
        surface <- loadInput (dir </> name <.> "musicxml")
        case observeDerivation' (anaDerivation a) of
          Left _err -> do
            putStrLn $ "could not observe trace for " <> name <> ", skipping."
            pure Nothing
          Right trace -> pure $ Just (name, a, trace, surface)
      else do
        putStrLn $ "derivation for " <> name <> " is incomplete, skipping."
        pure Nothing

loadDir :: FilePath -> [String] -> IO [Piece]
loadDir dir exclude = do
  files <- FP.getDirectoryFiles dir ["*.analysis.json"]
  let getName file = FP.match "*.analysis.json" file >>= listToMaybe
      names =
        -- exclude duplicats
        filter (`L.notElem` exclude) $ catMaybes $ getName <$> files
  -- print names
  items <- mapM (loadItem dir) names
  pure $ catMaybes items

loadExamples :: IO (Trace PVParams, Trace PVParams)
loadExamples = do
  Just (_, _, bwv940Trace, _) <- loadItem
    (dataDir </> "bach" </> "fünf-kleine-präludien")
    "BWV_0940"
  Just (_, _, rareTrace, _) <- loadItem (dataDir </> "theory-article")
                                        "10c_rare_int"
  pure (bwv940Trace, rareTrace)

-- learning
-- --------

learn :: Hyper PVParams -> [Piece] -> IO (Hyper PVParams)
learn = foldM train
 where
  train prior (name, _, trace, _) =
    case getPosterior prior trace sampleDerivation' of
      Nothing -> do
        putStrLn $ "couldn't compute posterior for " <> name <> ", skipping."
        pure prior
      Just post -> do
        -- putStrLn $ "learned from " <> name <> "."
        pure post

-- evaluating
-- ----------

leaveOneOut :: [a] -> [(a, [a])]
leaveOneOut dataset = go dataset [] []
 where
  go []       _    splits = splits
  go (x : xs) done splits = go xs (x : done) ((x, xs <> done) : splits)

derivationLogProb
  :: MWC.GenIO -> Hyper PVParams -> Trace PVParams -> IO [Double]
derivationLogProb gen hyper trace = do
  probs <- replicateM nSamples $ MWC.sample (sampleProbs @PVParams hyper) gen
  let estimates = mapMaybe
        (\params -> snd <$> evalTraceLogP params trace sampleDerivation')
        probs
  pure $! estimates

derivationLogProb'
  :: MWC.GenIO -> Hyper PVParams -> Trace PVParams -> IO Double
derivationLogProb' gen hyper trace = do
  estimates <- derivationLogProb gen hyper trace
  pure $! meanOfLogs estimates

meanOfLogs :: (RealFloat b, Foldable t, Functor t) => t b -> b
meanOfLogs logs =
  Log.ln $! Log.sum (Log.Exp <$> logs) / fromIntegral (length logs)

countNotes :: Path [a] b -> Int
countNotes (PathEnd notes       ) = length notes
countNotes (Path notes edges rst) = length notes + countNotes rst

sampleBaselines
  :: (StatefulGen g IO)
  => g
  -> MWC.GenIO
  -> Hyper PVParams
  -> Piece
  -> IO [Double]
sampleBaselines gen genMWC posterior (name, _, _, surface) = do
  derivsTry <- replicateM nBaselines $ runExceptT $ Greedy.parseRandom'
    gen
    protoVoiceEvaluator
    surface
  let derivs = rights derivsTry
  putStrLn $ "  collected " <> show (length derivs) <> " samples for " <> name
  let baselines = flip fmap derivs $ \ana -> do
        trace <- observeDerivation' (anaDerivation ana)
        pure $ derivationLogProb' genMWC posterior trace
  baselines <- traverse sequence baselines -- maaaaagic
  fmap catMaybes $ forM baselines $ \case
    Left  err -> putStrLn err >> pure Nothing
    Right val -> pure $ Just val

summarize :: [Double] -> (Double, Double)
summarize xs = (Stats.mean sample, Stats.stdDev sample)
  where sample = V.fromList xs

comparePerNote
  :: (StatefulGen g IO)
  => g
  -> MWC.GenIO
  -> Hyper PVParams
  -> (Piece, [Piece])
  -> IO (Double, Double, Double, Double, [Double])
comparePerNote gen genMWC prior (test@(tstName, _, tstTrace, tstSurface), train)
  = do
    putStrLn $ "testing on " <> tstName
    posterior <- learn prior train
    -- compute normalized logprobs
    let nnotes = int2Double $ countNotes tstSurface
    logpPrior <- derivationLogProb' genMWC prior tstTrace
    logpPost  <- derivationLogProb' genMWC posterior tstTrace
    let logppnPrior = logpPrior / nnotes
        logppnPost  = logpPost / nnotes
    baselines <- sampleBaselines gen genMWC posterior test
    let (basemean, basestd) = summarize baselines
        basemeanpn          = basemean / nnotes
        basestdpn           = basestd / nnotes
    putStrLn $ "  nnotes: " <> show nnotes
    putStrLn $ "  logppn prior: " <> show logppnPrior <> " nats"
    putStrLn $ "  logppn posterior: " <> show logppnPost <> " nats"
    putStrLn
      $  "  logppn baseline: mean="
      <> show basemeanpn
      <> ", std="
      <> show basestdpn
    putStrLn $ "  baselines top 5: " <> show
      (drop 95 $ (/ nnotes) <$> L.sort baselines)
    putStrLn $ "  perplexity prior: " <> show (exp $ negate logppnPrior)
    putStrLn $ "  perplexity posterior: " <> show (exp $ negate logppnPost)
    putStrLn $ "  Δlogppn: " <> show (logppnPost - logppnPrior) <> " nats"
    pure (logpPrior, logpPost, nnotes, basemean, baselines)

-- plotting
-- --------

a % b = a Plt.% Plt.mp Plt.# b
infixl 5 %

showSplits
  :: [Double] -> [[Double]] -> [Double] -> [String] -> IO (Either String String)
showSplits logppns blogppns counts testpieces =
  Plt.file "splits.svg"
    $ Plt.readData (logppns, blogppns, counts, testpieces)
    % "import numpy as np"
    % "import pandas as pd"
    % "import seaborn as sns"
    % "(logppns, blogppns, counts, pieces) = tuple(map(np.array, data))"
    % "baselines = pd.concat([pd.DataFrame({'logppn': bppns, 'piece': piece}) for bppns, piece in zip(blogppns, pieces)])"
    % "testscores = pd.DataFrame({'logppn': logppns, 'piece': pieces})"
    % "colors = sns.color_palette()"
    % "fig, ax = plot.subplots(figsize=(10,5))"
    % "sns.boxplot(ax=ax, y='piece', x='logppn', data=baselines, whis=(2,98), color=colors[0], fliersize=4)"
    % "sns.scatterplot(ax=ax, y='piece', x='logppn', data=testscores, color=colors[1])"
    % "ax.invert_yaxis()"
    % "fig.tight_layout()"
    % "fig.savefig('splits.pdf')"
    % "fig.savefig('splits.png')"


{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}

module HackageBenchmark (
    hackageBenchmarkMain

-- Exposed for testing:
  , CabalResult(..)
  , isSignificantTimeDifference
  , combineTrialResults
  , isSignificantResult
  , shouldContinueAfterFirstTrial
  ) where

import Control.Monad (forM_, replicateM, unless, when)
import qualified Data.ByteString as B
import Data.List (nub, unzip4)
import Data.Maybe (isJust)
import Data.Monoid ((<>))
import Data.String (fromString)
import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime)
import qualified Data.Vector.Unboxed as V
import Options.Applicative
import Statistics.Sample (mean, stdDev)
import Statistics.Test.MannWhitneyU ( PositionTest(..), TestResult(..)
                                    , mannWhitneyUCriticalValue
                                    , mannWhitneyUtest)
import Statistics.Types (PValue, mkPValue)
import System.Directory (getTemporaryDirectory)
import System.Exit (ExitCode(..), exitFailure)
import System.FilePath ((</>))
import System.IO ( BufferMode(LineBuffering), hPutStrLn, hSetBuffering, stderr
                 , stdout)
import System.Process ( StdStream(CreatePipe), CreateProcess(..), callProcess
                      , createProcess, readProcess, shell, waitForProcess )
import Text.Printf (printf)

import Distribution.Package (PackageName, mkPackageName, unPackageName)

data Args = Args {
    argCabal1                      :: FilePath
  , argCabal2                      :: FilePath
  , argCabal1Flags                 :: [String]
  , argCabal2Flags                 :: [String]
  , argPackages                    :: [PackageName]
  , argMinRunTimeDifferenceToRerun :: Double
  , argPValue                      :: PValue Double
  , argTrials                      :: Int
  , argPrintTrials                 :: Bool
  , argPrintSkippedPackages        :: Bool
  , argTimeoutSeconds              :: Int
  }

data CabalTrial = CabalTrial NominalDiffTime CabalResult

data CabalResult
  = Solution
  | NoInstallPlan
  | BackjumpLimit
  | Unbuildable
  | UnbuildableDep
  | ComponentCycle
  | ModReexpIssue
  | PkgNotFound
  | Timeout
  | Unknown
  deriving (Eq, Show)

hackageBenchmarkMain :: IO ()
hackageBenchmarkMain = do
  hSetBuffering stdout LineBuffering
  args@Args {..} <- execParser parserInfo
  checkArgs args
  printConfig args
  pkgs <- getPackages args
  putStrLn ""

  let -- The maximum length of the heading and package names.
      nameColumnWidth :: Int
      nameColumnWidth =
          maximum $ map length $ "package" : map unPackageName pkgs
      runCabal1 = runCabal argTimeoutSeconds argCabal1 argCabal1Flags
      runCabal2 = runCabal argTimeoutSeconds argCabal2 argCabal2Flags

  -- When the output contains both trails and summaries, label each row as
  -- "trial" or "summary".
  when argPrintTrials $ putStr $ printf "%-16s " "trial/summary"
  putStrLn $
      printf "%-*s %-14s %-14s %11s %11s %11s %11s %11s"
             nameColumnWidth "package" "result1" "result2"
             "mean1" "mean2" "stddev1" "stddev2" "speedup"

  forM_ pkgs $ \pkg -> do
    let printTrial msgType result1 result2 time1 time2 =
            putStrLn $
            printf "%-16s %-*s %-14s %-14s %10.3fs %10.3fs"
                   msgType nameColumnWidth (unPackageName pkg)
                   (show result1) (show result2)
                   (diffTimeToDouble time1) (diffTimeToDouble time2)

    CabalTrial t1 r1 <- runCabal1 pkg
    CabalTrial t2 r2 <- runCabal2 pkg
    if not $
       shouldContinueAfterFirstTrial argMinRunTimeDifferenceToRerun t1 t2 r1 r2
    then when argPrintSkippedPackages $
         if argPrintTrials
         then printTrial "trial (skipping)" r1 r2 t1 t2
         else putStrLn $ printf "%-*s (first run times were too similar)"
                                nameColumnWidth (unPackageName pkg)
    else do
      when argPrintTrials $ printTrial "trial" r1 r2 t1 t2
      (ts1, ts2, rs1, rs2) <- (unzip4 . ((t1, t2, r1, r2) :) <$>)
                            . replicateM (argTrials - 1) $ do
        CabalTrial t1' r1' <- runCabal1 pkg
        CabalTrial t2' r2' <- runCabal2 pkg
        when argPrintTrials $ printTrial "trial" r1' r2' t1' t2'
        return (t1', t2', r1', r2')

      let result1 = combineTrialResults rs1
          result2 = combineTrialResults rs2
          times1 = V.fromList (map diffTimeToDouble ts1)
          times2 = V.fromList (map diffTimeToDouble ts2)
          mean1 = mean times1
          mean2 = mean times2
          stddev1 = stdDev times1
          stddev2 = stdDev times2
          speedup = mean1 / mean2

      when argPrintTrials $ putStr $ printf "%-16s " "summary"
      if isSignificantResult result1 result2
          || isSignificantTimeDifference argPValue ts1 ts2
      then putStrLn $
           printf "%-*s %-14s %-14s %10.3fs %10.3fs %10.3fs %10.3fs %10.3f"
                  nameColumnWidth (unPackageName pkg)
                  (show result1) (show result2) mean1 mean2 stddev1 stddev2 speedup
      else when (argPrintTrials || argPrintSkippedPackages) $
           putStrLn $
           printf "%-*s (not significant)" nameColumnWidth (unPackageName pkg)
  where
    checkArgs :: Args -> IO ()
    checkArgs Args {..} = do
      let die msg = hPutStrLn stderr msg >> exitFailure
      unless (argTrials > 0) $ die "--trials must be greater than 0."
      unless (argMinRunTimeDifferenceToRerun >= 0) $
          die "--min-run-time-percentage-difference-to-rerun must be non-negative."
      unless (isSampleLargeEnough argPValue argTrials) $
          die "p-value is too small for the number of trials."

    printConfig :: Args -> IO ()
    printConfig Args {..} = do
      putStrLn "Comparing:"
      putStrLn $ "1: " ++ argCabal1 ++ " " ++ unwords argCabal1Flags
      callProcess argCabal1 ["--version"]
      putStrLn $ "2: " ++ argCabal2 ++ " " ++ unwords argCabal2Flags
      callProcess argCabal2 ["--version"]
      -- TODO: Print index state.
      putStrLn "Base package database:"
      callProcess "ghc-pkg" ["list"]

    getPackages :: Args -> IO [PackageName]
    getPackages Args {..} = do
      pkgs <-
          if null argPackages
          then do
            putStrLn $ "Obtaining the package list (using " ++ argCabal1 ++ ") ..."
            list <- readProcess argCabal1 ["list", "--simple-output"] ""
            return $ nub [mkPackageName $ head (words line) | line <- lines list]
          else do
            putStrLn "Using given package list ..."
            return argPackages
      putStrLn $ "Done, got " ++ show (length pkgs) ++ " packages."
      return pkgs

runCabal :: Int -> FilePath -> [String] -> PackageName -> IO CabalTrial
runCabal timeoutSeconds cabal flags pkg = do
  ((exitCode, err), time) <- timeEvent $ do
    tmpDir <- getTemporaryDirectory

    let timeout = "timeout --foreground -sINT " ++ show timeoutSeconds
        cabalCmd = unwords $
            [ cabal

              -- A non-existent store directory prevents cabal from reading the
              -- store, which would cause the size of the store to affect run
              -- time.
            , "--store-dir=" ++ (tmpDir </> "non-existent-store-dir")

            , "v2-install"

              -- These flags prevent a Cabal project or package environment from
              -- affecting the install plan.
            , "--ignore-project"
            , "--package-env=non-existent-package-env"

              -- --lib allows solving for packages with libraries or
              -- executables.
            , "--lib"

            , unPackageName pkg

            , "--dry-run"

              -- The test doesn't currently handle stdout, so we suppress it
              -- with silent. nowrap simplifies parsing the errors messages.
            , "-vsilent+nowrap"]

             ++ flags
        cmd = (shell (timeout ++ " " ++ cabalCmd)) { std_err = CreatePipe }

    -- TODO: Read stdout and compare the install plans.
    (_, _, Just errh, ph) <- createProcess cmd
    err <- B.hGetContents errh
    (, err) <$> waitForProcess ph
  let exhaustiveMsg =
          "After searching the rest of the dependency tree exhaustively"
      result
        | exitCode == ExitSuccess                                                          = Solution
        | exitCode == ExitFailure 124                                                      = Timeout
        | fromString exhaustiveMsg `B.isInfixOf` err                                       = NoInstallPlan
        | fromString "Backjump limit reached" `B.isInfixOf` err                            = BackjumpLimit
        | fromString "none of the components are available to build" `B.isInfixOf` err     = Unbuildable
        | fromString "Dependency on unbuildable" `B.isInfixOf` err                         = UnbuildableDep
        | fromString "Dependency cycle between the following components" `B.isInfixOf` err = ComponentCycle
        | fromString "Problem with module re-exports" `B.isInfixOf` err                    = ModReexpIssue
        | fromString "There is no package named" `B.isInfixOf` err                         = PkgNotFound
        | otherwise                                                                        = Unknown
  return (CabalTrial time result)

isSampleLargeEnough :: PValue Double -> Int -> Bool
isSampleLargeEnough pvalue trials =
    -- mannWhitneyUCriticalValue, which can fail with too few samples, is only
    -- used when both sample sizes are less than or equal to 20.
    trials > 20 || isJust (mannWhitneyUCriticalValue (trials, trials) pvalue)

isSignificantTimeDifference :: PValue Double -> [NominalDiffTime] -> [NominalDiffTime] -> Bool
isSignificantTimeDifference pvalue xs ys =
  let toVector = V.fromList . map diffTimeToDouble
  in case mannWhitneyUtest SamplesDiffer pvalue (toVector xs) (toVector ys) of
       Nothing             -> error "not enough data for mannWhitneyUtest"
       Just Significant    -> True
       Just NotSignificant -> False

-- Should we stop after the first trial of this package to save time? This
-- function skips the package if the results are uninteresting and the times are
-- within --min-run-time-percentage-difference-to-rerun.
shouldContinueAfterFirstTrial :: Double
                              -> NominalDiffTime
                              -> NominalDiffTime
                              -> CabalResult
                              -> CabalResult
                              -> Bool
shouldContinueAfterFirstTrial 0                            _  _  _       _       = True
shouldContinueAfterFirstTrial _                            _  _  Timeout Timeout = False
shouldContinueAfterFirstTrial maxRunTimeDifferenceToIgnore t1 t2 r1      r2      =
    isSignificantResult r1 r2
 || abs (t1 - t2) / min t1 t2 >= realToFrac (maxRunTimeDifferenceToIgnore / 100)

isSignificantResult :: CabalResult -> CabalResult -> Bool
isSignificantResult r1 r2 = r1 /= r2 || not (isExpectedResult r1)

-- Is this result expected in a benchmark run on all of Hackage?
isExpectedResult :: CabalResult -> Bool
isExpectedResult Solution       = True
isExpectedResult NoInstallPlan  = True
isExpectedResult BackjumpLimit  = True
isExpectedResult Timeout        = True
isExpectedResult Unbuildable    = True
isExpectedResult UnbuildableDep = True
isExpectedResult ComponentCycle = True
isExpectedResult ModReexpIssue  = True
isExpectedResult PkgNotFound    = False
isExpectedResult Unknown        = False

-- Combine CabalResults from multiple trials. Ignoring timeouts, all results
-- should be the same. If they aren't the same, we returns Unknown.
combineTrialResults :: [CabalResult] -> CabalResult
combineTrialResults rs
  | allEqual rs                          = head rs
  | allEqual [r | r <- rs, r /= Timeout] = Timeout
  | otherwise                            = Unknown
  where
    allEqual :: Eq a => [a] -> Bool
    allEqual xs = length (nub xs) == 1

timeEvent :: IO a -> IO (a, NominalDiffTime)
timeEvent task = do
  start <- getCurrentTime
  r <- task
  end <- getCurrentTime
  return (r, diffUTCTime end start)

diffTimeToDouble :: NominalDiffTime -> Double
diffTimeToDouble = fromRational . toRational

parserInfo :: ParserInfo Args
parserInfo = info (argParser <**> helper)
     ( fullDesc
    <> progDesc ("Find differences between two cabal commands when solving"
                   ++ " for all packages on Hackage.")
    <> header "hackage-benchmark" )

argParser :: Parser Args
argParser = Args
    <$> strOption
         ( long "cabal1"
        <> metavar "PATH"
        <> help "First cabal executable")
    <*> strOption
         ( long "cabal2"
        <> metavar "PATH"
        <> help "Second cabal executable")
    <*> option (words <$> str)
         ( long "cabal1-flags"
        <> value []
        <> metavar "FLAGS"
        <> help "Extra flags for the first cabal executable")
    <*> option (words <$> str)
         ( long "cabal2-flags"
        <> value []
        <> metavar "FLAGS"
        <> help "Extra flags for the second cabal executable")
    <*> option (map mkPackageName . words <$> str)
         ( long "packages"
        <> value []
        <> metavar "PACKAGES"
        <> help ("Space separated list of packages to test, or all of Hackage"
                   ++ " if unspecified"))
    <*> option auto
         ( long "min-run-time-percentage-difference-to-rerun"
        <> showDefault
        <> value 0.0
        <> metavar "PERCENTAGE"
        <> help ("Stop testing a package when the difference in run times in"
                   ++ " the first trial are within this percentage, in order to"
                   ++ " save time"))
    <*> option (mkPValue <$> auto)
         ( long "pvalue"
        <> showDefault
        <> value (mkPValue 0.05)
        <> metavar "DOUBLE"
        <> help ("p-value used to determine whether to print the results for"
                   ++ " each package"))
    <*> option auto
         ( long "trials"
        <> showDefault
        <> value 10
        <> metavar "N"
        <> help "Number of trials for each package")
    <*> switch
         ( long "print-trials"
        <> help "Whether to include the results from individual trials in the output")
    <*> switch
         ( long "print-skipped-packages"
        <> help "Whether to include skipped packages in the output")
    <*> option auto
         ( long "timeout"
        <> showDefault
        <> value 90
        <> metavar "SECONDS"
        <> help "Maximum time to run a cabal command, in seconds")


{-# language GADTs                 #-}
{-# language StandaloneDeriving    #-}
{-# language TypeOperators         #-}
{-# language FlexibleInstances     #-}
{-# language FlexibleContexts      #-}
{-# language UndecidableInstances  #-}
{-# language MultiParamTypeClasses #-}
{-# language TypeFamilies          #-}

{-# options_ghc -fwarn-incomplete-patterns #-}

module Feldspar.Software.Primitive where

import Feldspar.Representation
import Data.Struct

import Data.Array ((!))
import Data.Bits (Bits)
import Data.Complex
import Data.Int
import Data.Word
import Data.List (genericTake)
import Data.Typeable hiding (TypeRep)
import Data.Constraint hiding (Sub)
import qualified Data.Bits as Bits

-- syntactic.
import Language.Syntactic
import Language.Syntactic.Functional
import Language.Syntactic.Functional.Tuple
import qualified Language.Syntactic as Syn

-- imperative-edsl.
import Language.Embedded.Expression
import qualified Language.Embedded.Imperative.CMD as Imp (IArr(..))

--------------------------------------------------------------------------------
-- * Software primitives.
--------------------------------------------------------------------------------

--------------------------------------------------------------------------------
-- ** Software primitive types.

-- | Representation of supported, primitive software types.
data SoftwarePrimTypeRep a
  where
    -- booleans
    BoolST   :: SoftwarePrimTypeRep Bool
    -- signed numbers.
    Int8ST   :: SoftwarePrimTypeRep Int8
    Int16ST  :: SoftwarePrimTypeRep Int16
    Int32ST  :: SoftwarePrimTypeRep Int32
    Int64ST  :: SoftwarePrimTypeRep Int64
    -- unsigned numbers.
    Word8ST  :: SoftwarePrimTypeRep Word8
    Word16ST :: SoftwarePrimTypeRep Word16
    Word32ST :: SoftwarePrimTypeRep Word32
    Word64ST :: SoftwarePrimTypeRep Word64
    -- floating point numbers.
    FloatST  :: SoftwarePrimTypeRep Float
    DoubleST :: SoftwarePrimTypeRep Double
    -- complex numbers.
    ComplexFloatST  :: SoftwarePrimTypeRep (Complex Float)
    ComplexDoubleST :: SoftwarePrimTypeRep (Complex Double)

deriving instance Eq       (SoftwarePrimTypeRep a)
deriving instance Show     (SoftwarePrimTypeRep a)
deriving instance Typeable (SoftwarePrimTypeRep a)

instance Inhabited (Complex Float)
  where
    reset = 0 :+ 0

instance Inhabited (Complex Double)
  where
    reset = 0 :+ 0

--------------------------------------------------------------------------------

-- | Class of supported, primitive software types.
class (Eq a, Show a, Typeable a, Inhabited a) => SoftwarePrimType a
  where
    softwareRep :: SoftwarePrimTypeRep a

instance SoftwarePrimType Bool   where softwareRep = BoolST
instance SoftwarePrimType Int8   where softwareRep = Int8ST
instance SoftwarePrimType Int16  where softwareRep = Int16ST
instance SoftwarePrimType Int32  where softwareRep = Int32ST
instance SoftwarePrimType Int64  where softwareRep = Int64ST
instance SoftwarePrimType Word8  where softwareRep = Word8ST
instance SoftwarePrimType Word16 where softwareRep = Word16ST
instance SoftwarePrimType Word32 where softwareRep = Word32ST
instance SoftwarePrimType Word64 where softwareRep = Word64ST
instance SoftwarePrimType Float  where softwareRep = FloatST
instance SoftwarePrimType Double where softwareRep = DoubleST
instance SoftwarePrimType (Complex Float)  where softwareRep = ComplexFloatST
instance SoftwarePrimType (Complex Double) where softwareRep = ComplexDoubleST

-- | Compare two primitive software types for equality.
softwarePrimTypeEq :: SoftwarePrimTypeRep a -> SoftwarePrimTypeRep b -> Maybe (Dict (a ~ b))
softwarePrimTypeEq (BoolST)   (BoolST)   = Just Dict
softwarePrimTypeEq (Int8ST)   (Int8ST)   = Just Dict
softwarePrimTypeEq (Int16ST)  (Int16ST)  = Just Dict
softwarePrimTypeEq (Int32ST)  (Int32ST)  = Just Dict
softwarePrimTypeEq (Int64ST)  (Int64ST)  = Just Dict
softwarePrimTypeEq (Word8ST)  (Word8ST)  = Just Dict
softwarePrimTypeEq (Word16ST) (Word16ST) = Just Dict
softwarePrimTypeEq (Word32ST) (Word32ST) = Just Dict
softwarePrimTypeEq (Word64ST) (Word64ST) = Just Dict
softwarePrimTypeEq (FloatST)  (FloatST)  = Just Dict
softwarePrimTypeEq (DoubleST) (DoubleST) = Just Dict
softwarePrimTypeEq (ComplexFloatST)  (ComplexFloatST)  = Just Dict
softwarePrimTypeEq (ComplexDoubleST) (ComplexDoubleST) = Just Dict
softwarePrimTypeEq _          _          = Nothing

-- | Construct the primitive software type representation of 'a'.
softwarePrimTypeOf :: SoftwarePrimType a => a -> SoftwarePrimTypeRep a
softwarePrimTypeOf _ = softwareRep

-- | Construct a primitive software type witness from its representation.
softwarePrimWitType :: SoftwarePrimTypeRep a -> Dict (SoftwarePrimType a)
softwarePrimWitType BoolST   = Dict
softwarePrimWitType Int8ST   = Dict
softwarePrimWitType Int16ST  = Dict
softwarePrimWitType Int32ST  = Dict
softwarePrimWitType Int64ST  = Dict
softwarePrimWitType Word8ST  = Dict
softwarePrimWitType Word16ST = Dict
softwarePrimWitType Word32ST = Dict
softwarePrimWitType Word64ST = Dict
softwarePrimWitType FloatST  = Dict
softwarePrimWitType DoubleST = Dict
softwarePrimWitType ComplexFloatST  = Dict
softwarePrimWitType ComplexDoubleST = Dict

--------------------------------------------------------------------------------
-- ** Software primitive expressions.

-- | Software primitive symbols.
data SoftwarePrim sig
  where
    -- free variables and literals.
    FreeVar :: (SoftwarePrimType a) => String -> SoftwarePrim (Full a)
    Lit     :: (Show a, Eq a)       => a      -> SoftwarePrim (Full a)
    -- numerical operations.
    Neg  :: (SoftwarePrimType a, Num a) => SoftwarePrim (a :-> Full a)
    Abs  :: (SoftwarePrimType a, Num a) => SoftwarePrim (a :-> Full a)
    Sign :: (SoftwarePrimType a, Num a) => SoftwarePrim (a :-> Full a)
    Add  :: (SoftwarePrimType a, Num a) => SoftwarePrim (a :-> a :-> Full a)
    Sub  :: (SoftwarePrimType a, Num a) => SoftwarePrim (a :-> a :-> Full a)
    Mul  :: (SoftwarePrimType a, Num a) => SoftwarePrim (a :-> a :-> Full a)
    -- integral operations.
    Div  :: (SoftwarePrimType a, Integral a) => SoftwarePrim (a :-> a :-> Full a)
    Mod  :: (SoftwarePrimType a, Integral a) => SoftwarePrim (a :-> a :-> Full a)
    Quot :: (SoftwarePrimType a, Integral a) => SoftwarePrim (a :-> a :-> Full a)
    Rem  :: (SoftwarePrimType a, Integral a) => SoftwarePrim (a :-> a :-> Full a)
    --
    FDiv :: (SoftwarePrimType a, Fractional a) => SoftwarePrim (a :-> a :-> Full a)
    -- floating point operators.
    Pi    :: (SoftwarePrimType a, Floating a) => SoftwarePrim (Full a)
    Exp   :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Log   :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Sqrt  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Pow   :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> a :-> Full a)
    Sin   :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Cos   :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Tan   :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Asin  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Acos  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Atan  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Sinh  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Cosh  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Tanh  :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Asinh :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Acosh :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    Atanh :: (SoftwarePrimType a, Floating a) => SoftwarePrim (a :-> Full a)
    -- complex operators.
    Complex   :: (SoftwarePrimType a, SoftwarePrimType (Complex a), Num a) =>
      SoftwarePrim (a :-> a :-> Full (Complex a))
    Real      :: (SoftwarePrimType a, SoftwarePrimType (Complex a)) =>
      SoftwarePrim (Complex a :-> Full a)
    Imag      :: (SoftwarePrimType a, SoftwarePrimType (Complex a)) =>
      SoftwarePrim (Complex a :-> Full a)
    Polar     :: (SoftwarePrimType a, SoftwarePrimType (Complex a), Floating a) =>
      SoftwarePrim (a :-> a :-> Full (Complex a))
    Magnitude :: (SoftwarePrimType a, SoftwarePrimType (Complex a), RealFloat a) =>
      SoftwarePrim (Complex a :-> Full a)
    Phase     :: (SoftwarePrimType a, SoftwarePrimType (Complex a), RealFloat a) =>
      SoftwarePrim (Complex a :-> Full a)
    Conjugate :: (SoftwarePrimType a, SoftwarePrimType (Complex a), Num a) =>
      SoftwarePrim (Complex a :-> Full (Complex a))
    -- type casting.
    I2N   :: (SoftwarePrimType a, Integral a, SoftwarePrimType b, Num b) =>
      SoftwarePrim (a :-> Full b)
    I2B   :: (SoftwarePrimType a, Integral a) =>
      SoftwarePrim (a :-> Full Bool)
    B2I   :: (SoftwarePrimType a, Integral a) =>
      SoftwarePrim (Bool :-> Full a)
    Round :: (SoftwarePrimType a, RealFrac a, SoftwarePrimType b, Num b) =>
      SoftwarePrim (a :-> Full b)
    -- logical operations.
    Not     :: SoftwarePrim (Bool :-> Full Bool)
    And     :: SoftwarePrim (Bool :-> Bool :-> Full Bool)
    Or      :: SoftwarePrim (Bool :-> Bool :-> Full Bool)
    -- bitwise logical operations.
    BitAnd   :: (SoftwarePrimType a, Bits a) => SoftwarePrim (a :-> a :-> Full a)
    BitOr    :: (SoftwarePrimType a, Bits a) => SoftwarePrim (a :-> a :-> Full a)
    BitXor   :: (SoftwarePrimType a, Bits a) => SoftwarePrim (a :-> a :-> Full a)
    BitCompl :: (SoftwarePrimType a, Bits a) => SoftwarePrim (a :-> Full a)
    ShiftL   :: (SoftwarePrimType a, Bits a, SoftwarePrimType b, Integral b) =>
      SoftwarePrim (a :-> b :-> Full a)
    ShiftR   :: (SoftwarePrimType a, Bits a, SoftwarePrimType b, Integral b) =>
      SoftwarePrim (a :-> b :-> Full a)
    RotateL  :: (SoftwarePrimType a, Bits a, SoftwarePrimType b, Integral b) =>
      SoftwarePrim (a :-> b :-> Full a)
    RotateR  :: (SoftwarePrimType a, Bits a, SoftwarePrimType b, Integral b) =>
      SoftwarePrim (a :-> b :-> Full a)
    -- relational operations.
    Eq  :: (SoftwarePrimType a, Eq a)  => SoftwarePrim (a :-> a :-> Full Bool)
    Neq :: (SoftwarePrimType a, Eq a)  => SoftwarePrim (a :-> a :-> Full Bool)
    Lt  :: (SoftwarePrimType a, Ord a) => SoftwarePrim (a :-> a :-> Full Bool)
    Lte :: (SoftwarePrimType a, Ord a) => SoftwarePrim (a :-> a :-> Full Bool)
    Gt  :: (SoftwarePrimType a, Ord a) => SoftwarePrim (a :-> a :-> Full Bool)
    Gte :: (SoftwarePrimType a, Ord a) => SoftwarePrim (a :-> a :-> Full Bool)
    -- conditional.
    Cond :: SoftwarePrim (Bool :-> a :-> a :-> Full a)
    -- array indexing.
    ArrIx :: (SoftwarePrimType a) => Imp.IArr Index a ->
      SoftwarePrim (Index :-> Full a)

deriving instance Show     (SoftwarePrim a)
deriving instance Typeable (SoftwarePrim a)

--------------------------------------------------------------------------------

-- | Software primitive symbols.
type SoftwarePrimConstructs = SoftwarePrim

-- | Software primitive symbols tagged with their type representation.
type SoftwarePrimDomain = SoftwarePrimConstructs :&: SoftwarePrimTypeRep

-- | Software primitive expressions.
newtype Prim a = Prim { unPrim :: ASTF SoftwarePrimDomain a }

-- | Evaluate a closed, software primitive expression.
evalPrim :: Prim a -> a
evalPrim = go . unPrim
  where
    go :: AST SoftwarePrimDomain sig -> Denotation sig
    go (Sym (s :&: _)) = evalSym s
    go (f :$ a)        = go f $ go a

-- | Sugar a software primitive symbol as a smart constructor.
sugarSymPrim
  :: ( Signature sig
     , fi  ~ SmartFun dom sig
     , sig ~ SmartSig fi
     , dom ~ SmartSym fi
     , dom ~ SoftwarePrimDomain
     , SyntacticN f fi
     , sub :<: SoftwarePrimConstructs
     , SoftwarePrimType (DenResult sig)
     )
  => sub sig -> f
sugarSymPrim = sugarSymDecor softwareRep

--------------------------------------------------------------------------------

instance Syntactic (Prim a)
  where
    type Domain   (Prim a) = SoftwarePrimDomain
    type Internal (Prim a) = a
    desugar = unPrim
    sugar   = Prim

instance FreeExp Prim
  where
    type FreePred Prim = SoftwarePrimType
    constExp = sugarSymPrim . Lit
    varExp   = sugarSymPrim . FreeVar

instance EvalExp Prim
  where
    evalExp = evalPrim

--------------------------------------------------------------------------------
-- front-end.

instance (SoftwarePrimType a, Num a) => Num (Prim a)
  where
    fromInteger = constExp . fromInteger
    (+)         = sugarSymPrim Add
    (-)         = sugarSymPrim Sub
    (*)         = sugarSymPrim Mul
    negate      = sugarSymPrim Neg
    abs         = error "Num (Prim a): abs."
    signum      = error "Num (Prim a): signum."

--------------------------------------------------------------------------------
-- syntactic instances.

instance Eval SoftwarePrim
  where
    evalSym (FreeVar v) = error $ "evaluating free variable " ++ show v
    evalSym (Lit a)     = a
    evalSym Cond        = \c t f -> if c then t else f
    evalSym Neg         = negate
    evalSym Abs         = abs
    evalSym Sign        = signum
    evalSym Add         = (+)
    evalSym Sub         = (-)
    evalSym Mul         = (*)
    evalSym Div         = div
    evalSym Mod         = mod
    evalSym Quot        = quot
    evalSym Rem         = rem
    evalSym FDiv        = (/)
    evalSym Pi          = pi
    evalSym Exp         = exp
    evalSym Log         = log
    evalSym Sqrt        = sqrt
    evalSym Pow         = (**)
    evalSym Sin         = sin
    evalSym Cos         = cos
    evalSym Tan         = tan
    evalSym Asin        = asin
    evalSym Acos        = acos
    evalSym Atan        = atan
    evalSym Sinh        = sinh
    evalSym Cosh        = cosh
    evalSym Tanh        = tanh
    evalSym Asinh       = asinh
    evalSym Acosh       = acosh
    evalSym Atanh       = atanh
    evalSym Complex     = (:+)
    evalSym Polar       = mkPolar
    evalSym Real        = realPart
    evalSym Imag        = imagPart
    evalSym Magnitude   = magnitude
    evalSym Phase       = phase
    evalSym Conjugate   = conjugate
    evalSym I2N         = fromIntegral
    evalSym I2B         = (/=0)
    evalSym B2I         = \a -> if a then 1 else 0
    evalSym Round       = fromInteger . round
    evalSym Not         = not
    evalSym And         = (&&)
    evalSym Or          = (||)
    evalSym BitAnd      = (Bits..&.)
    evalSym BitOr       = (Bits..|.)
    evalSym BitXor      = Bits.xor
    evalSym BitCompl    = Bits.complement
    evalSym ShiftL      = \b i -> Bits.shiftL  b (fromIntegral i)
    evalSym ShiftR      = \b i -> Bits.shiftR  b (fromIntegral i)
    evalSym RotateL     = \b i -> Bits.rotateL b (fromIntegral i)
    evalSym RotateR     = \b i -> Bits.rotateR b (fromIntegral i)
    evalSym Eq          = (==)
    evalSym Neq         = (/=)
    evalSym Lt          = (<)
    evalSym Lte         = (<=)
    evalSym Gt          = (>)
    evalSym Gte         = (>=)
    evalSym (ArrIx (Imp.IArrRun arr)) = \i -> arr ! i
    evalSym (ArrIx _)   = error "eval of variable."

instance Symbol SoftwarePrim
  where
    symSig (FreeVar v) = signature
    symSig (Lit a)     = signature
    symSig Cond        = signature
    symSig Neg         = signature
    symSig Abs         = signature
    symSig Sign        = signature
    symSig Add         = signature
    symSig Sub         = signature
    symSig Mul         = signature
    symSig Div         = signature
    symSig Mod         = signature
    symSig Quot        = signature
    symSig Rem         = signature
    symSig FDiv        = signature
    symSig Pi          = signature
    symSig Exp         = signature
    symSig Log         = signature
    symSig Sqrt        = signature
    symSig Pow         = signature
    symSig Sin         = signature
    symSig Cos         = signature
    symSig Tan         = signature
    symSig Asin        = signature
    symSig Acos        = signature
    symSig Atan        = signature
    symSig Sinh        = signature
    symSig Cosh        = signature
    symSig Tanh        = signature
    symSig Asinh       = signature
    symSig Acosh       = signature
    symSig Atanh       = signature
    symSig Complex     = signature
    symSig Real        = signature
    symSig Imag        = signature
    symSig Polar       = signature
    symSig Magnitude   = signature
    symSig Phase       = signature
    symSig Conjugate   = signature
    symSig I2N         = signature
    symSig I2B         = signature
    symSig B2I         = signature
    symSig Round       = signature
    symSig Not         = signature
    symSig And         = signature
    symSig Or          = signature
    symSig BitAnd      = signature
    symSig BitOr       = signature
    symSig BitXor      = signature
    symSig BitCompl    = signature
    symSig ShiftL      = signature
    symSig ShiftR      = signature
    symSig RotateL     = signature
    symSig RotateR     = signature
    symSig Eq          = signature
    symSig Neq         = signature
    symSig Lt          = signature
    symSig Lte         = signature
    symSig Gt          = signature
    symSig Gte         = signature
    symSig (ArrIx a)   = signature

instance Render SoftwarePrim
  where
    renderSym  = show
    renderArgs = renderArgsSmart

instance Equality SoftwarePrim
  where
    equal s1 s2 = show s1 == show s2

instance StringTree SoftwarePrim
instance EvalEnv SoftwarePrim env

--------------------------------------------------------------------------------


{-# LANGUAGE BangPatterns, GADTs, StandaloneDeriving, OverloadedStrings #-}
-- |
-- Module      : Criterion
-- Copyright   : (c) 2009-2014 Bryan O'Sullivan
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : GHC
--
-- 'Versus'-type bgroup support

module Criterion.Versus
       (
         vscsv
       , versusReports
       , VersusReport(..)
       ) where

import Criterion.Types
import Criterion.IO.Printf (writeCsv)
import Criterion.Monad (Criterion)
import Control.Monad
import Control.Monad.Trans (liftIO)
import Control.Monad.Reader (asks)
import Control.Arrow ((&&&))
import Data.Function (on)
import qualified Data.Csv as Csv
import Data.List (groupBy, sortBy)
import Statistics.Resampling.Bootstrap (Estimate(..))
import Data.Aeson (ToJSON(..), object, encode, (.=))

data VersusReport where
  VersusReport :: (Show l, Ord l) => {
    vsReportDescription :: String
  , vsReportDataPoints :: [l]
  , vsReportData :: [(String, [Estimate])]
  , vsReportIndices :: [((String, l), Int)]
  } -> VersusReport
deriving instance Show VersusReport

{- This instance is actually incomplete -}
instance ToJSON VersusReport where
  toJSON VersusReport{
      vsReportDescription = desc
    , vsReportDataPoints  = dp
    , vsReportData        = d
    } = object [ "name"       .= toJSON desc
               , "dataPoints" .= toJSON (map show dp)
               , "data"       .= toJSON (map mkArr d)]
        where
          mkArr (alg, p) = object [ "alg"  .= toJSON alg
                                  , "data" .= toJSON p]

vscsv :: [VersusReport] -> Criterion ()
vscsv = mapM_ f
  where f VersusReport{
            vsReportDescription = d
          , vsReportData        = r
          , vsReportDataPoints  = p} = do
          file <- asks vsCsvFile
          writeCsv file [d]
          writeCsv file $ "name":(map show p)
          forM_ r $ \(a, m) -> writeCsv file $ a:(map (show . estPoint) m)

versusReports :: [VersusReport] -> [Report] -> [VersusReport]
versusReports vrpts rpts = map (vsReport rpts) vrpts

vsReport :: [Report] -> VersusReport -> VersusReport
vsReport rpts vr@VersusReport{vsReportIndices = indices} =
  vr{vsReportData = map f l}
  where
    valg = fst . fst
    venv = snd . fst
    l    = groupBy' ((==) `on` valg) valg indices
    rpts' = sortBy (compare `on` reportNumber) rpts
    f (alg, idx) = (alg, [anMean . reportAnalysis $ rpts'!!i
                         | (_, i)<-sortBy (compare `on` venv) idx])

groupBy' :: (a -> a -> Bool) -> (a -> b) -> [a] -> [(b, [a])]
groupBy' f g = map (g . head &&& id) . groupBy f