{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

-- :set -fobject-code in GHCi
module Picture where

import Foreign hiding (rotate)
import Foreign.C.Types

import Control.Applicative

import Data.Complex

default (Int, Float, Integer)

-- Strict
data Color =
  Color
    { r, g, b, a :: !Float
    }
  deriving (Show)

data CColor =
  CColor
    { rC, gC, bC :: !CUChar
    }
  deriving (Show)

instance Storable CColor where
  alignment _ = 1
  sizeOf _ = 3
  peek ptr =
    CColor <$> peekByteOff ptr 0 <*> peekByteOff ptr 1 <*> peekByteOff ptr 2
  poke ptr (CColor d c i) = do
    pokeByteOff ptr 0 d
    pokeByteOff ptr 1 c
    pokeByteOff ptr 2 i

type Point = (Float, Float)

type Image a = Point -> a

type Region = Image Bool

type Filter c = Image c -> Image c

type FilterC = Filter Color

type ImageC = Image Color

renderRegion :: Region -> Image Color
renderRegion r p =
  if r p
    then black
    else white

cOver :: Color -> Color -> Color
cOver (Color r1 g1 b1 a1) (Color r2 g2 b2 a2) =
  Color (h r1 r2) (h g1 g2) (h b1 b2) (h a1 a2)
  where
    h x1 x2 = x1 + (1 - a1) * x2

bilerpBRBW = bilerpC black red blue white

bilerpC ll lr ul ur (wx, wy) = lerpC wy (lerpC wx ll lr) (lerpC wx ul ur)

over :: ImageC -> ImageC -> ImageC
top `over` bot = \p -> top p `cOver` bot p

instance Semigroup Color where
  (<>) = cOver

instance Monoid Color where
  mempty = invisible

invisible = Color 0 0 0 0

black = Color 0 0 0 1

white = Color 1 1 1 1

red = Color 1 0 0 1

green = Color 0 1 0 1

blue = Color 0 0 1 1

yellow = Color 1 1 0 1

cond :: Image Bool -> Image c -> Image c -> Image c
cond =
  liftA3
    (\a b c ->
       if a
         then b
         else c)

lerpI :: Image Frac -> ImageC -> ImageC -> ImageC
lerpI = liftA3 lerpC

empty = const invisible

whiteI = const white

blackI = const black

redI = const red

blueI = const blue

greenI = const green

yellowI = const yellow

blackWhiteIm reg = cond reg blackI whiteI

blueYellowIm reg = cond reg blueI yellowI

crop :: Region -> FilterC
crop reg im = cond reg im mempty

type Warp = Point -> Point

type Vector = (Float, Float)

translateP :: Vector -> Warp
translateP (dx, dy) (x, y) = (x + dx, y + dy)

scaleP :: Vector -> Warp
scaleP (sx, sy) (x, y) = (sx * x, sy * y)

uscaleP :: Float -> Warp
uscaleP s = scaleP (s, s)

rotateP :: Float -> Warp
rotateP θ (x, y) = (x * cos θ - y * sin θ, y * cos θ + x * sin θ)

udisk :: Region
udisk p = distO p < 1

invWarp warp im = im . warp

translate :: Vector -> Filter c
translate (dx, dy) = invWarp (translateP (-dx, -dy))

scale :: Vector -> Filter c
scale (sx, sy) = invWarp (scaleP (1 / sx, 1 / sy))

-- Laws
-- uscale x . uscale y = uscale (x * y)
uscale :: Float -> Filter c
uscale s = invWarp (uscaleP (1 / s))

rotate :: Float -> Filter c
rotate θ = invWarp (rotateP (-θ))

swirlP :: Float -> Warp
swirlP r p = rotateP (distO p * (2 * pi / r)) p

swirl :: Float -> Filter c
swirl r = invWarp (swirlP (-r))

square n (x, y) = abs x <= n / 2 && abs y <= n / 2

-- Figure to visually check orientation
disks =
  renderRegion $
  f <$> udisk <*> translate (2, 2) udisk <*> translate (3, 2) udisk <*> vstrip
  where
    f a b c d = a || b || c || d

-- Figures from "Functional Images" by Conal Elliot.
fig1 = renderRegion $ vstrip

fig2 = renderRegion $ checker

fig3 = renderRegion $ altRings

fig4 = renderRegion $ polarChecker 10

fig5 = renderRegion $ uscale 0.05 gasket

fig6 = renderFrac $ wavDist

fig7 = bilerpBRBW

fig8 = lerpI wavDist (blackWhiteIm (polarChecker 10)) (blueYellowIm checker)

fig9 = ybRings

fig10 = renderRegion udisk

fig11 = renderRegion $ swirl 1 vstrip

fig13 t = renderRegion (swirlingXPos t)

fig14 = renderRegion $ annulus 0.5

fig17 = renderRegion $ shiftXor 2.6 altRings

fig18 = renderRegion $ xorgon 8 (7 / 4) altRings

-- Typo in original paper, 0.2 5 should be 0.25
fig19 = crop (wedgeAnnulus 0.25 10) ybRings

fig20 = crop (swirl 2 (wedgeAnnulus 0.25 10)) ybRings

fig22 = translate (-20, -20) tiledBilerp

fig23 = renderRegion $ uscale 3.5 (radInvert checker)

fig24 = rippleRad 8 0.3 ybRings

fig25 = \t -> rippleRad 8 (cos t / 2) ybRings

fig26 = rippleRad 8 0.3 $ cropRad 1 $ ybRings

fig27 = cropRad 1 $ rippleRad 8 0.3 $ ybRings

fig28 = swirl 8 $ rippleRad 5 0.3 $ cropRad 5 $ ybRings

-- Typo in original paper, washer (1/2) (pi/2) 1 tiledBilerp
-- TODO: Fix visual mismatch with paper.
fig32 = washer (1 / 2) (pi / 2) tiledBilerp

-- Convert an image to a static animation.
static :: Image c -> Anim c
static = const

-- The main animation to run.
-- mainAnim = static . renderFrac $ mandel
mainAnim :: Anim Color
mainAnim = fig32

calculate (x, y, t) = adjust (mainAnim t') (x', y')
  where
    duration = 5000 -- Duration of the animation in ms
    range = 2 * pi -- Time range of the animation
    -- Time normalized from 0 to 1
    normalizedTime = fromIntegral (t `mod` duration) / fromIntegral duration
    scaledTime = range * normalizedTime
    -- scaledTime = fromIntegral t / 10000
    (x', y', t') = (fromIntegral x, fromIntegral y, scaledTime)
    adjust = adjustToWindow
    -- Adjust an image to a window, by translating, scaling and flipping
    adjustToWindow :: Filter c
    adjustToWindow =
      translate (screenWidth / 2, screenHeight / 2) .
      uscale 60 . flipY
    flipY p (x, y) = p (x, -y)

(screenWidth, screenHeight) = (640, 480)

foreign export ccall fillPixelBuffer :: Ptr CColor -> CInt -> IO ()
fillPixelBuffer arr t = pokeArray arr (map calc l1)
  where
    l1 = [(x, y) | y <- [0 .. screenHeight - 1], x <- [0 .. screenWidth - 1]]
    calc (x, y) =
      CColor (truncate (255 * r)) (truncate (255 * g)) (truncate (255 * b))
      where
        Color r g b _ = calculate (x, y, t)

checker :: Region
checker (x, y) = even (floor x + floor y)

vstrip :: Region
vstrip (x, y) = abs x <= 1 / 2

overlay :: ImageC -> ImageC -> ImageC
overlay = liftA2 (<>)

-- Sierpinski triangle
gasket :: Region
gasket (x, y) = floor x .|. floor y == (floor x :: Integer)

altRings p = even (floor (distO p))

distO (x, y) = sqrt (x * x + y * y)

type Frac = Float -- in [0, 1]

wavDist :: Image Frac
wavDist p = (1 + cos (pi * distO p)) / 2

ybRings = lerpI wavDist blueI yellowI

renderFrac :: Image Frac -> ImageC
renderFrac f (x, y) = lerpC (f (x, y)) black white

lerpC :: Frac -> Color -> Color -> Color
lerpC w (Color r1 g1 b1 a1) (Color r2 g2 b2 a2) =
  Color (h r1 r2) (h g1 g2) (h b1 b2) (h a1 a2)
  where
    h x1 x2 = w * x1 + (1 - w) * x2

lighten x c = lerpC x c white

darken x c = lerpC x c black

type PolarPoint = (Float, Float)

fromPolar :: Point -> PolarPoint
fromPolar (ρ, θ) = (ρ * cos θ, ρ * sin θ)

toPolar :: PolarPoint -> Point
toPolar (x, y) = (distO (x, y), atan2 y x)

polarChecker :: Int -> Region
polarChecker = polarize checker

polarize :: Image c -> Int -> PolarPoint -> c
polarize pat n = pat . sc . toPolar
  where
    sc (ρ, θ) = (ρ, θ * n' / π)
    π = pi
    n' = fromIntegral n

-- Animations
type Time = Float

type Anim c = Time -> Image c

swirlingXPos :: Anim Bool
swirlingXPos t = swirl (t * t) xPos

xPos :: Region
xPos (x, y) = x > 0

-- Region algebra
intersect :: Region -> Region -> Region
intersect = liftA2 (&&)

union :: Region -> Region -> Region
union = liftA2 (||)

xorR :: Region -> Region -> Region
xorR = liftA2 xor
  where
    xor True b = not b
    xor _ b = b

compR :: Region -> Region
compR = fmap not

universeR :: Region
universeR = const True

emptyR :: Region
emptyR = const False

r \\ r' = r `intersect` compR r'

annulus :: Frac -> Region
annulus inner = udisk \\ uscale inner udisk

radReg :: Int -> Region
radReg n = test . toPolar
  where
    test (r, a) = even (floor (a * fromIntegral n / pi))

wedgeAnnulus :: Frac -> Int -> Region
wedgeAnnulus inner n = annulus inner `intersect` radReg n

shiftXor :: Float -> Region -> Region
shiftXor r reg = reg' r `xorR` reg' (-r)
  where
    reg' d = translate (d, 0) reg

-- Typo in the original paper, missing argument g
xorgon :: Int -> Float -> Region -> Region
xorgon n r g = xorRs (rf <$> [0 .. n - 1])
  where
    rf :: Int -> Region
    rf i = translate (fromPolar (r, a)) g
      where
        a = fromIntegral i * 2 * pi / fromIntegral n

xorRs :: [Region] -> Region
xorRs = foldr xorR emptyR

tileP :: Vector -> Warp
tileP (w, h) (x, y) = (wrap' w x, wrap' h y)

tiledBilerp = about (1 / 2, 1 / 2) (tile (1, 1)) bilerpBRBW

type HyperFilter c = Filter c -> Filter c

about :: Point -> HyperFilter c
about (x, y) filt = translate (x, y) . filt . translate (-x, -y)

wrap :: Float -> Float -> Float
wrap w x = w * fracPart (x / w)
  where
    fracPart t = t - fromIntegral (truncate t)

wrap' :: Float -> Float -> Float
wrap' w x = wrap w (x + w / 2) - w / 2

tile :: Vector -> Filter c
tile size = invWarp (tileP size)

swirlP' r = polarWarp (\(ρ, θ) -> (ρ, θ + ρ * (2 * pi / r)))

polarWarp warp = fromPolar . warp . toPolar

radInvertP :: Warp
radInvertP = polarWarp (\(ρ, θ) -> (1 / ρ, θ))

radInvert :: Filter c
radInvert = invWarp radInvertP

rippleRadP :: Int -> Float -> Warp
rippleRadP n s =
  polarWarp $ (\(ρ, θ) -> (ρ * (1 + s * sin (fromIntegral n * θ)), θ))

rippleRad :: Int -> Float -> Filter c
rippleRad n s = invWarp (rippleRadP n (-s))

cropRad :: Float -> FilterC
cropRad r = crop (uscale r udisk)

wiggleRotateP :: Float -> Float -> Time -> Warp
wiggleRotateP cycles θmax t = polarWarp warp
  where
    warp (r, a) = (r, a + θmax * sin (t + dt))
      where
        dt = 2 * pi * cycles * (r - 1 / 2)

wiggleRotate :: Float -> Float -> Time -> Filter c
wiggleRotate cycles θmax t = invWarp (wiggleRotateP cycles θmax t)

washer :: Float -> Float -> ImageC -> Time -> ImageC
washer cycles θmax im t = cropRad 1 $ wiggleRotate cycles θmax t $ im

-- Constants
maxIter :: Int -- Max iterations
maxIter = 750

width :: Int -- Image width
height :: Int -- Image height
width = 400

height = 400

-- Note: aspect ratio of (minX, minY), (maxX, maxY) must
-- match aspect ratio of (width, height)
minX :: Float -- Min x-coordinate of graph
maxX :: Float -- Max x-coordinate of graph
minY :: Float -- Min y-coordinate of graph
maxY :: Float -- Max y-coordinate of graph
-- For the zoomed in part of the Mandelbrot:
--minX = -0.826341244461360116
--maxX = -0.8026423086165848822
--minY = -0.2167936114403439588
--maxY = -0.193094675595568725
--For a full view of the mandelbrot
minX = -2.5

maxX = 1.5

minY = -2

maxY = 2

-- The actual fractal part. It basically works on a matrix, which we
-- will call M, that represents a grid of points on the
-- graph. Essentially, M[i, j] is (xList[j], yList[i])
xList :: [Float]
yList :: [Float]
xList = [minX,(minX + ((maxX - minX) / (fromIntegral width - 1))) .. maxX]

yList =
  reverse [minY,(minY + ((maxY - minY) / (fromIntegral height - 1))) .. maxY]

etaFraction :: Complex Float -> Float
etaFraction z = (log (log (magnitude z))) / log 2

smoothEta :: Int -> Complex Float -> Float -- Smooth escape time algorithm value
smoothEta iter z = (fromIntegral iter - etaFraction z) / fromIntegral maxIter

-- Gets the color for the point, in range [0, 1]
color :: Int -> Complex Float -> Float
color iter z = smoothEta iter z -- Smooth escape time algorithm (and invert)

-- color iter z = fromIntegral iter / fromIntegral maxIter
interpolate :: Float -> Int -- Adds an interpolation curve for interpolating color
interpolate v = truncate ((v ^ 12) * 255) -- Polynomial curve

-- interpolate v = chr (truncate(v * 255)) -- Linear
-- The actual fractal algorithm!
frac :: Complex Float -> Complex Float -> Int -> Int
frac c@(x :+ y) z iter
  | iter >= maxIter = 255 -- never escaped, return color value of 255
  | let p = sqrt ((x - 0.25) ^ 2 + y ^ 2)
     in x <= p - 2 * p ^ 2 + 0.25 && (x + 1) ^ 2 + y ^ 2 <= 0.625 = 255
  | otherwise =
    let z' = z * z + c
     in if ((realPart z') * (realPart z') + (imagPart z') * (imagPart z')) > 4
          then interpolate (color iter z')
          else frac c z' (iter + 1)

mandel :: Image Frac
mandel (x, y) = fromIntegral (frac (x :+ y) (0 :+ 0) 0) / 255

-- dotGridGradient :: Int -> Int -> Float -> Float -> IO Float
-- dotGridGradient ix iy x y = do
--   (v1, v2) <- do p' <- randomIO :: IO Float
--                  let (_, p) = properFraction p'
--                  pure (p, sqrt (1 - p ^ 2))
--   let c1 = dx * v1
--   let c2 = dy * v2
--   let res = (c1 + c2)
--   -- print res
--   pure res
--   where
--     dx = x - fromIntegral ix
--     dy = y - fromIntegral iy

-- perlin :: Image (IO Frac)
-- perlin (x, y) = do
--   let (x0, sx) = properFraction x
--   let (y0, sy) = properFraction y
--   let x1 = x0 + 1
--   let y1 = y0 + 1
--   n0 <- dotGridGradient x0 y0 x y
--   n1 <- dotGridGradient x1 y0 x y
--   let ix0 = lerp n0 n1 sx
--   n0' <- dotGridGradient x0 y1 x y
--   n1' <- dotGridGradient x1 y1 x y
--   let ix1 = lerp n0' n1' sx
--   let res' = lerp ix0 ix1 sy
--   let res = abs res'
--   pure res
--   where
--     lerp a0 a1 w = (1.0 - w) * a0 + w * a1


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}

-- This is a simple generative adversarial network to make pictures
-- of numbers similar to those in MNIST.
--
-- It demonstrates a different usage of the library. Within about 15
-- minutes it was producing examples like this:
--
--               --.
--     .=-.--..#=###
--     -##==#########.
--     #############-
--   -###-.=..-.-==
--   ###-
--   .###-
--   .####...==-.
--   -####=--.=##=
--   -##=-     -##
--             =##
--           -##=
--           -###-
--         .####.
--         .#####.
-- ...---=#####-
-- .=#########.         .
--   .#######=.          .
--     . =-.
--
-- It's a 5!
--
import           Control.Applicative
import           Control.Monad
import           Control.Monad.Random

import           Codec.Compression.GZip ( decompress )
import           Data.Serialize ( Get )
import qualified Data.Serialize as Serialize
import qualified Data.ByteString.Lazy as B

import           Data.List ( foldl' )
import           Data.List.Split ( chunksOf )
import           Data.Maybe ( fromMaybe )
#if ! MIN_VERSION_base(4,13,0)
import           Data.Semigroup ( (<>) )
#endif

import           Data.Word ( Word32 , Word8 )
import qualified Data.Vector.Storable as V

import qualified Numeric.LinearAlgebra.Static as SA
import           Numeric.LinearAlgebra.Data ( toLists )

import           Options.Applicative
import           System.FilePath ( (</>) )

import           Grenade
import           Grenade.Utils.OneHot

type Discriminator =
  Network
    '[ Convolution 1 10 5 5 1 1, Pooling 2 2 2 2, Relu
     , Convolution 10 16 5 5 1 1, Pooling 2 2 2 2, Relu
     , Reshape, FullyConnected 256 80, Logit, FullyConnected 80 1, Logit]
    '[ 'D2 28 28
     , 'D3 24 24 10, 'D3 12 12 10, 'D3 12 12 10
     , 'D3 8 8 16, 'D3 4 4 16, 'D3 4 4 16
     , 'D1 256, 'D1 80, 'D1 80, 'D1 1, 'D1 1]

type Generator =
  Network
    '[ FullyConnected 80 256, Relu, Reshape
     , Deconvolution 16 10 5 5 2 2, Relu
     , Deconvolution 10 1 8 8 2 2, Logit]
    '[ 'D1 80
     , 'D1 256, 'D1 256, 'D3 4 4 16
     , 'D3 11 11 10, 'D3 11 11 10
     , 'D2 28 28, 'D2 28 28 ]

randomDiscriminator :: MonadRandom m => m Discriminator
randomDiscriminator = randomNetwork

randomGenerator :: MonadRandom m => m Generator
randomGenerator = randomNetwork

trainExample :: LearningParameters -> Discriminator -> Generator -> S ('D2 28 28) -> S ('D1 80) -> ( Discriminator, Generator )
trainExample rate discriminator generator realExample noiseSource
 = let (generatorTape, fakeExample)       = runNetwork generator noiseSource

       (discriminatorTapeReal, guessReal) = runNetwork discriminator realExample
       (discriminatorTapeFake, guessFake) = runNetwork discriminator fakeExample

       (discriminator'real, _)            = runGradient discriminator discriminatorTapeReal ( guessReal - 1 )
       (discriminator'fake, _)            = runGradient discriminator discriminatorTapeFake guessFake
       (_, push)                          = runGradient discriminator discriminatorTapeFake ( guessFake - 1)

       (generator', _)                    = runGradient generator generatorTape push

       newDiscriminator                   = foldl' (applyUpdate rate { learningRegulariser = learningRegulariser rate * 10}) discriminator [ discriminator'real, discriminator'fake ]
       newGenerator                       = applyUpdate rate generator generator'
   in ( newDiscriminator, newGenerator )

ganTest :: (Discriminator, Generator) -> Int -> FilePath -> LearningParameters -> IO (Discriminator, Generator)
ganTest (discriminator0, generator0) iterations dataDir rate = do
  -- Note that for this example we use only the samples, and not the labels
  trainData      <- fmap fst <$> readMNIST (dataDir </> "train-images-idx3-ubyte.gz")
                                           (dataDir </> "train-labels-idx1-ubyte.gz")

  foldM (runIteration trainData) ( discriminator0, generator0 ) [1..iterations]

    where

  showShape' :: S ('D2 a b) -> IO ()
  showShape' (S2D mm) = putStrLn $
    let m  = SA.extract mm
        ms = toLists m
        render n'  | n' <= 0.2  = ' '
                   | n' <= 0.4  = '.'
                   | n' <= 0.6  = '-'
                   | n' <= 0.8  = '='
                   | otherwise =  '#'

        px = (fmap . fmap) render ms
    in unlines px

  runIteration :: [S ('D2 28 28)] -> (Discriminator, Generator) -> Int -> IO (Discriminator, Generator)
  runIteration trainData ( !discriminator, !generator ) _ = do
    trained'    <- foldM ( \(!discriminatorX, !generatorX ) realExample ->
                             trainExample rate discriminatorX generatorX realExample <$> randomOfShape )
                         ( discriminator, generator ) trainData

    showShape' . snd . runNetwork (snd trained') =<< randomOfShape

    return trained'

data GanOpts = GanOpts FilePath Int LearningParameters (Maybe FilePath) (Maybe FilePath)

mnist' :: Parser GanOpts
mnist' = GanOpts <$> argument str (metavar "DATADIR")
                 <*> option auto (long "iterations" <> short 'i' <> value 15)
                 <*> (LearningParameters
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0.0005)
                       )
                 <*> optional (strOption (long "load"))
                 <*> optional (strOption (long "save"))

main :: IO ()
main = do
  GanOpts mnist iter rate load save <- execParser (info (mnist' <**> helper) idm)
  putStrLn "Training stupidly simply GAN"
  nets0 <- case load of
    Just loadFile -> netLoad loadFile
    Nothing -> (,) <$> randomDiscriminator <*> randomGenerator

  nets1 <- ganTest nets0 iter mnist rate
  case save of
    Just saveFile -> B.writeFile saveFile $ Serialize.runPutLazy (Serialize.put nets1)
    Nothing -> return ()

-- Adapted from https://github.com/tensorflow/haskell/blob/master/tensorflow-mnist/src/TensorFlow/Examples/MNIST/Parse.hs
-- Could also have used Data.IDX, although that uses a different Vector variant from that need for fromStorable
readMNIST :: FilePath -> FilePath -> IO [(S ( 'D2 28 28), S ( 'D1 10))]
readMNIST iFP lFP = do
  labels  <- readMNISTLabels lFP
  samples <- readMNISTSamples iFP
  return $ zip
    (fmap (fromMaybe (error "bad samples") . fromStorable) samples)
    (fromMaybe (error "bad labels") . oneHot . fromIntegral <$> labels)

-- | Check's the file's endianess, throwing an error if it's not as expected.
checkEndian :: Get ()
checkEndian = do
  magic <- Serialize.getWord32be
  when (magic `notElem` ([2049, 2051] :: [Word32]))
    $ error "Expected big endian, but image file is little endian."

-- | Reads an MNIST file and returns a list of samples.
readMNISTSamples :: FilePath -> IO [V.Vector Double]
readMNISTSamples path = do
  raw <- decompress <$> B.readFile path
  either fail ( return . fmap (V.map normalize) ) $ Serialize.runGetLazy getMNIST raw
 where
  getMNIST :: Get [V.Vector Word8]
  getMNIST = do
    checkEndian
    -- Parse header data.
    cnt    <- fromIntegral <$> Serialize.getWord32be
    rows   <- fromIntegral <$> Serialize.getWord32be
    cols   <- fromIntegral <$> Serialize.getWord32be
    -- Read all of the data, then split into samples.
    pixels <- Serialize.getLazyByteString $ fromIntegral $ cnt * rows * cols
    return $ V.fromList <$> chunksOf (rows * cols) (B.unpack pixels)

  normalize :: Word8 -> Double
  normalize = (/ 255) . fromIntegral
  -- There are other normalization functions in the literature, such as
  -- normalize = (/ 0.3081) . (`subtract` 0.1307) . (/ 255) . fromIntegral
  -- but we need values in the range [0..1] for the showShape' pretty printer

-- | Reads a list of MNIST labels from a file and returns them.
readMNISTLabels :: FilePath -> IO [Word8]
readMNISTLabels path = do
  raw <- decompress <$> B.readFile path
  either fail return $ Serialize.runGetLazy getLabels raw
 where
  getLabels :: Get [Word8]
  getLabels = do
    checkEndian
    -- Parse header data.
    cnt <- fromIntegral <$> Serialize.getWord32be
    -- Read all of the labels.
    B.unpack <$> Serialize.getLazyByteString cnt

netLoad :: FilePath -> IO (Discriminator, Generator)
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $
    Serialize.runGetLazy (Serialize.get :: Get (Discriminator, Generator)) modelData


{-|
Module      : MachineLearning.TIR
Description : TIR expression data structures
Copyright   : (c) Fabricio Olivetti de Franca, 2022
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

The TIR expression  represents a function of the form:

\[
f(x) = g(\sum_{i}{w_i \cdot t_i(\prod_{j}{x_j^{k_{ij}})}} / (1 + \sum_{i}{w_i \cdot t_i(\prod_{j}{x_j^{k_{ij}})}})
\]

with \(t_i\) being a transformation function, \(g\) an invertible function, \(w_i\) a linear coefficient, \(k_{ij}\)
the interaction strength.

Any given expression can be represented by two lists of terms, with each term
being composed of a transformatioon function and an interaction.
The transformation function is represented by a `Function` sum type.
The interaction is represented as a list of tuple of ints where the key is the
predictor index and the value is the strength of the predictor in this
term. Strengths with a value of zero are omitted.
-}
module MachineLearning.TIR where

import Control.Evolution
import Control.Monad.State.Strict
import Data.List.Split
import Data.SRTree
import System.Random
import Control.DeepSeq                           (NFData, rnf)
import Data.List                                 (delete)
import Data.SRTree.Print                         (showDefault)
import Data.Vector                               (Vector)
import qualified Data.Vector             as V
import qualified Data.Vector.Storable    as VS
import qualified Numeric.LinearAlgebra   as LA
import MachineLearning.Utils.Config

-- | `TIR` is a record type composed of the external
-- function `_funY` of type `Function`, numerator `_p`
-- and denominator `_q` of type `Sigma`
data TIR = TIR { _funY :: Function
               , _p :: Sigma
               , _q :: Sigma
               } deriving Show

instance NFData TIR where
  rnf _ = ()

-- | `Sigma` is just a list of terms `Pi`
type Sigma = [Pi]
-- | `Pi` is a triple composed of a coefficient, a `Function`
-- and a list of tuples where `(ix, k)` represents `x ! ix ^ k`.
type Pi    = (Double, Function, [(Int, Int)])

-- | generates a random integer within the specified range.
randomRng :: (Int, Int) -> Rnd Int
randomRng rng = state $ randomR rng
{-# INLINE randomRng #-}

-- | generates a random integer within the specified range excluding zero.
randomRngNZ :: (Int, Int) -> Rnd Int
randomRngNZ rng = do
  x <- randomRng rng
  if x == 0
    then randomRngNZ rng
    else pure x
{-# INLINE randomRngNZ #-}

-- | picks a random element from a list.
randomFrom :: [a] -> Rnd a
randomFrom xs = do
  ix <- randomRng (0, length xs - 1)
  pure (xs !! ix)
{-# INLINE randomFrom #-}

-- | returns a random index of variables provided by the mutation configuration.
randomVar :: MutationCfg -> Rnd (Maybe Int, MutationCfg)
randomVar params = do
  let vars = _vars params
      n    = length vars
  ix <- randomRng (0, n)
  if ix == n
     then pure (Nothing, params)
     else do let x = vars !! ix
             pure (Just x, params{ _vars=delete x vars })

-- | returns a list of random interactions (tuples of variables indeces and exponentes)
-- with parameters provided by the mutation configuration.
randomVars :: MutationCfg -> Rnd [(Int, Int)]
randomVars params = do
  (v, params') <- randomVar params
  k            <- randomRngNZ $ _kRange params
  case v of
    Nothing  -> pure []
    Just var -> do vs <- randomVars params'
                   pure $ (var, k) : vs

-- | returns a random `Pi`
randomPi :: MutationCfg -> Rnd (Maybe Pi)
randomPi params = do
  pis <- randomVars params
  f   <- randomFrom $ _funs params
  if null pis
    then pure Nothing
    else pure $ Just (1.0, f, pis)

-- | returns a random `Sigma`
randomSigma :: MutationCfg -> Int -> Rnd (Sigma, Int)
randomSigma params budget | budget <= 0 = pure ([], budget)
randomSigma params budget = do
  n <- randomRng (0, budget)
  if n == budget
     then pure ([], budget)
     else do term             <- randomPi params
             (terms, budget') <- randomSigma params (budget - spentBudget term)
             case term of
               Nothing -> pure (terms, budget')
               Just t  -> pure (t:terms, budget')

  where
    spentBudget Nothing           = 0
    spentBudget (Just (_, _, ps)) = 1 -- length ps

-- | returns a random `TIR` expression
randomTIR :: MutationCfg -> Rnd TIR
randomTIR params = do
  yf           <- randomFrom $ _yfuns params
  (p, budget') <- randomSigma params $ _budget params
  (q, _)       <- randomSigma params budget'
  if null p
    then randomTIR params
    else pure (TIR yf p q)

-- | We store thee dataset as a vector of columns.
-- Each vector is stored a `Storable`-based vector.
type Column a   = LA.Vector a

-- | A dataset is a `Vector` of `Column`
type Dataset a  = Vector (Column a)

-- | A constraint is a function that gets a symbolic tree
-- as an input and returns non negative `Double` representing
-- how much a constraint was violated.
type Constraint = SRTree Int Double -> Double

-- | An individual in the population is composed of
-- the chromossome, a vector of fitness, a list of
-- coefficients (for multiclass problems it stores
-- one vector of coefficient per class),
-- the constraint violation, the size of the expression,
-- and the penalty value.
data Individual = Individual { _chromo  :: TIR
                             , _fit     :: [Double]
                             , _weights :: [LA.Vector Double]
                             , _constr  :: Double
                             , _len     :: Int
                             , _penalty :: Double
                             }

-- | creates an unevaluated individual.
createIndividual :: TIR -> Individual
createIndividual tir = Individual tir [] [] 0.0 0 0.0

-- | calculates the penalized fitness.
penalizedFit :: Individual -> Double
penalizedFit t = (head . _fit) t + _penalty t
{-# INLINE penalizedFit #-}

-- | replaces the coefficients of a TIR expression
replaceConsts :: TIR -> V.Vector Double -> TIR
replaceConsts (TIR g p q) ws = TIR g p' q'
  where
    (p', ws1) = runState (traverse replaceWeight p) (V.toList ws)
    (q', ws2) = runState (traverse replaceWeight q) ws1

replaceWeight :: Pi -> State [Double] Pi
replaceWeight (w, g, h) = state $ \ws -> case ws of
                                           (wi:ws') -> ((wi, g, h), ws')
                                           []       -> error $ show h -- ((w, g, h), [])

instance Eq Individual where
    t1 == t2 = penalizedFit t1 == penalizedFit t2
instance Ord Individual where
    t1 <= t2 = penalizedFit t1 <= penalizedFit t2

instance NFData Individual where
  rnf _ = ()

instance Solution Individual where
  _getFitness = head . _fit
  _isFeasible = (<1e-12) . _constr

-- | creates a symbolic tree from a TIR expression.
assembleTree :: Double -> TIR -> SRTree Int Double
assembleTree bias (TIR f p q) = Fun f ((Const bias + assemble p) / (1 + assemble q))
  where
    -- assemble :: Sigma ix val -> SRTree ix val
    assemble []      = 0
    assemble [p']    = mk p'
    assemble (p':ps) = mk p' + assemble ps

    -- mk :: Pi ix val -> SRTree ix val
    mk (v, g, ts) = Const v * Fun g (foldr (\(ix, k) acc -> acc * Pow (Var ix) k) 1 ts)

-- | pretty print a solution.
prettyPrintsolution :: Individual -> String
prettyPrintsolution sol | Prelude.null (_fit sol) = error "unevaluated solution"
prettyPrintsolution sol = concat [ "Expression:\n", (showDefault . assembleTree bias . _chromo) sol, "\n"
                                 , "Fitness: ", (show . head . _fit) sol, "\n"
                                 , "Constraints: ", (show . _constr) sol, "\n"
                                 , "Length: ", (show . _len) sol, "\n"
                                 , "Penalty: ", (show . _penalty) sol, "\n"
                                 ]

  where bias = V.head $ VS.convert $ head $ _weights sol


{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Basic types used throughout the library.
--
--   Includes a probability type (based on sampling functions, with a tweak for
--   collected samples) and associated instances, as well as types used for
--   MCMC.

module Math.Probably.Types where

import Control.Applicative
import Control.Monad.State.Strict
import Numeric.LinearAlgebra
import System.Random.Mersenne.Pure64

type Seed = PureMT

data Prob a =
    Sampler { unSampler :: Seed -> (a, Seed) }
  | Samples [a]

instance Functor Prob where
  fmap f (Sampler sf) = Sampler $ \rs -> let (x,rs') = sf rs in (f x, rs')
  fmap f (Samples xs) = Samples $ map f xs

instance Applicative Prob where
  pure x = Sampler (\rs-> (x, rs))
  (Sampler sff) <*> (Sampler sfx) = Sampler $ \rs ->
    let (f ,rs') = sff rs
        (x, rs'') = sfx rs'
    in (f x, rs'')
  _ <*> _ = error "Prob (<*>): unsupported pattern"

instance Monad Prob where
  return = pure
  (Sampler sf) >>= f = Sampler $ \rs->
    let (x, rs'::Seed) = sf rs
        nextProb = f x
    in case nextProb of
         Sampler g -> g rs'
         Samples xs -> primOneOf xs rs'

  (Samples xs) >>= f = Sampler $ \rs->
    let (x, rs'::Seed) = primOneOf xs rs
        nextProb = f x
    in case nextProb of
         Sampler g -> g rs'
         Samples ys -> primOneOf ys rs'

primOneOf :: [a] -> Seed -> (a, Seed)
primOneOf xs seed
  = let (u, nextSeed) = randomDouble seed
        idx = floor $ realToFrac u * realToFrac (length xs )
    in (xs !! idx, nextSeed)

type DiscreteParams   = Vector Int
type ContinuousParams = Vector Double
type Parameters       = (DiscreteParams, ContinuousParams)

type LogObjective     = Parameters -> Double
type Gradient         = ContinuousParams -> ContinuousParams

type Particle         = (ContinuousParams, ContinuousParams)

type Transition t     = StateT (Chain t) Prob Parameters

-- | State of a Markov chain.  Note that the objective function itself is
--   included in the state, which allows the possibility of annealing schedules
--   and the like.
data Chain t = Chain {
    parameterSpacePosition :: Parameters
  , objectiveFunction      :: Target
  , objectiveValue         :: Double
  , tunables               :: t
  }

-- | A target to sample, consisting of a log objective function and possibly
--   its gradient.  The gradient is only taken with respect to continuous
--   parameters.
data Target = Target {
    logObjective :: Parameters -> Double
  , gradient     :: Maybe (ContinuousParams -> ContinuousParams)
  }

-- | Convenience constructor for targets and gradients.
createTargetWithGradient :: LogObjective -> Gradient -> Target
createTargetWithGradient f g = Target f (Just g)

-- | Convenience constructor for targets without gradients.
createTargetWithoutGradient :: LogObjective -> Target
createTargetWithoutGradient f = Target f Nothing

handleGradient :: Maybe t -> t
handleGradient Nothing  = error "handleGradient: no gradient provided"
handleGradient (Just g) = g

-- | The dual-averaging implementation of NUTS requires quite a few tuning
--   parameters, which are all held in this type.
data DualAveragingParameters = DualAveragingParameters {
    mAdapt    :: !Int
  , delta     :: !Double
  , mu        :: !Double
  , gammaP    :: !Double
  , tau0      :: !Double
  , kappa     :: !Double
  , daStep    :: !Double
  , daStepAvg :: !Double
  , daH       :: !Double
  } deriving (Eq, Show)

defaultDualAveragingParameters :: Double -> Int -> DualAveragingParameters
defaultDualAveragingParameters e burnInPeriod = DualAveragingParameters {
    mu        = log (10 * e)
  , delta     = 0.5
  , mAdapt    = burnInPeriod
  , gammaP    = 0.05
  , tau0      = 10
  , kappa     = 0.75
  , daStep    = e
  , daStepAvg = e
  , daH       = 0
  }


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TemplateHaskell       #-}

module Lib where

import           Control.Monad.Except                      (MonadError,
                                                            runExceptT,
                                                            throwError)
import           Control.Monad.IO.Class                    (MonadIO, liftIO)
import           Data.ByteString                           (ByteString)
import           Data.FileEmbed                            (embedFile)
import           Data.Monoid                               ((<>))
import           Data.Scientific                           (Scientific,
                                                            toRealFloat)
import           Debug.Trace
import           GHC.TypeLits                              ()
import           Graphics.Plot                             (mplot)
import           Graphics.Rendering.Chart.Backend.Diagrams (toFile)
import           Graphics.Rendering.Chart.Easy             (def, layout_title,
                                                            line, plot, points,
                                                            re, (.=))
import           Numeric.GSL.Minimization                  (MinimizeMethod (NMSimplex),
                                                            MinimizeMethodD (SteepestDescent, VectorBFGS2),
                                                            minimizeD,
                                                            minimizeV,
                                                            minimizeVD)
import           Numeric.LinearAlgebra                     (Container, accum,
                                                            matFunc, tr, ( #> ),
                                                            (<.>))
import qualified Numeric.LinearAlgebra                     as Matrix
import           Numeric.LinearAlgebra.Data                (Matrix, Vector,
                                                            asColumn, asRow,
                                                            cols, cond,
                                                            dropColumns,
                                                            dropRows, rows,
                                                            takeColumns,
                                                            takeRows, (><),
                                                            (|||))
import qualified Numeric.LinearAlgebra.Data                as MData
import           Numeric.LinearAlgebra.Static              (L, dim, matrix, mul)
import qualified Numeric.LinearAlgebra.Static              as Static
import           Text.Megaparsec                           (Dec, ParseError,
                                                            Token, char,
                                                            newline, parse,
                                                            sepBy)
import           Text.Megaparsec.ByteString                (Parser)
import           Text.Megaparsec.Lexer                     (number)

data Error =
  ParseFailed (ParseError (Token ByteString) Dec)
  deriving (Show, Eq)

------------------------------------------------------------
-- Parsing.
rawDataParser :: Parser [[Scientific]]
rawDataParser =
  filter (not . null) <$> (number `sepBy` char ',') `sepBy` newline

toMatrix :: [[Scientific]] -> Matrix Double
toMatrix = Matrix.fromLists . fmap (fmap toRealFloat)

wineData :: ByteString
wineData = $(embedFile "data/wine.data")

wineParser :: Parser (Matrix Double)
wineParser = toMatrix <$> rawDataParser

loadData
  :: (MonadError Error m, MonadIO m)
  => m (Matrix Double)
loadData = either (throwError . ParseFailed) pure (parse wineParser "" wineData)

------------------------------------------------------------
-- Matrices
a :: L 3 2
a = matrix [1 .. 6]

b :: L 2 1
b = matrix [1 .. 2]

c :: L 3 1
c = mul a b

------------------------------------------------------------
-- Linear Regression
initialTheta :: Int -> Vector Double
initialTheta n = MData.vector $ replicate n 0

------------------------------------------------------------
-- Charts
signal :: [Double] -> [(Double, Double)]
signal xs =
  [(x, (sin (x * 3.14159 / 45) + 1) / 2 * sin (x * 3.14159 / 5)) | x <- xs]

writeChart :: IO ()
writeChart =
  toFile def "example.svg" $ do
    layout_title .= "Amplitude Modulation"
    plot (line "am" [signal [0,0.5 .. 800]])
    plot (points "am points" (signal [0,7 .. 800]))

------------------------------------------------------------
-- Main
main :: IO ()
main = do
  result <-
    runExceptT $ do
      dataset <- loadData
      process dataset
  print result

data Result = Result
  { crossValidationSet    :: Matrix Double
  , crossValidationResult :: Vector Double
  , finalTheta            :: Vector Double
  , finalHyp              :: Vector Double
  , finalCost             :: Double
  } deriving (Show)

process
  :: (MonadIO m, MonadError Error m)
  => Matrix Double -> m Result
process dataSet = do
  let is n v = cond v n 0 1 0
  let target = 1
  let scale :: Matrix Double =
        accum
          (MData.ident 13)
          (*)
          [((0, 0), 0.1), ((3, 3), 0.1), ((4, 4), 0.01), ((12, 12), 0.001)]
  let m = rows dataSet
  let f = cols dataSet
  let y = is target $ MData.flatten $ takeColumns 1 dataSet
  let x = Matrix.col (replicate m 1.0) ||| ((dropColumns 1 dataSet) <> scale)
  let cvSize = 10
  let trainingX = dropRows cvSize x
  let cvX = takeRows cvSize x
  let trainingY = MData.fromList $ drop cvSize $ MData.toList y
  let cvY = MData.fromList $ take cvSize $ MData.toList y
  let theta = initialTheta f
  let (finalTheta, path) =
        minimizeVD
          SteepestDescent -- VectorBFGS2
          10e-6
          50
          10e-6
          0.1
          (costFn x y)
          (gradFn x y)
          theta
  pure $
    Result
      cvX
      cvY
      finalTheta
      (hypothesis cvX finalTheta)
      (costFn cvX cvY finalTheta)

sigmoid
  :: Floating a
  => a -> a
sigmoid term = 1.0 / (1.0 + exp (-term))

hypothesis :: Matrix Double -> Vector Double -> Vector Double
hypothesis x theta = sigmoid (x #> theta)

costFn :: Matrix Double -> Vector Double -> Vector Double -> Double
costFn x y theta = traceShowId $ all
  where
    all :: Double
    all = vsum (leftTerm - rightTerm) / fromIntegral (rows x)
    h = hypothesis x theta
    leftTerm = (-y) * log h
    rightTerm = (1.0 - y) * log (1.0 - h)

gradFn :: Matrix Double -> Vector Double -> Vector Double -> Vector Double
gradFn x y theta = (tr x #> (hypothesis x theta - y)) / fromIntegral (rows x)

vsum :: Vector Double -> Double
vsum = sum . MData.toList


#!/usr/bin/env stack
-- stack runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main where

import           Codec.Picture.Types
import           Control.Lens                  ()
import           Control.Monad
import           Data.Function
import           Data.List
import           Data.List.NonEmpty            (NonEmpty)
import qualified Data.List.NonEmpty            as NE
import           Data.Maybe
import           Data.Ratio
import qualified Data.Text                     as T
import           Data.Tuple
import qualified Data.Vector                   as V
import           Debug.Trace
import           Linear.Matrix                 hiding (trace)
import           Linear.Metric
import           Linear.V2
import           Linear.V3
import           Linear.Vector
import           Numeric.LinearAlgebra         hiding (polar, scale, (<>))
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix hiding (polar, scale, (<>))
import           Reanimate
import           Reanimate.Animation
import           Reanimate.Math.Balloon
import           Reanimate.Math.Common
import           Reanimate.Math.Triangulate
import           Reanimate.Math.Polygon
import           Reanimate.Math.EarClip
import           Reanimate.Math.SSSP
import           Reanimate.Math.Render
import           Reanimate.Math.Visibility
import           Reanimate.Math.Compatible
import           Reanimate.Morph.Common
import           Reanimate.Morph.Linear
import           Reanimate.Morph.LeastDifference
import           Reanimate.PolyShape           (svgToPolygons)
import           Reanimate.Debug
import qualified Reanimate.Math.Compatible as Compat

extraPoints = 0

-- p1 = pScale 2 $ pAtCenter $ pAddPoints extraPoints (pSetOffset shape13 0)
-- p1 = pSetOffset (addPoints 2 shape13) 0
-- p2 = pScale 0.5 $ centerPolygon shape20
-- p1 = centerPolygon shape2
-- p2 = pScale 2 $ pAtCenter $ pAddPoints extraPoints (pSetOffset shape14 0 )

-- p1 = pAddPoints extraPoints (pAtCenter $ unsafeSVGToPolygon 0.1 $ scale 6 $ latex "S")
-- p2 = pAddPoints extraPoints (pAtCenter $ unsafeSVGToPolygon 0.1 $ scale 6 $ latex "C")
p1 = pAddPoints extraPoints (pAtCenter $ unsafeSVGToPolygon 0.1 $ scale 6 $ latex "X")
p2 = pAddPoints extraPoints (pAtCenter $ unsafeSVGToPolygon 0.1 $ scale 6 $ latex "I")

p1_ = castPolygon p1
p2_ = castPolygon p2
polys = triangulate_ p1_ p2_
p1_circ = circumference (map fst polys') p1_
p2_circ = circumference (map snd polys') p2_

polys' = alignPolygons polys p1_ p2_
polys'' = compatTriagPairs polys'

-- (p1_circ, p2_circ) = closestLinearCorrespondence p1_circ' p2_circ'

-- (p1s, p2s) = unzip $ triangulate p1 p2

main :: IO ()
-- main = reanimate $ playTraces $ seq (last $ triangulate [] p1 p2) ()
main = reanimate $ scene $ do
  bg <- newSpriteSVG $ mkBackground "black"
  spriteZ bg (-1)

  -- let leastDiffTrig = triangulate p1 p2
  let lst = polys'
  -- let lst = take 10 $ uncurry Compat.compatiblyTriangulateP (polys'!!3) -- polys''
  -- let lst = polys''
  forM_ (zip [0..] lst {-([polys''!!2]++p_snd)-}) $ \(n,(l, r)) -> do
    let c = promotePixel $ turbo (n/fromIntegral (length lst-1))
    newSpriteSVG_ $
      translate (-2) 0 $ withFillColorPixel c $ mkGroup
      [ polygonShape (castPolygon l)
      -- , polygonNumDots (castPolygon l)
      ]
    newSpriteSVG_ $
      translate 2 0 $ withFillColorPixel c $ mkGroup
      [ polygonShape (castPolygon r)
      -- , polygonNumDots (castPolygon r)
      ]
    nums <- newSpriteSVG $ mkGroup
      [ translate (-2) 0 $ polygonNumDots (castPolygon l)
      , translate 2 0 $ polygonNumDots (castPolygon r)
      ]
    wait (1/60)
    destroySprite nums

  -- fork $ play $ staticFrame (1/60) $ mkGroup
  --   [ translate 2 0 $ mkGroup
  --     [ withFillColor "grey" $ polygonShape p2_circ
  --     , polygonNumDots p2_circ
  --     ]
  --   , translate 6 0 $ mkGroup
  --     [ withFillColor "grey" $ polygonShape p2
  --     , polygonNumDots p2
  --     ]
  --   ]
  -- play $ staticFrame (1/60) $ mkGroup
  --   [ translate (-2) 0 $ mkGroup
  --     [ withFillColor "grey" $ polygonShape p1_circ
  --     , polygonNumDots p1_circ
  --     ]
  --   , translate (-6) 0 $ mkGroup
  --     [ withFillColor "grey" $ polygonShape p1
  --     , polygonNumDots p1
  --     ]
  --   ]

showP p = mkGroup
  [ withFillColor "grey" $ polygonShape p
  , polygonNumDots p ]


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE TypeApplications    #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

module Neldermead (
  Simplex(..)
  , haftkaGurdalSimplex
  ) where

import Data.Ord (comparing)
import Data.Function (on)
import Data.Bool (bool)
import Data.Maybe (fromMaybe)
import Data.Monoid (Sum)
import qualified Data.Foldable as FLD
import qualified Data.Proxy as PXY
import Data.Coerce

import Control.Arrow ((&&&))
import Control.Monad.State.Strict

import qualified Numeric.Interval.NonEmpty as IVL

import qualified Data.Finite as F
import Data.Singletons (Sing, sing, fromSing)
import Data.Singletons.TypeLits

import qualified Data.Vector.Sized as V
import qualified Data.Vector.Storable.Sized as SV
--import qualified Data.Vector.Mutable.Sized as MV
import Numeric.LinearAlgebra.Static
import Numeric.LinearAlgebra.Static.Vector
import GHC.TypeNats

import qualified Control.Lens as L
import qualified Control.Lens.Iso as LI
import qualified Control.Lens.Traversal as LT

--TODO in future - mark Vertex with the objective function
--in any constructor, must take objective and then mark it

--TODO in future, allow more flexibility in objective?
--TODO better field names or allow overloaded
--TODO lenses
data NMVertex n = NMVertex { getX :: !(R n)
                           , getFx :: ℝ
                           }

newtype NMSimplex n = NMSimplex { fOrderedVertices :: V.Vector (n + 1) (NMVertex n) }

newtype Simplex n = Simplex { xA :: L n (n + 1) }
    deriving (Show)

newtype BoxConstraints n = BoxConstraints { unBoxConstraints :: V.Vector n (IVL.Interval ℝ) }

data NMScalars = NMScalars {
  reflectWeight :: ℝ
  , expandWeight :: ℝ
  , contractWeight :: ℝ
  , shrinkWeight :: ℝ
}

--TODO convergeTest should be monadic to allow writing details
--
data NMEnv n = MEnv {
  scalars :: NMScalars
  , simplexConvergeTest :: NMSimplex n -> Bool
}

newtype OnBounds = OnBounds { unOnBounds :: Bool }

makeNMVertex :: KnownNat n => (R n -> ℝ) -> R n -> NMVertex n
makeNMVertex f x = NMVertex { getX = x, getFx = f x }

--TODO actually need to change case 2 to keep if better than second worst (not best)
nextNMSimplex :: forall n. (KnownNat n, 1 <= n)
  => NMScalars -> (R n -> ℝ) -> NMSimplex n -> NMSimplex n
nextNMSimplex env f (NMSimplex vtxs)
  | comparing getFx vtxReflect vtxl == LT = undefined --expand case
  | comparing getFx vtxReflect vtxl2 == LT = undefined -- just keep swap out reflect and re-order
  | otherwise = undefined --shrink case
  where
    vtxl = V.head vtxs
    vtxl2 = V.index vtxs $ F.natToFinite (PXY.Proxy @1)
    vtxh = V.last vtxs
    xc = centroid $ getX <$> FLD.toList vtxs
    vtxReflect = makeNMVertex f $ affineCombWith (getX vtxh) (reflectWeight env) xc

--TODO may need MonadReader and MonadState, and allow objective to be in any monad
rVecIso :: KnownNat n => LI.Iso' (R n) (SV.Vector n ℝ)
rVecIso = LI.iso rVec vecR

projectToBox :: KnownNat n => BoxConstraints n -> R n -> R n
projectToBox (BoxConstraints bc) =
  L.over rVecIso $ SV.imap $ IVL.clamp . V.index bc

-- Return bool is weather any dimension was clamped
projectToBoxState :: KnownNat n
  => BoxConstraints n -> R n -> State OnBounds (R n)
projectToBoxState (BoxConstraints bc) = L.mapMOf rVecIso
  $ SV.imapM $ \i x -> do
    let bci = V.index bc i
    modify $ coerce (|| x `IVL.notMember` bci)
    return $ IVL.clamp bci x

--awkward because cant deconstruct interval into inf sup
--and cant zip storable and not storable
rightAngledSimplexInBox :: forall n. KnownNat n
  => BoxConstraints n -> R n -> R n -> Simplex n
rightAngledSimplexInBox (BoxConstraints bc) step x0 = Simplex
  $ ((x0 `outer` 1) + diag step') ||| col x0
  where
    (step' :: R n) = vecR
      $ SV.generate
      $ \i -> let stepBounds = V.index bc i - IVL.singleton (SV.index vx0 i)
                  si = abs $ SV.index vstep i
              in clampStep ((IVL.inf &&& IVL.sup) stepBounds) si
    clampStep (stepLb, stepUb) stepi
      | stepi <= stepUb = stepi
      | -stepi >= stepLb = -stepi
      | -stepLb < stepUb = max 0 $ 0.5 * stepUb
      | otherwise = min 0 $ 0.5 * stepLb
    vstep = rVec step
    vx0 = rVec x0

haftkaGurdalSimplex :: forall n. KnownNat n => ℝ -> R n -> Simplex n
haftkaGurdalSimplex size x0 = Simplex . colsL
  $ V.generate $ \j -> x0 L.& rVecIso L.%~ (SV.imap $ \i x -> x + step i j)
  where
    step i j
      | toInteger j == toInteger n' = 0
      | toInteger j == toInteger i = q + (c * l)
      | otherwise = q
    q = c * (sqrt (l + 1) - 1)
    c = size / (l * sqrt 2)
    l = fromIntegral n'
    n' = fromSing $ (sing :: Sing n)

centroid :: forall n t. (KnownNat n, Foldable t) => t (R n) -> R n
centroid = (/l) . FLD.foldl' (+) 0
  where
    l = fromIntegral . fromSing $ (sing :: Sing n)

centroidL :: forall m n. (KnownNat m, KnownNat n) => L m n -> R m
centroidL = (#> (1.0 / l))
  where
    l = fromIntegral . fromSing $ (sing :: Sing m)

realWithinRelTol :: ℝ -> ℝ -> ℝ -> Bool
realWithinRelTol relTol x y = (<= 0.5 * (x `absAdd` y)) . abs $ x - y
  where absAdd = (+) `on` abs

realWithinTols :: Maybe ℝ -> Maybe ℝ -> ℝ -> ℝ -> Bool
realWithinTols relTol absTol x y = relTolCheck && absTolCheck
  where
    relTolCheck = fromMaybe True
      $ (\t -> realWithinRelTol t x y) <$> relTol
    absTolCheck = fromMaybe True $ (<= abs (x - y)) <$> absTol

simplexRadiusTo :: (KnownNat n, Foldable t) => t (R n) -> R n -> R n
simplexRadiusTo xs c = vecR
  $ SV.generate
  $ \i -> FLD.foldl' (\acc -> (max `on` abs) acc . distci i . rVec) 0 xs
  where
    distci i v = SV.index v i - SV.index vc i
    vc = rVec c

--center + max step to given vertex for each dimension
simplexRadiusToL :: KnownNat n => Simplex n -> R n -> R n
simplexRadiusToL (Simplex xs) c = vecR $ SV.generate (V.index v)
  where
    v = (SV.maximumBy (comparing abs) . rVec) <$> lRows xsDiff -- Vector, need to convert to (R n)
    xsDiff = xs - (c `outer` 1) -- vector not promoted to matrix

affineCombWith :: KnownNat n => R n -> ℝ -> R n -> R n
affineCombWith x cy y = x + rnScale cy (y - x)

--put in helper class, maybe use coerce or something
rnScale :: KnownNat n => ℝ -> R n -> R n
rnScale c = L.over rVecIso $ SV.map (c *)

--TODO write version using mean
--toRows, map mean, list to vec
centroidL2 :: forall m n. Sing m -> Sing n -> L m n -> R m
centroidL2 (sm@SNat) (sn@SNat) = (#> (1.0 / l))
  where
    l = fromIntegral . fromSing $ (sing :: Sing m)


module EvalSpec (spec) where

import           Control.Monad.Error
import           Data.Complex
import           Data.Either
import           Data.IORef
import           Data.Ratio
import           Env
import           Eval
import           Read
import           System.IO.Unsafe
import           Test.Hspec
import           Types

-- same as Main.hs
primitiveBindings :: IO Env
primitiveBindings = nullEnv >>= flip bindVars (map makePrimitiveFunc primitives)
  where makePrimitiveFunc (var, func) = (var, PrimitiveFunc func)

runIOThrowsToLispVal :: IOThrowsError LispVal -> IO LispVal
runIOThrowsToLispVal action = fmap extractValue (runErrorT action)

exprToVal :: String -> IOThrowsError LispVal
exprToVal expr = liftThrows $ readExpr expr

evalExpr :: Env -> String -> IO LispVal
evalExpr env expr = runIOThrowsToLispVal $ exprToVal expr >>= eval env

re expr = unsafePerformIO $ do
  env <- primitiveBindings
  evalExpr env expr

le expr = unsafePerformIO $ do
  env <- primitiveBindings
  evalString env expr

spec :: Spec
spec = do
  describe "eval" $ do
    it "evaluates functions" $ do
      re "(+ 1 2 3)" `shouldBe` Number 6
      re "(- 3 1)" `shouldBe` Number 2
      re "(- 10 1 2)" `shouldBe` Number 7
      re "(* 1 2 3 4 5)" `shouldBe` Number 120
      re "(/ 23 7)" `shouldBe` Number 3
      re "(/ 51 7 3)" `shouldBe` Number 2
      re "(mod 7 3)" `shouldBe` Number 1
      re "(quotient 7 3)" `shouldBe` Number 2
      re "(remainder 7 3)" `shouldBe` Number 1
      re "(symbol? 'a)" `shouldBe` Bool True
      re "(symbol? 1)" `shouldBe` Bool False
      re "(string? \"foo\")" `shouldBe` Bool True
      re "(string? 1)" `shouldBe` Bool False
      re "(number? 1)" `shouldBe` Bool True
      re "(number? 'a)" `shouldBe` Bool False
      re "(boolean? #t)" `shouldBe` Bool True
      re "(boolean? 'f)" `shouldBe` Bool False

      re "(= 1 1)" `shouldBe` Bool True
      re "(= 1 2)" `shouldBe` Bool False
      re "(> 1 0)" `shouldBe` Bool True
      re "(> 1 2)" `shouldBe` Bool False
      re "(< 1 0)" `shouldBe` Bool False
      re "(< 1 2)" `shouldBe` Bool True
      re "(/= 1 1)" `shouldBe` Bool False
      re "(/= 1 2)" `shouldBe` Bool True
      re "(>= 1 2)" `shouldBe` Bool False
      re "(>= 1 1)" `shouldBe` Bool True
      re "(<= 1 0)" `shouldBe` Bool False
      re "(<= 1 1)" `shouldBe` Bool True
      re "(&& #t #t)" `shouldBe` Bool True
      re "(&& #t #f)" `shouldBe` Bool False
      re "(&& #f #f)" `shouldBe` Bool False
      re "(|| #t #t)" `shouldBe` Bool True
      re "(|| #t #f)" `shouldBe` Bool True
      re "(|| #f #f)" `shouldBe` Bool False

    it "string?" $ do
      re "(string=? \"a\" \"a\")" `shouldBe` Bool True
      re "(string=? \"a\" \"b\")" `shouldBe` Bool False
      re "(string<? \"a\" \"a\")" `shouldBe` Bool False
      re "(string<? \"a\" \"b\")" `shouldBe` Bool True
      re "(string>? \"a\" \"a\")" `shouldBe` Bool False
      re "(string>? \"c\" \"b\")" `shouldBe` Bool True
      re "(string<=? \"b\" \"a\")" `shouldBe` Bool False
      re "(string<=? \"b\" \"b\")" `shouldBe` Bool True
      re "(string>=? \"a\" \"b\")" `shouldBe` Bool False
      re "(string>=? \"b\" \"b\")" `shouldBe` Bool True

    it "if" $ do
      re "(if (> 2 3) \"foo\" \"bar\")" `shouldBe` String "bar"
      re "(if (= 3 3) (+ 2 3 (- 5 1)) \"unequal\")" `shouldBe` Number 9

    it "car, cdr" $ do
      re "(cdr '(a simple test))" `shouldBe` List [Atom "simple", Atom "test"]
      re "(car (cdr '(a simple test)))" `shouldBe` Atom "simple"
      re "(car '((this is) a test))" `shouldBe` List [Atom "this", Atom "is"]
      re "(cons '(this is) 'test)" `shouldBe` DottedList [List [Atom "this", Atom "is"]] (Atom "test")
      re "(cons '(this is) '())" `shouldBe` List [List [Atom "this", Atom "is"]]

    it "eqv?" $ do
      re "(eqv? 1 3)" `shouldBe` Bool False
      re "(eqv? 3 3)" `shouldBe` Bool True
      re "(eqv? 'atom 'atom)" `shouldBe` Bool True

    it "equal?" $ do
      re "(equal? 123 123)" `shouldBe` Bool True
      re "(equal? 123 \"123\")" `shouldBe` Bool True
      re "(equal? 123 \"foo\")" `shouldBe` Bool False

    it "throws" $ do
      le "(symbol? 1 2)" `shouldStartWith` "Expected 1 args;"
      le "(string? 1 2)" `shouldStartWith` "Expected 1 args;"
      le "(number? 1 2)" `shouldStartWith` "Expected 1 args;"
      le "(boolean? 1 2)" `shouldStartWith` "Expected 1 args;"
      le "(+ 2 \"two\")" `shouldStartWith` "Invalid type:"
      le "(+ 2)" `shouldStartWith` "Expected 2 args;"
      le "(what? 2)" `shouldBe` "Getting an unbound variable: : what?"
      le "(if 1 2 3)" `shouldStartWith` "Invalid type:"

  describe "env" $ do
    it "defines functions" $ do
      env <- primitiveBindings

      evalExpr env "(define (f x y) (+ x y))"
      unsafePerformIO (evalExpr env "(f 1 2)") `shouldBe` Number 3

      evalExpr env "(define (factorial x) (if (= x 1) 1 (* x (factorial (- x 1)))))"
      unsafePerformIO (evalExpr env "(factorial 10)") `shouldBe` Number 3628800

      evalExpr env "(define (counter inc) (lambda (x) (set! inc (+ x inc)) inc))"
      evalExpr env "(define my-count (counter 5))"
      unsafePerformIO (evalExpr env "(my-count 3)") `shouldBe` Number 8
      unsafePerformIO (evalExpr env "(my-count 6)") `shouldBe` Number 14
      unsafePerformIO (evalExpr env "(my-count 5)") `shouldBe` Number 19

    it "fails if function is not defined" $
      le "(f 1 2)" `shouldBe` "Getting an unbound variable: : f"

    it "defines functions" $ do
      env <- primitiveBindings
      evalExpr env "(define (f x y) (+ x y))"
      unsafePerformIO (evalExpr env "(f 1 2)") `shouldBe` Number 3

    it "fails if function is not defined" $
      le "(f 1 2)" `shouldBe` "Getting an unbound variable: : f"


module Algebra
    (Adjoin (..)
    ,FreeGroup (..)
    ,Group (..)
    ,Z2
    ,plugIn
    ,Algebra
    ,represent
    ,cutoff
    ,module N
    ) where

import Algebra.Adjoin
import Algebra.Group
import Algebra.FreeGroup
import Algebra.Z2
import Numeric.LinearAlgebra as N

type Algebra = Adjoin Z2 FreeGroup
{--type Algegra' = [(Z2,Vector Z,FreeGroup)]

plush :: (Z2,Vector Z,FreeGroup) -> Algebra' -> Algebra'
plush e [] = [e]
plush e1@(z1,v1,g1) (e2@(z2,v2,g2):ls) = if v1 == v2 && g1 == g2 then (if z1+z2 == 0 then (0,N.fromList $ map (\_ -> 0) $ N.tolist v1,mempty) else (z1+z2,v1,g1)):ls else e2:(plush e1 ls)

instance Num Algebra' where
    l1 + l2 = map (\e -> plush e l2) l1
    l1 * l2 = sum $ map (\(z1,v1,g1) -> map (\(z2,v2,g2) -> (z1*z2,v1+v2,g1<>g2)) l2) l1
    abs = map (\(z,v,g) -> (abs z,v,g))
    signum = map (\(z,v,g) -> (signum z,v,g))
--    fromInteger x = [(fromInteger x,
-}
cutoff :: R
cutoff = 0.0000000001

plugIn :: (Eq f,Fractional f) => (Char -> Adjoin f FreeGroup) -> Adjoin f FreeGroup -> Adjoin f FreeGroup
plugIn f = adjPlugIn (groupPlugIn f)

represent :: [Char] -> Algebra -> Maybe [Vector Z]
represent [] _ = Nothing
represent elems expr = do
                        { chain <- pairPlusChain expr
                        ; vectors <- mapM (\(_,g) -> toPseudoVector g) chain
                        ; return $ map (\l -> (length elems) |> [toEnum $ maybe 0 id $ lookup (elems !! i) l | i <- [0..]]) vectors
                        }


{-# LANGUAGE FlexibleInstances #-}

module School.Types.FloatEq
( FloatEq(..)
, compareDouble
) where

import Numeric.LinearAlgebra (Element)
import Numeric.LinearAlgebra.Data (Matrix, Vector, cols, toList, toLists)
import School.Train.TrainState (HandlerStore(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.UnitParams (UnitParams(..))

class FloatEq a where
  (~=) :: a -> a -> Bool
  (~/) :: a -> a -> Bool
  x ~= y = not (x ~/ y)
  x ~/ y = not (x ~= y)

instance FloatEq Int where
  (~=) = (==)

instance (FloatEq a) => FloatEq (Maybe a) where
  Nothing ~= Nothing = True
  (Just x) ~= (Just y) = x ~= y
  _ ~= _ = False

compareDouble :: Double -> Double -> Double -> Bool
compareDouble prec d1 d2 = if abs d1 < prec
  then abs d2 < prec
  else abs ((d2 - d1) / d1) < prec

doubleEq :: Double -> Double -> Bool
doubleEq = compareDouble 5e-11

instance FloatEq Double where
  (~=) = doubleEq

instance (FloatEq a) => FloatEq [a] where
  l ~= m = (length l == length m)
        && and (zipWith (~=) l m)

mToList :: (Element a) => Matrix a -> [a]
mToList = concat . toLists

instance (Element a, FloatEq a) => FloatEq (Matrix a) where
  m1 ~= m2 = (cols m1 == cols m2)
          && (mToList m1 ~= mToList m2)

instance (Element a, FloatEq a) => FloatEq (Vector a) where
  v1 ~= v2 = toList v1 ~= toList v2

instance (FloatEq b) => FloatEq (Either a b) where
  (Left _) ~= (Left _) = True
  (Right res1) ~= (Right res2) = res1 ~= res2
  _ ~= _ = False

instance (Element a, FloatEq a) => FloatEq (UnitActivation a) where
  (BatchActivation m1) ~= (BatchActivation m2) = m1 ~= m2
  _ ~= _ = False

instance (Element a, FloatEq a) => FloatEq (UnitGradient a) where
  (BatchGradient g1) ~= (BatchGradient g2) = g1 ~= g2
  _ ~= _ = False

instance (Element a, FloatEq a) => FloatEq (UnitParams a) where
  AffineParams { affineBias = b1, affineWeights = w1 } ~= AffineParams { affineBias = b2, affineWeights = w2 } =
    b1 ~= b2 && w1 ~= w2
  EmptyParams ~= EmptyParams = True
  _ ~= _ = False

instance (FloatEq a) => FloatEq (HandlerStore a) where
  (CostList l1) ~= (CostList l2) = l1 ~= l2
  NoStore ~= NoStore = True
  _ ~= _ = False


{-# language GADTs #-}
{-# language ConstraintKinds #-}
{-# language FlexibleContexts #-}

{-# language ScopedTypeVariables #-}
{-# LANGUAGE QuasiQuotes #-}
{-# language TypeApplications #-}

module FFT where

import Feldspar
import Feldspar.Software hiding (Arr)
import Feldspar.Software.Verify
import Feldspar.Software.Compile

import Feldspar.Array.Vector
import Feldspar.Array.Buffered

import Data.Bits (Bits)
import Data.Complex (Complex)

import Data.Selection
import Data.Default.Class

import Control.Monad

-- language-c-quote
import Language.C.Quote.GCC
import qualified Language.C.Syntax as C

-- imperative-edsl
import qualified Language.Embedded.Backend.C  as Imp

import Prelude hiding ((==), (/=), (>=), length)

--------------------------------------------------------------------------------
-- * FFT
--------------------------------------------------------------------------------

type SRef    a = Reference Software a
type SArr    a = Array  Software a
type SIArr   a = IArray Software a
type SStore  a = Store  Software a
type SSyntax a = Syntax SExp a

-- | Collection of constraints for immutable arrays.
type Immutable arr a = (
    Manifestable Software arr a
  , Finite   SExp arr
  , Indexed  SExp arr, ArrElem arr ~ a
  , Slicable SExp arr
  )

--------------------------------------------------------------------------------
-- ** Helper functions.

-- | Creates a value of type @a@ where the lowest @n@ bits are ones.
oneBits :: (Bits a, Num (SExp a), SType' a)
  => SExp Int32 -> SExp a
oneBits n = complement (ones .<<. n)

-- | Check if the @i@'th bit of @a@ is a one.
testBit :: (Integral a, Bits a, Num (SExp a), SType' a) =>
  SExp a -> SExp Index -> SExp Bool
testBit a i = i2b (a .&. (1 .<<. i2n i))

-- | Flip the @i@'th bit of @a@.
flipBit :: (Bits a, Num (SExp a), SType' a) =>
  SExp a -> SExp Index -> SExp a
flipBit a i = a `xor` (1 .<<. (i2n i))
--  SExp Index -> SExp a -> SExp a
--flipBit i a = a `xor` (1 .<<. (i2n i))

-- |
zeroBit :: (Bits a, Num (SExp a), SType' a) =>
  SExp Index -> SExp a -> SExp a
zeroBit i a = a + (a .&. (ones .<<. (i2n i)))

-- | Zeroes all but the lowest @i@ bits in @a@.
leastBits :: (Bits a, Num (SExp a), SType' a) =>
  SExp Int32 -> SExp a -> SExp a
leastBits i a = a .&. oneBits i

-- | Two to the power of @n@, i.e. @2^n@.
pow2 :: (Num (SExp a), Bits a, SType' a)
  => SExp Index -> SExp a
pow2  n = 1 .<<. i2n n
twoTo n = pow2 n

--------------------------------------------------------------------------------
-- ** Riffle network.

-- | Riffle indices.
rotBit :: SExp Index -> SExp Index -> SExp Index
rotBit k i = lefts .|. rights
  where
    k'     = i2n k
    ir     = i .>>. 1
    rights = ir .&. oneBits k'
    lefts  = (((ir .>>. k') .<<. 1) .|. (i .&. 1)) .<<. k'

-- | Riffle network for a pully array.
riffle :: (Immutable arr a, SSyntax a)
  => SExp Index -> arr -> SPull a
riffle k arr = Pull (length arr) $ \i -> arr ! rotBit k i

-- |
revBit :: (Immutable arr a, SSyntax a)
  => SStore a -> SExp Length -> arr -> Software (SManifest a)
revBit st n vec = loopStore st 1 1 (n-1) (step) vec
  where
    step i arr = return $ riffle i arr
    --step :: (Immutable arr a, SSyntax a) => SExp Index -> arr -> Software (SPush a)
--    step i arr = return $ unroll 2 $ riffle i arr
--    step k arr = return $ pairwise @Software (\i -> (i, rotBit k i)) $ riffle k arr

--------------------------------------------------------------------------------
-- ** Twiddle factors.

-- |
tw :: (Floating a, SType' a, SType' (Complex a))
  => SExp Index -> SExp Index -> SExp (Complex a)
tw n k = polar 1 (-2 * pi * i2n k / i2n n)

-- |
twids
  :: ( Immutable ts  (SExp (Complex a))
     , Immutable vec (SExp (Complex a))
     , SType' a
     , SType' (Complex a)
     , RealFloat a
     )
  => ts
  -> SExp Index
  -> SExp Index
  -> SExp Length
  -> vec
  -> SPull (SExp (Complex a))
twids ts n k l vec = Pull l $ \i ->
  let
    j = (leastBits (i2n k) i) .<<. (n'-1-k')
  in
    (testBit i k) ? ((ts!j) * (vec!i)) $ (vec!i)
  where
    n' = i2n n
    k' = i2n k

--------------------------------------------------------------------------------
-- ** Butterfly.

-- | Butterfly network.
bfly
  :: ( Immutable vec (SExp (Complex a))
     , RealFloat a
     , SType' a
     , SType' (Complex a)
     )
  => SExp Index -> vec -> SPull (SExp (Complex a))
bfly k as = Pull (length as) $ \i ->
    let a = as ! i
        b = as ! flipBit i k
    in  (testBit i k) ? (b-a) $ (a+b)

--------------------------------------------------------------------------------
-- ** FFT Core.

-- | Core of the FFT
fftCore
  :: ( Immutable ts  (SExp (Complex a))
     , Immutable vec (SExp (Complex a))
     , RealFloat a
     , SType' a
     , SType' (Complex a)
     )
  => SStore (SExp (Complex a))
  -> ts
  -> SExp Length
  -> vec
  -> Software (SManifest (SExp (Complex a)))
fftCore st ts n vec =
  let
    step k = return . twids ts n k (length vec) . bfly k
--    step k = return . pairwise @Software (\i -> (i, flipBit i k)) . twids ts n k (length vec) . bfly k
  in
    do arr <- loopStore st ((i2n n :: SExp Int32)-1) (-1) 0 (step . i2n) vec
--       revBit st n arr
       return arr

-- | Radix-2 Decimation-In-Frequency Fast Fourier Transformation of the given
-- complex vector. The given vector must be power-of-two sized, (for example 2,
-- 4, 8, 16, 32, etc.) The output is non-normalized.
fft
  :: ( Immutable vec (SExp (Complex a))
     , RealFloat a
     , SType' a
     , SType' (Complex a)
     )
  => SExp Length
  -> SStore (SExp (Complex a))
  -> vec
  -> Software (SManifest (SExp (Complex a)))
fft n st vec =
  do n  <- shareM (ilog2 (length vec))
     ts <- manifestFresh $ Pull (pow2 (n-1)) (tw (pow2 n))
     fftCore st ts n vec

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

fftCore2 ::
  SStore (SExp (Complex Double)) ->
  SIArr (SExp (Complex Double)) ->
  SExp Length ->
  SPull (SExp (Complex Double)) ->
  Software (SIArr (SExp (Complex Double)))
fftCore2 store iarr exp pull = do
  let iarrm :: SManifest (SExp (Complex Double))
      iarrm = M iarr
  (M ibrr) <- fftCore store iarrm exp pull
  return ibrr

fft2 ::
  SExp Length ->
  SStore (SExp (Complex Double)) ->
  SPull (SExp (Complex Double)) ->
  Software (SIArr (SExp (Complex Double)))
fft2 n store pull = do
  (M ts) <- manifestFresh $ Pull (twoTo (n-1)) (tw (twoTo n))
  fftCore2 store ts n pull

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

example :: Software ()
example = do
  size :: SExp Length <- fget stdin
  n :: SExp Length <- fget stdin
  assert (1 `shiftL` n == size) "2^n == size"
  assert (size >= 2) "not too small"
  st   :: SStore (SExp (Complex Double)) <- newInPlaceStore size
  arr  :: SArr   (SExp (Complex Double)) <- newArr size
  iarr :: SIArr  (SExp (Complex Double)) <- freezeArr arr
  fft n st iarr
  assert (value False) "oh no"
  return ()

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

bad :: Software ()
bad = do
  size :: SExp Length <- fget stdin
  fprintf stdout "%d" (share size id)
  return ()

--------------------------------------------------------------------------------
-- FFT Bench. Copy of https://github.com/Feldspar/raw-feldspar/blob/master/examples/FFT_bench.hs
--------------------------------------------------------------------------------

printTime_def = [cedecl|
void printTime(typename clock_t start, typename clock_t end)
{
  printf("CPU time (sec): %f\n", (double)(end-start) / CLOCKS_PER_SEC);
}
|]

sizeOf_double_complex :: SExp Length
sizeOf_double_complex = 16

-- | Measure the time for 100 runs of 'fftCore' (excluding initialization) for
-- arrays of the given size
benchmark :: SExp Length -> Software ()
benchmark n = do
  addInclude "<stdio.h>"
  addInclude "<string.h>"
  addInclude "<time.h>"

  addDefinition printTime_def

  start <- newObject "clock_t" False
  end   <- newObject "clock_t" False

  st :: SStore (SExp (Complex Double)) <- newStore n
  inp <- unsafeFreezeStore n st
  callProc "memset"
      [ iarrArg (manifest inp)
      , valArg (0 :: SExp Index)
      , valArg (n*sizeOf_double_complex)
      ]

  n  <- shareM (ilog2 (length inp))
  ts <- manifestFresh $ Pull (twoTo (n-1)) (tw (twoTo n))
  -- Change `manifestFresh` to `return` to avoid pre-computing twiddle factors

  callProcAssign start "clock" []

  for 0 1 99 $ \(_ :: SExp Index) ->
    void $ fftCore st ts n inp

  callProcAssign end "clock" []
  callProc "printTime" [objArg start, objArg end]

runBenchmark n = runCompiled'
    (def :: CompilerOpts) --{compilerAssertions = select []}
    -- Note: important to turn off assertions when running the benchmarks
    --       (in old Raw-Feldspar that is, Co-Feldspar gets rid of them)
    def {Imp.externalFlagsPre = ["-O3"], Imp.externalFlagsPost = ["-lm"]}
    (benchmark n)

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

dummy :: Software ()
dummy = do
  size :: SExp Length <- fget stdin
  n :: SExp Length <- fget stdin
  assert (1 `shiftL` n == size) "2^n == size"
  assert (size >= 2) "not too small"
  st   :: SStore (SExp (Complex Double)) <- newInPlaceStore size
  arr  :: SArr   (SExp (Complex Double)) <- newArr size
  iarr :: SIArr  (SExp (Complex Double)) <- freezeArr arr
  fft n st iarr
  fft n st iarr
  assert (value False) "oh no"
  return ()

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


{-# LANGUAGE CPP #-}
--------------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.HMatrix
Copyright   :  (c) Alberto Ruiz 2006-14
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

compatibility with previous version, to be removed

-}
--------------------------------------------------------------------------------

module Numeric.LinearAlgebra.HMatrix (
    module Numeric.LinearAlgebra,
    (¦),(——),ℝ,ℂ,(<·>),app,mul, cholSH, mbCholSH, eigSH', eigenvaluesSH', geigSH'
) where

import Numeric.LinearAlgebra
import Internal.Util
import Internal.Algorithms(cholSH, mbCholSH, eigSH', eigenvaluesSH', geigSH')
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

infixr 8 <·>
(<·>) :: Numeric t => Vector t -> Vector t -> t
(<·>) = dot

app :: Numeric t => Matrix t -> Vector t -> Vector t
app m v = m #> v

mul :: Numeric t => Matrix t -> Matrix t -> Matrix t
mul a b = a <> b


module Main where

import System.Environment (getArgs)
import qualified Statistics.Test.Types as Stats
import Jakway.Blackjack.Random.Checks
import Jakway.Blackjack.Random.Options
import Jakway.Blackjack.Util (die)     -- ^ die is implemented here for compatibility with older GHC versions
import Jakway.Blackjack.AI

main :: IO ()
main = do
    args <- getArgs
    conf <- getConfig args

    -- | TODO: make the RNG max range an option
    let rngMaxRange = toInteger $ 2 ^ 25

    let distrib = distribution conf
        pval = pvalue conf
        n = sampleSize conf
        r = rngSampleSize conf
        dealerAI = BasicDealer
        playerAIs = [BasicPlayer]

    if (n <= 104) then die "Samples must be >104." else return ()

    case distrib of EvenDistribution -> testDeckEvenDistribution pval n dealerAI playerAIs >>= printResult
                    RNGDistribution -> testRNGDistribution pval n r rngMaxRange dealerAI playerAIs >>= printResult
                    _ -> die "Fatal error, unknown test!"

printResult :: Stats.TestResult -> IO ()
printResult (Stats.Significant) = putStrLn $ "The test result is: Significant.  There is evidence to conclude that the deck IS NOT an unbiased source of randomness."
printResult (Stats.NotSignificant) = putStrLn $ "The test result is: Not Significant.  There is insufficient evidence to conclude that the deck is not an unbiased source of randomness."


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE Strict #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

import qualified Data.Vector as V
import qualified NumHask.Array.Dynamic as D
import qualified NumHask.Array.Fixed as F
import qualified NumHask.Array.HMatrix as H
import qualified Numeric.LinearAlgebra as HMatrix
import Perf
import NumHask.Prelude as P hiding (option)
import Data.Text (unpack, Text)
import Options.Applicative
import Control.Monad

data RunType = RunMMult | RunDotSum | RunDotVector deriving (Eq, Show)

data StatType = StatAverage | StatMedian | StatBest deriving (Eq, Show)

data Options = Options
  { optionRuns :: Int,
    optionMSize :: Int,
    optionMMult :: Bool,
    optionDotSum :: Bool,
    optionDotVector :: Bool,
    optionStatType :: StatType
  } deriving (Eq, Show)

options :: Parser Options
options = Options <$>
  option auto (long "runs" <> short 'r' <> help "number of runs to perform") <*>
  option auto (long "matrix-size" <> short 'm' <> help "size of square matrix") <*>
  switch (long "include matrix-multiplication calc" <> short 'x') <*>
  switch (long "dot-sum" <> help "include dot sum (+) calc" <> short 'd') <*>
  switch (long "dot-vector" <> help "include dot vector calc" <> short 'v') <*>
  stat

opts :: ParserInfo Options
opts = info (options <**> helper)
  (fullDesc <> progDesc "benchmark testing" <> header "A performance benchmark for numhask")

stat :: Parser StatType
stat =
  flag' StatBest (long "best" <> help "report upper decile") <|>
  flag' StatMedian (long "median" <> help "report median") <|>
  pure StatAverage

tickStat :: StatType -> [Cycle] -> Text
tickStat StatBest = tenth
tickStat StatMedian = median
tickStat StatAverage = average

main :: IO ()
main = do
  o <- execParser opts
  let !n = optionRuns o
  let s = optionStatType o
  _ <- warmup 100

  putStrLn "mmult 10x10"

  when (optionMMult o) $ do
    runHMatrix s n
    runNHHMatrix s n
    runF s n
    runD s n
    when (optionDotSum o) (runFDS s n)

  when (optionDotVector o) $ do
    runDotV s n
    runDotF s n
    runDotD s n

runHMatrix :: StatType -> Int -> IO ()
runHMatrix s n = do
  -- HMatrix 10x10 matrix multiplication
  let ticksH x = ticks n (x HMatrix.<>) x
  let !h10 = (10 HMatrix.>< 10) [1 :: HMatrix.R ..]
  th10 <- ticksH h10
  putStrLn $ unpack $ "hmatrix " <> tickStat s (fst th10)

runNHHMatrix :: StatType -> Int -> IO ()
runNHHMatrix s n = do
  -- NumHask.Array.HMatrix
  let ticksH' x = ticks n (x `H.mmult`) x
  let !h'10 = [1 .. 100] :: H.Array '[10, 10] Double
  th'10 <- ticksH' h'10
  putStrLn $ unpack $ "numhask-hmatrix " <> tickStat s (fst th'10)

runF :: StatType -> Int -> IO ()
runF s n = do
  -- NumHask.Array.Fixed
  let ticksF x = ticks n (x `F.mmult`) x
  let !f10 = [1 .. 100] :: F.Array '[10, 10] Double
  tf10 <- ticksF f10
  putStrLn $ unpack $ "Fixed " <> tickStat s (fst tf10)

runFDS :: StatType -> Int -> IO ()
runFDS s n = do
  -- NumHask.Array.Fixed.dot
  let ticksF x = ticks n (F.dot sum (+) x) x
  let !f10 = [1 .. 100] :: F.Array '[10, 10] Double
  tf10 <- ticksF f10
  putStrLn $ unpack $ "Fixed-dotsum " <> tickStat s (fst tf10)

runD :: StatType -> Int -> IO ()
runD s n = do
  -- NumHask.Array.Dynamic
  let !d10 = D.fromFlatList [10, 10] [1 .. 100] :: D.Array Double
  let ticksD x = ticks n (x `D.mmult`) x
  td10 <- ticksD d10
  putStrLn $ unpack $ "Dynamic " <> tickStat s (fst td10)

runDotV :: StatType -> Int -> IO ()
runDotV s n = do
  -- Vector dot
  let !vv = V.fromList [1 .. 100] :: V.Vector Double
  let ticksDotv x = ticks n (sum . V.zipWith (*) x) x
  tvd <- ticksDotv vv
  putStrLn $ unpack $ "Vector-dot " <> tickStat s (fst tvd)

runDotF :: StatType -> Int -> IO ()
runDotF s n = do
  -- Fixed Dot
  let !vf = P.fromList [1 .. 100] :: F.Array '[100] Double
  let ticksDotF x = ticks n (F.dot sum (+) x) x
  tfd <- ticksDotF vf
  putStrLn $ unpack $ "Fixed-dot " <> tickStat s (fst tfd)

runDotD :: StatType -> Int -> IO ()
runDotD s n = do
  -- Dynamic Dot
  let !vd = D.fromFlatList [100] [1 .. 100] :: D.Array Double
  let !vd2 = D.fromFlatList [100] [101 .. 200] :: D.Array Double
  let ticksDotD x = ticks n (D.dot sum (+) vd2) x
  tdd <- ticksDotD vd
  putStrLn $ unpack $ "Dynamic-dot " <> tickStat s (fst tdd)
  pure ()


{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}

import qualified Prelude

import qualified Data.Complex as Complex

import Feldspar.Run
import Feldspar.Data.Vector
import Feldspar.Data.Buffered

import qualified Test.QuickCheck as QC
import qualified Test.QuickCheck.Monadic as QC

import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck

import qualified Tut1_HelloWorld            as Tut1
import qualified Tut2_ExpressionsAndTypes   as Tut2
import qualified Tut3_Vectors               as Tut3
import qualified Tut4_MemoryManagement      as Tut4
import qualified Tut5_Matrices              as Tut5
import qualified Tut6_Testing               as Tut6
import qualified Tut7_ImperativeProgramming as Tut7
import qualified Tut8_SequentialVectors     as Tut8
import qualified Concurrent
import DFT
import FFT

almostEq a b
    =          Complex.magnitude d Prelude.< 1e-7
    Prelude.&& Complex.phase d     Prelude.< 1e-7
  where
    d = abs (a-b)

a ~= b = Prelude.and $ Prelude.zipWith almostEq a b

wrapStore :: (Syntax a, Finite (vec a), MonadComp m) =>
    (Store a -> vec a -> m b) -> vec a -> m b
wrapStore f v = do
    st <- newStore $ length v
    f st v

fftS u = wrapStore (flip fft u)  :: DManifest (Complex Double) -> _
ifftS  = wrapStore (flip ifft 1) :: DManifest (Complex Double) -> _

prop_fft_dft dft' fft' = QC.monadicIO $ do
    n   :: Int              <- QC.pick $ QC.choose (2,5)
    inp :: [Complex Double] <- QC.pick $ QC.vector (2 Prelude.^ n)
    outd <- QC.run $ dft' inp
    outf <- QC.run $ fft' inp
    QC.assert (outd ~= outf)

prop_inverse f fi = QC.monadicIO $ do
    n   :: Int              <- QC.pick $ QC.choose (2,5)
    inp :: [Complex Double] <- QC.pick $ QC.vector (2 Prelude.^ n)
    out1 <- QC.run $ f inp
    out2 <- QC.run $ fi out1
    QC.assert (inp ~= out2)

prop_fib fb1 fb2 = QC.monadicIO $ do
    n   <- QC.pick $ QC.choose (0,40)
    fs1 <- QC.run $ fb1 n
    fs2 <- QC.run $ fb2 n
    QC.assert (fs1 Prelude.== fs2)

main =
    marshalledM (return . dft)  $ \dft'  ->
    marshalledM (return . idft) $ \idft' ->
    marshalledM (fftS 1)        $ \fft1  ->
    marshalledM (fftS 2)        $ \fft2  ->
    marshalledM ifftS           $ \ifft' ->

    marshalled (return . Tut8.fibSeq) $ \fb1 ->
    marshalled (\n -> return $ Pull n Tut2.fib) $ \fb2 ->

      defaultMain $ testGroup "tests"
        [ testCase "Tut1"       Tut1.testAll
        , testCase "Tut2"       Tut2.testAll
        , testCase "Tut3"       Tut3.testAll
        , testCase "Tut4"       Tut4.testAll
        , testCase "Tut5"       Tut5.testAll
        , testCase "Tut6"       Tut6.testAll
        , testCase "Tut7"       Tut7.testAll
        , testCase "Tut8"       Tut8.testAll
        , testCase "Concurrent" Concurrent.testAll
        , testProperty "fft1_dft" $ prop_fft_dft dft' fft1
        , testProperty "fft2_dft" $ prop_fft_dft dft' fft2
        , testProperty "dft_idft" $ prop_inverse dft' idft'
        , testProperty "fft_ifft" $ prop_inverse fft1 ifft'
        , testProperty "fib"      $ prop_fib fb1 fb2
        ]

  where
    marshalledM = marshalled' def def {externalFlagsPost = ["-lm"]}


{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE RecordWildCards           #-}
{-# LANGUAGE ScopedTypeVariables       #-}

module Main where

import           Control.Concurrent (threadDelay)
import           Control.Concurrent.Async
import           Control.Concurrent.MVar
import           Control.Monad
import           Data.IORef
import           Data.Time.Clock.POSIX
import           Data.Time.Units
import           Data.Vector (Vector, fromList)
import qualified Data.Vector as V (length)
import           Data.Word (Word32)
import           Network.QDisc.Fair
import           Network.Transport.TCP (QDisc (..), simpleOnePlaceQDisc, simpleUnboundedQDisc)
import           Statistics.Distribution
import           Statistics.Distribution.Exponential
import           Statistics.Distribution.Normal
import           Statistics.Distribution.Uniform
import qualified Statistics.Sample as Sample
import           System.Environment (getArgs)
import           System.Random.MWC

-- | A writer is determined by some continuous distribution giving the duration
--   (in microseconds) between successive data being made available.
data SimulationWriter = forall distr . ContGen distr => SimulationWriter distr

-- | A reader is determined by some continuous distribution giving the duration
--   (in microseconds) between successive reads (how long the reader thread
--   works between taking events).
data SimulationReader = forall distr . ContGen distr => SimulationReader distr

data Scenario = Scenario {
      sim_reader  :: SimulationReader
    , sim_writers :: [SimulationWriter]
    }

data SimulationParameters = SimulationParameters {
      sim_scenario :: Scenario
      -- | How long the simulation should run.
    , sim_duration :: Second
      -- | The QDisc to use.
    , sim_qdisc    :: QDisc ()
      -- | A seed for randomness.
    , sim_seed     :: Word32
    }

type Latency = Double

-- | The output of a simulation is, for each writer, the samples of actual
--   delays observed when trying to write: how long it was blocked on
--   trying to enqueue. The number of samples is the number of writes it
--   made.
data SimulationOutput = SimulationOutput {
      sim_writer_outputs :: [Vector Latency]
    }

instance Show SimulationOutput where
    show (SimulationOutput vecs) = concat $ flip fmap vecs $ \vec -> concat [
          "Samples (writes): ", show (V.length vec), "\n"
        , "Latency (microseconds):\n"
        , "  Mean: ", show (Sample.mean vec), "\n"
        , "  Std. Dev.: ", show (Sample.stdDev vec), "\n\n"
        ]

data QDiscChoice = Fair | OnePlace | Unbounded

makeQDisc :: QDiscChoice -> IO (QDisc t)
makeQDisc choice = case choice of
    Fair      -> fairQDisc (const (return Nothing))
    Unbounded -> simpleUnboundedQDisc
    OnePlace  -> simpleOnePlaceQDisc

unfairScenario :: Scenario
unfairScenario = Scenario {
      -- Difference between fair and one-place is very clear here.
      -- The two fast writers get the same number of writes in either case.
      -- One-place gives the slow writer half as many writes, but fair gives
      -- it a lot more. The cycle time (time to read from all 3 writers) is
      -- the slow writer's write delay, so we expect that the slow writer
      -- should be able to write on every cycle.
      sim_reader = SimulationReader (uniformDistr 5000 5001)
    , sim_writers = [
            SimulationWriter (uniformDistr 4999 5000)
          , SimulationWriter (uniformDistr 4999 5000)
          , SimulationWriter (uniformDistr 14999 15000)
          ]
    }

-- | A normally-distributed reader (mean and std. dev. configurable) and
--   n exponentially-distributed writers (means configurable).
typicalScenario :: (Double, Double) -> [Double] -> Scenario
typicalScenario (rmean, rstd_dev) writers = Scenario {
      sim_reader = SimulationReader (normalDistr rmean rstd_dev)
    , sim_writers = flip fmap writers $ \wmean -> SimulationWriter (exponential (1/wmean))
    }

simpleSimulationParameters :: QDisc () -> Second -> Scenario -> SimulationParameters
simpleSimulationParameters qdisc duration scenario = SimulationParameters {
      sim_scenario = scenario
    , sim_duration = duration
    , sim_qdisc = qdisc
    , sim_seed = 42
    }

-- | Run a simulation.
simulate :: SimulationParameters -> IO SimulationOutput
simulate SimulationParameters{..} = do

    let Scenario{..} = sim_scenario

    -- All threads will wait on this.
    -- Threads will loop, and at each iteration will read it. If it's True,
    -- they'll stop.
    startStop :: MVar Bool <- newEmptyMVar

    -- Spawn the reader.
    refs <- withAsync (reader 0 sim_reader) $ \readerThread -> do
        -- Spawn the writers.
        writerThreadsAndRefs <- forM (zip [1..] sim_writers) $ \(seed', sim_writer) -> do
            (ref, doWrites) <- writer startStop seed' sim_writer
            --withAsync doWrites $ \thread -> return (ref, thread)
            thread <- async doWrites
            return (ref, thread)
        -- Duration is in seconds.
        putMVar startStop False
        threadDelay $ fromIntegral sim_duration * 1000000
        swapMVar startStop True
        forM writerThreadsAndRefs $ \(ref, thread) -> do
            wait thread
            return ref

    -- Reader and writers are all killed. Results available in the IORefs.
    vectors <- forM refs (fmap fromList . readIORef)

    return $ SimulationOutput ((fmap . fmap) fromIntegral vectors)

    where

    reader :: Word32 -> SimulationReader -> IO ()
    reader seed' (SimulationReader distribution) = do
        gen <- initialize (fromList [sim_seed, seed'])
        let readLoop = do
                () <- qdiscDequeue sim_qdisc
                -- Unlike for writers, the reader delays are always respected,
                -- as the delay is independent of how long the reader waits for
                -- a value from qdiscDequeue.
                delay :: Microsecond <- fromIntegral . round <$> genContVar distribution gen
                threadDelay (fromIntegral delay)
                readLoop
        readLoop

    writer :: MVar Bool -> Word32 -> SimulationWriter -> IO (IORef [Microsecond], IO ())
    writer control seed' (SimulationWriter distribution) = do
        ref <- newIORef []
        gen <- initialize (fromList [sim_seed, seed'])
        -- The writer's distribution determines the duration between the points
        -- in time when new data are made available, and this is independent
        -- of how long it takes to actually do the write. For instance, if
        -- the writer is blocked on enqueue for 2 seconds, and the distribution
        -- says there will be data available every 1 second, then 2 data will
        -- be immediately available after the enqueue finishes.
        let writeLoop :: Microsecond -> IO ()
            writeLoop surplusWait = do
                stop <- readMVar control
                unless stop $ do
                    nextDelay :: Microsecond <- fromIntegral . round <$> genContVar distribution gen
                    let actualDelay = nextDelay - surplusWait
                    when (actualDelay > 0) $ do
                        threadDelay (fromIntegral actualDelay)
                    start <- getPOSIXTime
                    qdiscEnqueue sim_qdisc undefined undefined ()
                    end <- getPOSIXTime
                    let latency :: Microsecond
                        !latency = fromIntegral . round $ (realToFrac (end - start) :: Double) * 1000000
                    modifyIORef' ref ((:) latency)
                    let !surplusWait' = max 0 (surplusWait + latency - nextDelay)
                    writeLoop surplusWait'
        return (ref, writeLoop 0)

main :: IO ()
main = do
    args <- getArgs
    qdiscChoice <- case args of
        "unbounded" : _ -> putStrLn "Using unbounded QDisc" >> return Unbounded
        "one_place" : _ -> putStrLn "Using one-place QDisc" >> return OnePlace
        _               -> putStrLn "Using fair QDisc" >> return Fair
    duration :: Int <- case args of
        _ : n : _ -> case reads n of
            [(n', "")] -> return n'
        _ -> return 10
    qdisc <- makeQDisc qdiscChoice
    --let scenario = typicalScenario (1000, 25) [800, 900, 1000, 1100, 1200]
    let scenario = unfairScenario
    results <- simulate $ simpleSimulationParameters qdisc (fromIntegral duration) scenario
    putStrLn ""
    print results


{-# LANGUAGE AllowAmbiguousTypes, FlexibleContexts, NamedFieldPuns #-}

module School.TestUtils
( CostFunc
, addClasses
, assertRight
, def
, diffCost
, diffInput
, doCost
, dummyList
, dummyMatrix
, empty
, fromLeft
, fromRight
, getRandDouble
, isSorted
, jTest
, matIndexes
, randomAffineParams
, randomMatrix
, randomMatrixL
, randomNNInts
, randomVector
, singleClassOutput
, testIOCatch
, testState
, unitCorrect
, weight1
, whenPrint
) where

import Conduit (ConduitM, liftIO, runConduit)
import Control.Exception (catch)
import Control.Monad (when)
import Data.Default.Class (def)
import Data.Either (either)
import Data.List (sort)
import Data.Void (Void)
import Numeric.LinearAlgebra ((><), (|>), Element, IndexOf, Matrix, R, Vector, accum, assoc, fromColumns, fromList, fromRows, size, sumElements, toColumns)
import School.Train.AppTrain (AppTrain, runAppTrain)
import School.Train.TrainState (TrainState)
import School.Types.Slinky (Slinky)
import School.Unit.CostFunction (CostFunction(..))
import School.Unit.CostParams (CostParams)
import School.Unit.Unit (Unit(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitParams (UnitParams(..))
import School.Unit.WeightDecay (weightDecay)
import School.Utils.Double (doubleRange)
import System.Exit (ExitCode(..))
import System.Random (getStdRandom, randomR)
import Test.QuickCheck.Monadic (PropertyM, assert)

testIOCatch :: IO a -> PropertyM IO ExitCode
testIOCatch action = liftIO $
  catch (action >> (return ExitSuccess))
        return

assertRight :: (b -> Bool)
            -> Either a b
            -> PropertyM IO ()
assertRight f = either (const $ assert False)
                       (\x -> assert $ f x)

testState :: ConduitM () Void (AppTrain R) b
          -> TrainState R
          -> PropertyM IO (Either String (b, TrainState R))
testState conduit state =
  liftIO . (runAppTrain state) . runConduit $ conduit

isSorted :: (Ord a) => [a] -> Bool
isSorted xs = xs == (reverse . sort $ xs)

dummyList :: (Num a) => Int -> Int -> [a]
dummyList r c = map fromIntegral [ r*i + j | i <- [1..r], j <- [1..c] ]

dummyMatrix :: Int -> Int -> Matrix Double
dummyMatrix r c = r >< c $ dummyList r c

getRandDouble :: IO Double
getRandDouble =
  getStdRandom (randomR doubleRange)

randomNNInts :: Int -> Int -> IO [Int]
randomNNInts maxInt n =
  sequence . (replicate n) $ getStdRandom (randomR (0, maxInt))

matMax :: Double
matMax = 1e3

randomMatrix :: Int -> Int -> IO (Matrix R)
randomMatrix nRows nCols = do
  let nEls = nRows * nCols
  inputList <- sequence . (replicate nEls) $ getStdRandom (randomR (-matMax, matMax))
  return $ (nRows >< nCols) inputList

randomMatrixL :: Double -> Int -> Int -> IO (Matrix R)
randomMatrixL limit nRows nCols = do
  let nEls = nRows * nCols
  inputList <- sequence . (replicate nEls) $ getStdRandom (randomR (-limit, limit))
  return $ (nRows >< nCols) inputList

randomVector :: Int -> IO (Vector R)
randomVector n = do
  inputList <- sequence . (replicate n) $ getStdRandom (randomR (-matMax, matMax))
  return $ n |> inputList

randomAffineParams :: Int
                   -> Int
                   -> IO (UnitParams R)
randomAffineParams fSize oSize = do
  affineBias <- randomVector oSize
  affineWeights <- randomMatrix oSize fSize
  return AffineParams { affineBias
                      , affineWeights
                      }

toPrint :: (Show a) => String -> a -> IO ()
toPrint tag val = putStrLn $ tag ++ " " ++ (show val)

whenPrint :: (Show a) => Bool -> [String] -> [a] -> IO ()
whenPrint cond tags vals = do
  when cond (sequence_ $ zipWith toPrint tags vals)

type CostFunc = UnitActivation R -> Double

jTest :: CostFunc
jTest (BatchActivation m) = sumElements m
jTest _ = 0

diffInput :: Unit R
          -> UnitParams R
          -> UnitActivation R
          -> Double
          -> IndexOf Matrix
          -> Double
diffInput unit params input eps idx =
  (jAdd - jSub) / (2*eps) where
    outAdd = apply unit params (alterInput eps idx input)
    outSub = apply unit params (alterInput (-eps) idx input)
    jAdd = jTest outAdd
    jSub = jTest outSub

fromRight :: b -> Either a b -> b
fromRight b = either (const b) id

fromLeft :: a -> Either a b -> a
fromLeft b = either id (const b)

diffCost :: CostFunction R m
         -> UnitActivation R
         -> Double
         -> IndexOf Matrix
         -> Slinky CostParams
         -> Double
diffCost costFunc input eps idx costParams = let
  jAdd = computeCost costFunc (alterInput eps idx input) costParams
  jSub = computeCost costFunc (alterInput (-eps) idx input) costParams
  in (fromRight 0 jAdd - fromRight 0 jSub) / (2*eps)

alterInput :: AlterInput
alterInput change idx (BatchActivation m) =
  BatchActivation $ accum m (+) [(idx, change)]
alterInput _ _ _ = ApplyFail "alterInput error"

type AlterInput = Double
               -> IndexOf Matrix
               -> UnitActivation R
               -> UnitActivation R

matIndexes :: Int -> Int -> [IndexOf Matrix]
matIndexes r c = [ (j, k) | j <- [0..r-1], k <- [0..c-1] ]

empty :: (Element a) => Matrix a
empty = (0><0) []

weight1 :: (Monad m) => CostFunction R m
weight1 = weightDecay 1

doCost :: (Element a, Num a)
       => CostFunction a m
       -> UnitActivation a
       -> Slinky CostParams
       -> (a, UnitGradient a)
doCost costFunction activation params =
  fromRight (0, BatchGradient empty) result where
    result = do
      cost <- computeCost costFunction activation params
      grad <- derivCost costFunction activation params
      return (cost, grad)

appCol :: (Element a)
       => Vector a
       -> Matrix a
       -> Matrix a
appCol v =
  fromColumns . (flip (++) $ [v]) . toColumns

addClasses :: (Element a, Num a)
           => [Int]
           -> Matrix a
           -> Matrix a
addClasses classes mat = appCol v mat where
  v = fromList . (map fromIntegral) $ classes

singleClassOutput :: Int -> [Int] -> UnitActivation R
singleClassOutput nClasses =
  BatchActivation . fromRows . (map (\idx -> assoc nClasses 0 [(idx, 1)]))

unitCorrect :: [Int] -> Unit R
unitCorrect classes = Unit { apply, deriv } where
  apply _ (BatchActivation input) =
    singleClassOutput n classes
    where (_, n) = size input
  apply _ _ = ApplyFail ""
  deriv = undefined