{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ >= 708

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}

{- |
Module      :  Internal.Static
Copyright   :  (c) Alberto Ruiz 2006-14
License     :  BSD3
Stability   :  provisional

-}

module Internal.Static where

import GHC.TypeLits
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra hiding (konst,size,R,C)
import Internal.Vector as D hiding (R,C)
import Internal.ST
import Control.DeepSeq
import Data.Proxy(Proxy)
import Foreign.Storable(Storable)
import Text.Printf

import Data.Binary
import GHC.Generics (Generic)
import Data.Proxy (Proxy(..))

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

type ℝ = Float
type ℂ = Complex Float

newtype Dim (n :: Nat) t = Dim t
  deriving (Show, Generic)

instance (KnownNat n, Binary a) => Binary (Dim n a) where
  get = do
    k <- get
    let n = natVal (Proxy :: Proxy n)
    if n == k
      then Dim <$> get
      else fail ("Expected dimension " ++ (show n) ++ ", but found dimension " ++ (show k))

  put (Dim x) = do
    put (natVal (Proxy :: Proxy n))
    put x

lift1F
  :: (c t -> c t)
  -> Dim n (c t) -> Dim n (c t)
lift1F f (Dim v) = Dim (f v)

lift2F
  :: (c t -> c t -> c t)
  -> Dim n (c t) -> Dim n (c t) -> Dim n (c t)
lift2F f (Dim u) (Dim v) = Dim (f u v)

instance NFData t => NFData (Dim n t) where
    rnf (Dim (force -> !_)) = ()

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

newtype R n = R (Dim n (Vector ℝ))
  deriving (Num,Fractional,Floating,Generic,Binary)

newtype C n = C (Dim n (Vector ℂ))
  deriving (Num,Fractional,Floating,Generic)

newtype L m n = L (Dim m (Dim n (Matrix ℝ)))
  deriving (Generic, Binary)

newtype M m n = M (Dim m (Dim n (Matrix ℂ)))
  deriving (Generic)

mkR :: Vector ℝ -> R n
mkR = R . Dim

mkC :: Vector ℂ -> C n
mkC = C . Dim

mkL :: Matrix ℝ -> L m n
mkL x = L (Dim (Dim x))

mkM :: Matrix ℂ -> M m n
mkM x = M (Dim (Dim x))

instance NFData (R n) where
    rnf (R (force -> !_)) = ()

instance NFData (C n) where
    rnf (C (force -> !_)) = ()

instance NFData (L n m) where
    rnf (L (force -> !_)) = ()

instance NFData (M n m) where
    rnf (M (force -> !_)) = ()

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

type V n t = Dim n (Vector t)

ud :: Dim n (Vector t) -> Vector t
ud (Dim v) = v

mkV :: forall (n :: Nat) t . t -> Dim n t
mkV = Dim

vconcat :: forall n m t . (KnownNat n, KnownNat m, Numeric t)
    => V n t -> V m t -> V (n+m) t
(ud -> u) `vconcat` (ud -> v) = mkV (vjoin [u', v'])
  where
    du = fromIntegral . natVal $ (undefined :: Proxy n)
    dv = fromIntegral . natVal $ (undefined :: Proxy m)
    u' | du /= 1 && LA.size u == 1 = LA.konst (u D.@> 0) du
       | otherwise = u
    v' | dv /= 1 && LA.size v == 1 = LA.konst (v D.@> 0) dv
       | otherwise = v

gvec2 :: Storable t => t -> t -> V 2 t
gvec2 a b = mkV $ runSTVector $ do
    v <- newUndefinedVector 2
    writeVector v 0 a
    writeVector v 1 b
    return v

gvec3 :: Storable t => t -> t -> t -> V 3 t
gvec3 a b c = mkV $ runSTVector $ do
    v <- newUndefinedVector 3
    writeVector v 0 a
    writeVector v 1 b
    writeVector v 2 c
    return v

gvec4 :: Storable t => t -> t -> t -> t -> V 4 t
gvec4 a b c d = mkV $ runSTVector $ do
    v <- newUndefinedVector 4
    writeVector v 0 a
    writeVector v 1 b
    writeVector v 2 c
    writeVector v 3 d
    return v

gvect :: forall n t . (Show t, KnownNat n, Numeric t) => String -> [t] -> V n t
gvect st xs'
    | ok = mkV v
    | not (null rest) && null (tail rest) = abort (show xs')
    | not (null rest) = abort (init (show (xs++take 1 rest))++", ... ]")
    | otherwise = abort (show xs)
  where
    (xs,rest) = splitAt d xs'
    ok = LA.size v == d && null rest
    v = LA.fromList xs
    d = fromIntegral . natVal $ (undefined :: Proxy n)
    abort info = error $ st++" "++show d++" can't be created from elements "++info

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

type GM m n t = Dim m (Dim n (Matrix t))

gmat :: forall m n t . (Show t, KnownNat m, KnownNat n, Numeric t) => String -> [t] -> GM m n t
gmat st xs'
    | ok = Dim (Dim x)
    | not (null rest) && null (tail rest) = abort (show xs')
    | not (null rest) = abort (init (show (xs++take 1 rest))++", ... ]")
    | otherwise = abort (show xs)
  where
    (xs,rest) = splitAt (m'*n') xs'
    v = LA.fromList xs
    x = reshape n' v
    ok = null rest && ((n' == 0 && dim v == 0) || n'> 0 && (rem (LA.size v) n' == 0) && LA.size x == (m',n'))
    m' = fromIntegral . natVal $ (undefined :: Proxy m) :: Int
    n' = fromIntegral . natVal $ (undefined :: Proxy n) :: Int
    abort info = error $ st ++" "++show m' ++ " " ++ show n'++" can't be created from elements " ++ info

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

class Num t => Sized t s d | s -> t, s -> d
  where
    konst     ::  t   -> s
    unwrap    ::  s   -> d t
    fromList  :: [t]  -> s
    extract   ::  s   -> d t
    create    ::  d t -> Maybe s
    size      ::  s   -> IndexOf d

singleV v = LA.size v == 1
singleM m = rows m == 1 && cols m == 1

instance KnownNat n => Sized ℂ (C n) Vector
  where
    size _ = fromIntegral . natVal $ (undefined :: Proxy n)
    konst x = mkC (LA.scalar x)
    unwrap (C (Dim v)) = v
    fromList xs = C (gvect "C" xs)
    extract s@(unwrap -> v)
      | singleV v = LA.konst (v!0) (size s)
      | otherwise = v
    create v
        | LA.size v == size r = Just r
        | otherwise = Nothing
      where
        r = mkC v :: C n

instance KnownNat n => Sized ℝ (R n) Vector
  where
    size _ = fromIntegral . natVal $ (undefined :: Proxy n)
    konst x = mkR (LA.scalar x)
    unwrap (R (Dim v)) = v
    fromList xs = R (gvect "R" xs)
    extract s@(unwrap -> v)
      | singleV v = LA.konst (v!0) (size s)
      | otherwise = v
    create v
        | LA.size v == size r = Just r
        | otherwise = Nothing
      where
        r = mkR v :: R n

instance (KnownNat m, KnownNat n) => Sized ℝ (L m n) Matrix
  where
    size _ = ((fromIntegral . natVal) (undefined :: Proxy m)
             ,(fromIntegral . natVal) (undefined :: Proxy n))
    konst x = mkL (LA.scalar x)
    fromList xs = L (gmat "L" xs)
    unwrap (L (Dim (Dim m))) = m
    extract (isDiag -> Just (z,y,(m',n'))) = diagRect z y m' n'
    extract s@(unwrap -> a)
        | singleM a = LA.konst (a `atIndex` (0,0)) (size s)
        | otherwise = a
    create x
        | LA.size x == size r = Just r
        | otherwise = Nothing
      where
        r = mkL x :: L m n

instance (KnownNat m, KnownNat n) => Sized ℂ (M m n) Matrix
  where
    size _ = ((fromIntegral . natVal) (undefined :: Proxy m)
             ,(fromIntegral . natVal) (undefined :: Proxy n))
    konst x = mkM (LA.scalar x)
    fromList xs = M (gmat "M" xs)
    unwrap (M (Dim (Dim m))) = m
    extract (isDiagC -> Just (z,y,(m',n'))) = diagRect z y m' n'
    extract s@(unwrap -> a)
        | singleM a = LA.konst (a `atIndex` (0,0)) (size s)
        | otherwise = a
    create x
        | LA.size x == size r = Just r
        | otherwise = Nothing
      where
        r = mkM x :: M m n

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

instance (KnownNat n, KnownNat m) => Transposable (L m n) (L n m)
  where
    tr a@(isDiag -> Just _) = mkL (extract a)
    tr (extract -> a) = mkL (tr a)
    tr' = tr

instance (KnownNat n, KnownNat m) => Transposable (M m n) (M n m)
  where
    tr a@(isDiagC -> Just _) = mkM (extract a)
    tr (extract -> a) = mkM (tr a)
    tr' a@(isDiagC -> Just _) = mkM (extract a)
    tr' (extract -> a) = mkM (tr' a)

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

isDiag :: forall m n . (KnownNat m, KnownNat n) => L m n -> Maybe (ℝ, Vector ℝ, (Int,Int))
isDiag (L x) = isDiagg x

isDiagC :: forall m n . (KnownNat m, KnownNat n) => M m n -> Maybe (ℂ, Vector ℂ, (Int,Int))
isDiagC (M x) = isDiagg x

isDiagg :: forall m n t . (Numeric t, KnownNat m, KnownNat n) => GM m n t -> Maybe (t, Vector t, (Int,Int))
isDiagg (Dim (Dim x))
    | singleM x = Nothing
    | rows x == 1 && m' > 1 || cols x == 1 && n' > 1 = Just (z,yz,(m',n'))
    | otherwise = Nothing
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m) :: Int
    n' = fromIntegral . natVal $ (undefined :: Proxy n) :: Int
    v = flatten x
    z = v `atIndex` 0
    y = subVector 1 (LA.size v-1) v
    ny = LA.size y
    zeros = LA.konst 0 (max 0 (min m' n' - ny))
    yz = vjoin [y,zeros]

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

instance KnownNat n => Show (R n)
  where
    show s@(R (Dim v))
      | singleV v = "(" ++ show (v!0) ++ " :: R " ++ show d ++ ")"
      | otherwise = "(vector " ++ show v ++ " :: R " ++ show d ++")"
      where
        d = size s

instance KnownNat n => Show (C n)
  where
    show s@(C (Dim v))
      | singleV v = "(" ++ show (v!0) ++ " :: C " ++ show d ++ ")"
      | otherwise = "(vector " ++ show v ++ " :: C " ++ show d ++")"
      where
        d = size s

instance (KnownNat m, KnownNat n) => Show (L m n)
  where
    show (isDiag -> Just (z,y,(m',n'))) = printf "(diag %s %s :: L %d %d)" (show z) (show y) m' n'
    show s@(L (Dim (Dim x)))
       | singleM x = printf "(%s :: L %d %d)" (show (x `atIndex` (0,0))) m' n'
       | otherwise = "(matrix" ++ dropWhile (/='\n') (show x) ++ " :: L " ++ show m' ++ " " ++ show n' ++ ")"
      where
        (m',n') = size s

instance (KnownNat m, KnownNat n) => Show (M m n)
  where
    show (isDiagC -> Just (z,y,(m',n'))) = printf "(diag %s %s :: M %d %d)" (show z) (show y) m' n'
    show s@(M (Dim (Dim x)))
       | singleM x = printf "(%s :: M %d %d)" (show (x `atIndex` (0,0))) m' n'
       | otherwise = "(matrix" ++ dropWhile (/='\n') (show x) ++ " :: M " ++ show m' ++ " " ++ show n' ++ ")"
      where
        (m',n') = size s

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

instance (Num (Vector t), Numeric t )=> Num (Dim n (Vector t))
  where
    (+) = lift2F (+)
    (*) = lift2F (*)
    (-) = lift2F (-)
    abs = lift1F abs
    signum = lift1F signum
    negate = lift1F negate
    fromInteger x = Dim (fromInteger x)

instance (Num (Vector t), Num (Matrix t), Fractional t, Numeric t) => Fractional (Dim n (Vector t))
  where
    fromRational x = Dim (fromRational x)
    (/) = lift2F (/)

instance (Fractional t, Floating (Vector t), Numeric t) => Floating (Dim n (Vector t)) where
    sin   = lift1F sin
    cos   = lift1F cos
    tan   = lift1F tan
    asin  = lift1F asin
    acos  = lift1F acos
    atan  = lift1F atan
    sinh  = lift1F sinh
    cosh  = lift1F cosh
    tanh  = lift1F tanh
    asinh = lift1F asinh
    acosh = lift1F acosh
    atanh = lift1F atanh
    exp   = lift1F exp
    log   = lift1F log
    sqrt  = lift1F sqrt
    (**)  = lift2F (**)
    pi    = Dim pi

instance (Num (Matrix t), Numeric t) => Num (Dim m (Dim n (Matrix t)))
  where
    (+) = (lift2F . lift2F) (+)
    (*) = (lift2F . lift2F) (*)
    (-) = (lift2F . lift2F) (-)
    abs = (lift1F . lift1F) abs
    signum = (lift1F . lift1F) signum
    negate = (lift1F . lift1F) negate
    fromInteger x = Dim (Dim (fromInteger x))

instance (Num (Vector t), Num (Matrix t), Fractional t, Numeric t) => Fractional (Dim m (Dim n (Matrix t)))
  where
    fromRational x = Dim (Dim (fromRational x))
    (/) = (lift2F.lift2F) (/)

instance (Num (Vector t), Floating (Matrix t), Fractional t, Numeric t) => Floating (Dim m (Dim n (Matrix t))) where
    sin   = (lift1F . lift1F) sin
    cos   = (lift1F . lift1F) cos
    tan   = (lift1F . lift1F) tan
    asin  = (lift1F . lift1F) asin
    acos  = (lift1F . lift1F) acos
    atan  = (lift1F . lift1F) atan
    sinh  = (lift1F . lift1F) sinh
    cosh  = (lift1F . lift1F) cosh
    tanh  = (lift1F . lift1F) tanh
    asinh = (lift1F . lift1F) asinh
    acosh = (lift1F . lift1F) acosh
    atanh = (lift1F . lift1F) atanh
    exp   = (lift1F . lift1F) exp
    log   = (lift1F . lift1F) log
    sqrt  = (lift1F . lift1F) sqrt
    (**)  = (lift2F . lift2F) (**)
    pi    = Dim (Dim pi)

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

adaptDiag f a@(isDiag -> Just _) b | isFull b = f (mkL (extract a)) b
adaptDiag f a b@(isDiag -> Just _) | isFull a = f a (mkL (extract b))
adaptDiag f a b = f a b

isFull m = isDiag m == Nothing && not (singleM (unwrap m))

lift1L f (L v) = L (f v)
lift2L f (L a) (L b) = L (f a b)
lift2LD f = adaptDiag (lift2L f)

instance (KnownNat n, KnownNat m) =>  Num (L n m)
  where
    (+) = lift2LD (+)
    (*) = lift2LD (*)
    (-) = lift2LD (-)
    abs = lift1L abs
    signum = lift1L signum
    negate = lift1L negate
    fromInteger = L . Dim . Dim . fromInteger

instance (KnownNat n, KnownNat m) => Fractional (L n m)
  where
    fromRational = L . Dim . Dim . fromRational
    (/) = lift2LD (/)

instance (KnownNat n, KnownNat m) => Floating (L n m) where
    sin   = lift1L sin
    cos   = lift1L cos
    tan   = lift1L tan
    asin  = lift1L asin
    acos  = lift1L acos
    atan  = lift1L atan
    sinh  = lift1L sinh
    cosh  = lift1L cosh
    tanh  = lift1L tanh
    asinh = lift1L asinh
    acosh = lift1L acosh
    atanh = lift1L atanh
    exp   = lift1L exp
    log   = lift1L log
    sqrt  = lift1L sqrt
    (**)  = lift2LD (**)
    pi    = konst pi

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

adaptDiagC f a@(isDiagC -> Just _) b | isFullC b = f (mkM (extract a)) b
adaptDiagC f a b@(isDiagC -> Just _) | isFullC a = f a (mkM (extract b))
adaptDiagC f a b = f a b

isFullC m = isDiagC m == Nothing && not (singleM (unwrap m))

lift1M f (M v) = M (f v)
lift2M f (M a) (M b) = M (f a b)
lift2MD f = adaptDiagC (lift2M f)

instance (KnownNat n, KnownNat m) =>  Num (M n m)
  where
    (+) = lift2MD (+)
    (*) = lift2MD (*)
    (-) = lift2MD (-)
    abs = lift1M abs
    signum = lift1M signum
    negate = lift1M negate
    fromInteger = M . Dim . Dim . fromInteger

instance (KnownNat n, KnownNat m) => Fractional (M n m)
  where
    fromRational = M . Dim . Dim . fromRational
    (/) = lift2MD (/)

instance (KnownNat n, KnownNat m) => Floating (M n m) where
    sin   = lift1M sin
    cos   = lift1M cos
    tan   = lift1M tan
    asin  = lift1M asin
    acos  = lift1M acos
    atan  = lift1M atan
    sinh  = lift1M sinh
    cosh  = lift1M cosh
    tanh  = lift1M tanh
    asinh = lift1M asinh
    acosh = lift1M acosh
    atanh = lift1M atanh
    exp   = lift1M exp
    log   = lift1M log
    sqrt  = lift1M sqrt
    (**)  = lift2MD (**)
    pi    = M pi

instance Additive (R n) where
    add = (+)

instance Additive (C n) where
    add = (+)

instance (KnownNat m, KnownNat n) => Additive (L m n) where
    add = (+)

instance (KnownNat m, KnownNat n) => Additive (M m n) where
    add = (+)

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

class Disp t
  where
    disp :: Int -> t -> IO ()

instance (KnownNat m, KnownNat n) => Disp (L m n)
  where
    disp n x = do
        let a = extract x
        let su = LA.dispf n a
        printf "L %d %d" (rows a) (cols a) >> putStr (dropWhile (/='\n') $ su)

instance (KnownNat m, KnownNat n) => Disp (M m n)
  where
    disp n x = do
        let a = extract x
        let su = LA.dispcf n a
        printf "M %d %d" (rows a) (cols a) >> putStr (dropWhile (/='\n') $ su)

instance KnownNat n => Disp (R n)
  where
    disp n v = do
        let su = LA.dispf n (asRow $ extract v)
        putStr "R " >> putStr (tail . dropWhile (/='x') $ su)

instance KnownNat n => Disp (C n)
  where
    disp n v = do
        let su = LA.dispcf n (asRow $ extract v)
        putStr "C " >> putStr (tail . dropWhile (/='x') $ su)

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

overMatL' :: (KnownNat m, KnownNat n)
          => (LA.Matrix ℝ -> LA.Matrix ℝ) -> L m n -> L m n
overMatL' f = mkL . f . unwrap
{-# INLINE overMatL' #-}

overMatM' :: (KnownNat m, KnownNat n)
          => (LA.Matrix ℂ -> LA.Matrix ℂ) -> M m n -> M m n
overMatM' f = mkM . f . unwrap
{-# INLINE overMatM' #-}

#else

module Numeric.LinearAlgebra.Static.Internal where

#endif


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Recurrent.Layers.BasicRecurrent (
    BasicRecurrent (..)
  ) where

import           Data.Proxy
import           Data.Singletons.TypeLits hiding (natVal)

import           Numeric.LinearAlgebra.Static

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Recurrent.Core

data BasicRecurrent :: Nat -- Input layer size
                    -> Nat -- Output layer size
                    -> * where
  BasicRecurrent :: ( KnownNat input
                    , KnownNat output
                    , KnownNat matrixCols
                    , matrixCols ~ (input + output))
                 => !(R output)   -- Bias neuron weights
                 -> !(R output)   -- Bias neuron momentum
                 -> !(L output matrixCols) -- Activation
                 -> !(L output matrixCols) -- Momentum
                 -> BasicRecurrent input output

data BasicRecurrent' :: Nat -- Input layer size
                     -> Nat -- Output layer size
                     -> * where
  BasicRecurrent' :: ( KnownNat input
                     , KnownNat output
                     , KnownNat matrixCols
                     , matrixCols ~ (input + output))
                  => !(R output)   -- Bias neuron gradients
                  -> !(L output matrixCols)
                  -> BasicRecurrent' input output

instance Show (BasicRecurrent i o) where
  show BasicRecurrent {} = "BasicRecurrent"

instance (KnownNat i, KnownNat o, KnownNat (i + o)) => UpdateLayer (BasicRecurrent i o) where
  type Gradient (BasicRecurrent i o) = (BasicRecurrent' i o)

  runUpdate LearningParameters {..} (BasicRecurrent oldBias oldBiasMomentum oldActivations oldMomentum) (BasicRecurrent' biasGradient activationGradient) =
    let newBiasMomentum = konst learningMomentum * oldBiasMomentum - konst learningRate * biasGradient
        newBias         = oldBias + newBiasMomentum
        newMomentum     = konst learningMomentum * oldMomentum - konst learningRate * activationGradient
        regulariser     = konst (learningRegulariser * learningRate) * oldActivations
        newActivations  = oldActivations + newMomentum - regulariser
    in BasicRecurrent newBias newBiasMomentum newActivations newMomentum

instance (KnownNat i, KnownNat o, KnownNat x, KnownNat (x*o), x ~ (i+o)) => RandomLayer (BasicRecurrent i o) where
  createRandomWith m gen = do
    wB <- getRandomVector i o m gen
    wN <- getRandomMatrix i o m gen
    let bm = konst 0
        mm = konst 0
    return $ BasicRecurrent wB bm wN mm
      where i = natVal (Proxy :: Proxy i)
            o = natVal (Proxy :: Proxy o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentUpdateLayer (BasicRecurrent i o) where
  type RecurrentShape (BasicRecurrent i o) = S ('D1 o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentLayer (BasicRecurrent i o) ('D1 i) ('D1 o) where

  type RecTape (BasicRecurrent i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
  -- Do a matrix vector multiplication and return the result.
  runRecurrentForwards (BasicRecurrent wB _ wN _) (S1D lastOutput) (S1D thisInput) =
    let thisOutput = S1D $ wB + wN #> (thisInput # lastOutput)
    in ((S1D lastOutput, S1D thisInput), thisOutput, thisOutput)

  -- Run a backpropogation step for a full connected layer.
  runRecurrentBackwards (BasicRecurrent _ _ wN _) (S1D lastOutput, S1D thisInput) (S1D dRec) (S1D dEdy) =
    let biasGradient        = (dRec + dEdy)
        layerGrad           = (dRec + dEdy) `outer` (thisInput # lastOutput)
        -- calcluate derivatives for next step
        (backGrad, recGrad) = split $ tr wN #> (dRec + dEdy)
    in  (BasicRecurrent' biasGradient layerGrad, S1D recGrad, S1D backGrad)

-------------------- GNum instances --------------------

instance (KnownNat i,KnownNat o,KnownNat (i+o)) => GNum (BasicRecurrent i o) where
  n |* (BasicRecurrent wB mB a nM) = (BasicRecurrent (fromRational n * wB) mB a nM)
  (BasicRecurrent wB mB a nM) |+ (BasicRecurrent wB2 mB2 a2 nM2)  = (BasicRecurrent (0.5 * (wB + wB2)) (0.5* (mB+mB2)) (0.5 * (a+a2)) (0.5 * (nM + nM2)))
  gFromRational r = (BasicRecurrent (fromRational r) 0 (fromRational r) 0)


{-# OPTIONS_GHC -fno-cse #-}
{-# OPTIONS_GHC -fno-full-laziness #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE BangPatterns #-}
module Trace(
      clearTrace
    , traceNames
    , CanBeTraced
    , traceValues
    , trace
    , traceSample
    , forceSignal
    ) where

import Common(HasDoubleRepresentation(..))
import System.IO.Unsafe
import qualified Data.Map as M
import Data.IORef
import Control.Applicative((<$>))
import Signal(Signal,mapS,fromListS,takeS)
import Fixed(Resolution(..))
import Statistics.Sample.KernelDensity
import qualified Data.Vector.Unboxed as U
import Viewer
import Plot
import Graphics.PDF
import Displayable
import Internal
import Control.DeepSeq
import Text.Printf
import qualified Graphics.PDF as PDF

import qualified Debug.Trace as T
debug a = T.trace (show a) a

data CanBeTraced = forall a. (HasDoubleRepresentation a,Resolution a) => CBT !a

instance HasDoubleRepresentation CanBeTraced where
    toDouble (CBT a) = toDouble a
    fromDouble = error "Can't create a CanBeTraced from a double"

instance Resolution CanBeTraced where
    smallestValue (CBT a) = CBT (smallestValue a)
    maxValue (CBT a) = CBT (maxValue a)
    minValue (CBT a) = CBT (minValue a)
    signedFormat (CBT a) = signedFormat a
    bitWidth (CBT a) = bitWidth a

{-# NOINLINE globalTrace #-}
globalTrace  :: IORef (M.Map String [CanBeTraced])
globalTrace = unsafePerformIO $ newIORef (M.empty)

clearTrace :: IO ()
clearTrace = writeIORef globalTrace M.empty

traceNames :: IO [String]
traceNames = M.keys <$> (readIORef  globalTrace)

getTrace :: String -> IO (Maybe [CanBeTraced])
getTrace name = M.lookup name <$> (readIORef globalTrace)

removeZeroAndLog :: [Double] -> [Double] -> ([Double],[Double])
removeZeroAndLog a b = unzip . map (\(t,x) -> (t, log x)) . filter ((> 0) . snd ) . zip a $ b

traceValues :: String -> IO ()
traceValues s = do
    let statStyle = LabelStyle 10 Centered PDF.S
    r <- getTrace s
    case r of
        Nothing -> return ()
        Just l -> do
            let h = head l
                sm = toDouble (smallestValue h)
                mi = toDouble (minValue h)
                ma = toDouble (maxValue h)
                triplePartition ta tb !la !lb !lc [] = (la,lb,lc)
                triplePartition ta tb !la !lb !lc (a:h) | ta a = triplePartition ta tb (a:la) lb lc h
                                                        | tb a = triplePartition ta tb la (a:lb) lc h
                                                        | otherwise = triplePartition ta tb la lb (a:lc) h
                (negatives, nulls, positives) = triplePartition (<0) (==0) [] [] [] (map toDouble l)
                positiveSat | null positives = 0 :: Double
                            | otherwise = 100*fromIntegral (length (filter (== ma) positives)) / fromIntegral (length positives)
                negativeSat | null negatives = 0 :: Double
                            | otherwise = 100*fromIntegral (length (filter (== mi) negatives)) / fromIntegral (length negatives)
                nullValues | null nulls = 0 :: Double
                           | otherwise = 100*fromIntegral (length nulls) / fromIntegral (length l)
                drawHist bmi bma l h label theTitle = do
                            let targetLine wi hi (ptF,_) = do
                                    withNewContext $ do
                                        strokeColor $ Rgb 1.0 0.8 0.8
                                        let start = (log bma) - (log bmi)
                                            (xa :+ ya) = ptF ((log bmi),(-start))
                                            (xb :+ yb) = ptF ((log bma),0)
                                        stroke $ Line  xa ya xb yb
                                tickv mav mbv =
                                    map (\t -> (fromIntegral t)/(fromIntegral 10)*(mbv-mav) + mav) ([0..10] :: [Int])
                                tickh _ _ =
                                    let mav = log bmi
                                        mbv = log bma
                                    in
                                    map (\t -> (fromIntegral t)/(fromIntegral 10)*(mbv-mav) + mav) ([0..10] :: [Int])
                                plotStyle =  (defaultPlotStyle { title = theTitle
                                                               , verticalLabel = Just "Log density"
                                                               , horizontalLabel = Just $ "Log amplitude (" ++ label ++ ")"
                                                               , horizontalTickValues = tickh
                                                               , verticalTickValues = tickv
                                                               , epilog = targetLine
                                                               , axis = False
                                                               , horizontalBounds = Just (log bmi, log bma)
                                                               })
                            if null l
                                then do
                                    (\_ -> \_ -> return (R Nothing Nothing), \_ -> \_ -> \_ -> return ())
                                else do
                                    let (t,p) = kde_ 32 (log bmi) (log bma) (U.fromList (map log l))
                                        lt' = U.toList t
                                        lp' = let templ = U.toList p
                                                  s = sum templ
                                              in
                                              map (/ s) templ
                                        (lt,lp) = removeZeroAndLog lt' lp'
                                    let r = (log bma - log bmi) / fromIntegral (length lp)
                                    drawing $ discreteSignalsWithStyle (length lt) plotStyle [ AS (fromListS 0 lp)]
            let (wi,he) = (\(x,y) -> (fromIntegral x, fromIntegral y)) $ viewerSize
            if (signedFormat h)
                then do
                    let posD = drawHist sm ma positives (he / 2 - 10) "pos" (Just "Log density of log amplitude")
                        negD = drawHist sm (-mi) (map negate negatives) (he / 2 - 10) "neg" Nothing
                        pict = Pair posD negD $ \pos -> \neg -> do
                            withNewContext $ do
                                applyMatrix $ translate (0 :+ (he/2 + 10))
                                pos
                            withNewContext $ do
                                neg
                            drawStringLabel statStyle ((printf "negative max: %2.3f %%" negativeSat) :: String)
                                             (leftMargin defaultPlotStyle + (wi / 4.0))  (he / 2 - 10) 100 20
                            drawStringLabel statStyle ((printf "nulls: %2.3f %%" nullValues) :: String)
                                             (leftMargin defaultPlotStyle + (2.0*wi/4.0)) (he/2 - 10) 100 20
                            drawStringLabel statStyle ((printf "positive max: %2.3f %%" positiveSat) :: String)
                                             (leftMargin defaultPlotStyle + (3.0*wi/4.0)) (he/2 - 10) 100 20
                    display pict
                else do
                        let dD = drawHist sm ma positives (he - 40) "unsigned" (Just "Log density of log amplitude")
                            pict = Mono dD $ \d -> do
                                withNewContext $ do
                                    applyMatrix $ translate (0 :+ 40)
                                    d
                                drawStringLabel statStyle ((printf "nulls: %2.3f %%" nullValues) :: String)
                                                    (leftMargin defaultPlotStyle + wi / 3.0) 10 100 20
                                drawStringLabel statStyle ((printf "positive max: %2.3f %%" positiveSat) :: String)
                                                    (leftMargin defaultPlotStyle + (2*wi/3.0)) 10 100 20
                        display pict

{-# NOINLINE traceSample #-}
traceSample :: (HasDoubleRepresentation a, Resolution a) => String -> a -> a
traceSample s a = unsafePerformIO $ do
    nv <- M.insertWith' (++) s [CBT a] <$> (readIORef globalTrace)
    writeIORef globalTrace nv
    return a

{-# NOINLINE trace #-}
trace :: (HasDoubleRepresentation a, Resolution a)
      => String
      -> Signal a
      -> Signal a
trace s = mapS (traceSample s)

{-# NOINLINE myTake #-}
myTake :: NFData a => Int -> [a] -> [a]
myTake 0 l = []
myTake i (h:l) = h:myTake (i-1) l
myTake _ [] = []

{-# NOINLINE forceSignal #-}
forceSignal :: (NFData a, Show a)
            => Int
            -> Signal a
            -> IO ()
forceSignal i (Signal s) = do
    let l = myTake i s
    l `deepseq` return ()


{-# LANGUAGE RecordWildCards #-}

module Main where

import           Control.Monad                  ( when
                                                , unless
                                                )
import qualified Data.Text.IO                  as TIO
import qualified Data.ByteString.Lazy          as BL
                                                ( readFile )
import           Data.Csv                       ( decodeByName )
import           QuoteData
import           Statistics
import           StatReport
import           Charts
import           Params
import           Data.Bool
import           HtmlReport

main :: IO ()
main = cmdLineParser >>= work

work :: Params -> IO ()
work params = do
  csvData <- BL.readFile (fname params)
  case decodeByName csvData of
    Left  err         -> putStrLn err
    Right (_, quotes) -> generateReports params quotes

generateReports :: (Functor t, Foldable t) => Params -> t QuoteData -> IO ()
generateReports Params {..} quotes = do
  TIO.putStr $ statReport statInfo'
  when prices $ plotChart title quotes [Open, Close, High, Low] fname_prices
  when volumes $ plotChart title quotes [Volume] fname_volumes
  unless no_text $ TIO.putStr $ statReport statInfo'
  when html $ BL.writeFile fname_html $ htmlReport title quotes statInfo' images
 where
  statInfo' = statInfo quotes
  withCompany pref = if company /= "" then pref ++ company else ""
  img_suffix    = withCompany "_" ++ ".svg"
  fname_prices  = "prices" ++ img_suffix
  fname_volumes = "volumes" ++ img_suffix
  title         = "Historical Quotes" ++ withCompany " for "
  images        = concat
    $ zipWith (bool []) [[fname_prices], [fname_volumes]] [prices, volumes]
  fname_html = "report" ++ withCompany "_" ++ ".html"


-----------------------------------------------------------------------------
-- |
-- Module      : Berp.Base.Truth
-- Copyright   : (c) 2010 Bernie Pope
-- License     : BSD-style
-- Maintainer  : florbitous@gmail.com
-- Stability   : experimental
-- Portability : ghc
--
-- Implementation of the truth predicate on Python objects.
--
-----------------------------------------------------------------------------

module Berp.Base.Truth (truth) where

import {-# SOURCE #-} Berp.Base.HashTable as HashTable (sizeIO)
import {-# SOURCE #-} Berp.Base.HashSet as HashSet (sizeIO)
import Berp.Base.SemanticTypes (Object (..))
import Data.Complex (Complex (..))
import System.IO.Unsafe (unsafePerformIO)
import Data.IORef (readIORef)

-- XXX incomplete
truth :: Object -> Bool
truth (Bool { object_bool = b }) = b
truth (Integer { object_integer = i }) = i /= 0
truth (Float { object_float = f }) = f /= 0
truth (Complex { object_complex = c }) = c /= (0 :+ 0)
truth None = False
truth obj@(Tuple {}) = object_length obj /= 0
truth obj@(String {}) = object_string obj /= []
truth obj@(List {}) = unsafePerformIO $ do
   numElems <- readIORef $ object_list_num_elements obj
   return (numElems /= 0)
truth obj@(Dictionary {}) = unsafePerformIO $ do
   numElems <- HashTable.sizeIO $ object_hashTable obj
   return (numElems /= 0)
truth obj@(Set {}) = unsafePerformIO $ do
   numElems <- HashSet.sizeIO $ object_hashSet obj
   return (numElems /= 0)
truth _other = True

{-
   Python Language Reference, section 5.10 "Boolean operations":

   In the context of Boolean operations, and also when expressions
   are used by control flow statements, the following values are
   interpreted as false: False, None, numeric zero of all types,
   and empty strings and containers (including strings, tuples,
   lists, dictionaries, sets and frozensets). All other values
   are interpreted as true. User-defined objects can customize their
   truth value by providing a __bool__() method.
-}


{-# language TypeSynonymInstances #-}
{-# language LiberalTypeSynonyms #-}

module FrequencyResponse
  (Analyzer (..),
  AnalyzerAcyclic (..),
  Spectrum (..),
  SpecLin (..),
  SpecWide ,
  getResponse,
  getResponseLin,
  getResponseAcyclic)

where

import CLaSH.Prelude (Default(..))
import Prelude
import CLaSH.Signal.MultiSignal
import Data.Complex as C
import Data.Maybe
import Control.Applicative
import Data.Map.Strict as Map

-- | Type representing input and output signal in circuitry
-- p of kind (* -> *) represents underlaying Signal of type Prependable (ZipList for example)
-- s is type of Spectrum
data Analyzer p s = Analyzer {unAnalyzer :: ([Maybe (s -> s)] , p s)}

instance (Functor p, Prependable p) => Prependable (Analyzer p) where
     prepend a da = Analyzer (Nothing : f , (prepend a $ fmap fx d))
      where
        (f,d) = unAnalyzer da
        fx = fromMaybe id $ foldr1 (<|>) f

instance (Applicative p,Num n) => Num (Analyzer p n) where
    (+) = liftA2FA (+)
    (-) = liftA2FA (-)
    (*) = liftA2FA (*)
    fromInteger = makeAnalyzer . fromInteger
    negate = liftA1FA negate
    abs = liftA1FA abs
    signum = liftA1FA signum

instance (Applicative p,Fractional n) => Fractional (Analyzer p n) where
    (/) = liftA2FA (/)
    fromRational = makeAnalyzer . fromRational

liftA2FA op da db = Analyzer (zipWith (<|>) fa fb, liftA2 op a b) where
       (fa,a) = unAnalyzer da
       (fb,b) = unAnalyzer db

liftA1FA op f = Analyzer (repeat Nothing , fmap op s)   where
   (_,s) = unAnalyzer f

makeAnalyzer s = Analyzer ( repeat Nothing , pure s )
{- inline liftA2FA, liftA1FA-}

-- | spectrum is result of frequency mixing two spectrums
-- n represents type represented as harmonics (Int or better, Rational)
-- and a is type of complex number (frequently Double)
data Spectrum n a = Spectrum {getSpectrum :: Map.Map n (Complex a)}

instance (Ord n,Num n,Default a,RealFloat a) => Default (Spectrum n a) where
    def = dcSpectrum def

instance (Ord n,Num n,RealFloat a) => Num (Spectrum n a) where
    (+) = dotMergeSpec (+)
    (-) = dotMergeSpec (-)
    (*) = mixSpectrum
    negate = fmapSpec negate
    abs = error "no abs function for spectrum available"
    signum = error "no sig num function for spectrum available"
    fromInteger = dcSpectrum . fromInteger

instance (Show n, Show a, Ord n,Num n,Fractional a,RealFloat a) => Fractional (Spectrum n a) where
    (/) = spectrumDiv
    fromRational = dcSpectrum . fromRational

dotMergeSpec op (Spectrum sa) (Spectrum sb) = Spectrum $ (Map.unionWith op) sa sb
fmapSpec f = Spectrum . fmap f . getSpectrum
constSpectrum freq a = Spectrum $ Map.singleton freq (a :+ 0 )
dcSpectrum a = constSpectrum 0 a

mixSpectrum (Spectrum sa) (Spectrum sb) = Spectrum
        $ Map.fromListWith (+)
        $ concat [mixFeq a b | a <- Map.toList sa , b <- Map.toList sb]

mixFeq (fa,pa) (fb,pb)
   | fa > fb = mixFeqSorted (fb,pb) (fa,pa)
   | True =  mixFeqSorted (fa,pa) (fb,pb)

mixFeqSorted (0,pa) (fb,pb) = [(fb,pa*pb)]
mixFeqSorted (fa,pa) (fb,pb) = [(fb - fa,r1),(fb+fa,r2)] where
   (ma,pha) = polar pa
   (mb,phb) = polar pb
   k = ma * mb / 2
   r1
      | fa == fb  = mkPolar (k * sin (phb - pha + pi/2)) 0
      | otherwise = mkPolar k (phb - pha + pi/2)
   r2 = mkPolar ((-1)*k) (phb + pha + pi/2)

delaySpec ph (Spectrum s) = Spectrum $ Map.mapWithKey f s where
    f 0 a = a
    f k a = a * mkPolar 1 (fromRational k * ph)

spectrumDiv (Spectrum a) sb = Spectrum (fmap (/ dc) a) where
    dc = getDcOrError sb

getDcOrError (Spectrum a) =
     case toList a of
            ((0,dc):[]) -> dc
            s  -> error ("spectrum is not dc, but" ++ show s )

makeDcSig :: (Applicative p, Num n, Num a) => a -> Analyzer p (Spectrum n a)
makeDcSig = makeAnalyzer . dcSpectrum

-- | spectrum for linear system, faster version of Spectrum
-- underlaying type is tuple insted of Map
newtype SpecLin = SpecLin {getSpecLin :: (Double,Complex Double)}

instance  Default SpecLin where
    def = SpecLin (def,def :+ def)

instance  Num SpecLin where
    (+) = dotMergeSpecLin (+)
    (-) = dotMergeSpecLin (-)
    (*) = mixSpecLin
    negate = fmapSpecLin negate
    abs = error "no abs function for spectrum available"
    signum = error "no signum function for spectrum available"
    fromInteger = dcSpecLin . fromInteger

instance Fractional SpecLin where
    (/) = undefined -- dcSpecLin
    fromRational = dcSpecLin . fromRational

dotMergeSpecLin op (SpecLin (dcA,acA)) (SpecLin (dcB,acB)) = SpecLin (op dcA dcB , acr) where
    acr = op (realPart acA) (realPart acB) :+ op (imagPart acA) (imagPart acB)

mixSpecLin (SpecLin (dcA,acA)) (SpecLin (dcB,acB))
   | (acA /= 0 ) && (acB /= 0) = error "can not mix two linear spectrum"
   | otherwise = SpecLin (dcr , acr) where
       dcr = dcA * dcB
       acr = fmap (* dcA) acB + fmap (* dcB) acA

dcSpecLin dc = SpecLin (dc,0)

fmapSpecLin f (SpecLin (dc,ac)) =  SpecLin (f dc,fmap f ac)

delayLinSpec ph (SpecLin (dc,ac)) = SpecLin (dc,ac * mkPolar 1 ph)
-- end of SpecLin

-- for acyclic circuits we can implement faster version of Prependable
-- It has nice property that first value of output is correct.
data AnalyzerAcyclic s = AnalyzerAcyclic {unAnalyzerAcyclic :: (Maybe (s -> s), s)}

-- Prependable instance has funny laws (no laws?)
-- it does not follow list isomorphism, because it is single value
instance Prependable AnalyzerAcyclic where
     prepend _ da = AnalyzerAcyclic (f , fx d)
      where
        (f,d) = unAnalyzerAcyclic da
        fx = fromMaybe id f

instance (Num n) => Num (AnalyzerAcyclic n) where
    (+) = opAcyclic (+)
    (-) = opAcyclic (-)
    (*) = opAcyclic (*)
    fromInteger = pureAcyclic . fromInteger
    negate = liftAcyclic negate
    abs = liftAcyclic abs
    signum = liftAcyclic signum

instance (Fractional n) => Fractional (AnalyzerAcyclic n) where
    (/) = opAcyclic (/)
    fromRational = pureAcyclic . fromRational

pureAcyclic n = AnalyzerAcyclic (Nothing,n)
opAcyclic op (AnalyzerAcyclic (fa,a)) (AnalyzerAcyclic (fb,b)) = (AnalyzerAcyclic (fa <|> fb ,op a b))
liftAcyclic f (AnalyzerAcyclic (fa,a)) = (AnalyzerAcyclic (fa,f a))

-- end of AnalyzerAcyclic

type SpecWide = Spectrum Rational Double

-- | run analysis for both acyclic\/cyclic and both linear\/nonlinear circuits
-- It is required to skip first N resulting values
-- where N represent delay of circuitry (approx total number of registers)
getResponse :: (Analyzer ZipList SpecWide -> Analyzer ZipList SpecWide) -- ^ circuitry
            -> Double  -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
            -> [Spectrum Rational Double] -- ^ response as stream
getResponse f ph = getZipList $ snd $ unAnalyzer $ f $ Analyzer ([ Just (delaySpec ph)] , pure (constSpectrum 1 1))

-- | run analysis for both acyclic\/cyclic and only linear circuits
-- faster version of getResponse for acyclic linear ciruits
getResponseLin :: (Analyzer ZipList SpecLin -> Analyzer ZipList SpecLin) -- ^ circuitry
          -> Double -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
          -> [SpecLin] -- ^ response as stream
getResponseLin f ph = getZipList $ snd $ unAnalyzer $ f $ Analyzer ([Just (delayLinSpec ph)] , pure (SpecLin (0,1 :+ 0)))

-- | run analysis both linear\/nonlinear and only acyclic circuit
-- calculate response of acyclic circuit
getResponseAcyclic :: (AnalyzerAcyclic SpecWide -> AnalyzerAcyclic SpecWide) -- ^ circuitry
                      -> Double -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
                      -> Spectrum Rational Double -- ^ response as stream
getResponseAcyclic f ph =  snd $ unAnalyzerAcyclic $ f $ AnalyzerAcyclic (Just (delaySpec ph) ,constSpectrum 1 1)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE ScopedTypeVariables   #-}

module GrenadeExtras.Zip (
  Zip (..),
) where

import           Data.Serialize (Serialize, put, get)

--import           Data.Singletons
import           GHC.TypeLits

import           Grenade.Core (UpdateLayer, Layer, Shape(D1), S(S1D), Gradient,
                               runForwards, runBackwards, runUpdate, createRandom, Tape)
import           GrenadeExtras.GradNorm (GradNorm, normSquared)

import           Numeric.LinearAlgebra.Static ((#), split, R) -- row, (===), splitRows, unrow

data Zip :: Shape -> Shape -> * -> Shape -> Shape -> * -> * where
  Zip :: x -> y -> Zip ix ox x iy oy y

instance (Show x, Show y) => Show (Zip ix ox x iy oy y) where
  show (Zip x y) = "Zip\n" ++ show x ++ "\n" ++ show y

instance (UpdateLayer x, UpdateLayer y) => UpdateLayer (Zip ix ox x iy oy y) where
  type Gradient (Zip ix ox x iy oy y) = (Gradient x, Gradient y)
  runUpdate lr (Zip x y) (x', y') = Zip (runUpdate lr x x') (runUpdate lr y y')
  createRandom = Zip <$> createRandom <*> createRandom

instance ( Layer x ('D1 p) ('D1 m)
         , Layer y ('D1 q) ('D1 n)
         , KnownNat r
         , KnownNat p
         , KnownNat q
         , r ~ (p + q)
         , q ~ (r - p)
         , (p <=? r) ~ 'True
         , KnownNat o
         , KnownNat m
         , KnownNat n
         , o ~ (m + n)
         , n ~ (o - m)
         , (m <=? o) ~ 'True
         ) => Layer (Zip ('D1 p) ('D1 m) x ('D1 q) ('D1 n) y) ('D1 r) ('D1 o) where
   type Tape (Zip ('D1 p) ('D1 m) x ('D1 q) ('D1 n) y) ('D1 r) ('D1 o) = (Tape x ('D1 p) ('D1 m), Tape y ('D1 q) ('D1 n))

   runForwards (Zip x y) (S1D input) =
     let (inputX :: R p, inputY :: R q) = split input
         (xT, xOut :: S ('D1 m)) = runForwards x (S1D inputX)
         (yT, yOut :: S ('D1 n)) = runForwards y (S1D inputY)
     in case (xOut, yOut) of
         (S1D xOut', S1D yOut') ->
             ((xT, yT), S1D (xOut' # yOut'))

   runBackwards (Zip x y) (xTape, yTape) (S1D o) =
     let (ox :: R m , oy :: R n) = split o
         (x', xB :: S ('D1 p)) = runBackwards x xTape (S1D ox)
         (y', yB :: S ('D1 q)) = runBackwards y yTape (S1D oy)
      in case (xB, yB) of
           (S1D xB', S1D yB') ->
             ((x', y'), S1D (xB' # yB'))

instance (Serialize a, Serialize b) => Serialize (Zip sia sa a sib sb b) where
  put (Zip a b) = put a *> put b
  get = Zip <$> get <*> get

instance (GradNorm a, GradNorm b) => GradNorm (Zip sia sa a sib sb b) where
  normSquared (Zip a b) = normSquared a + normSquared b

instance (Num a, Num b) => Num (Zip sia sa a sib sb b) where
  (Zip w1 m1) + (Zip w2 m2) = Zip (w1+w2) (m1+m2)
  (Zip w1 m1) - (Zip w2 m2) = Zip (w1-w2) (m1-m2)
  (Zip w1 m1) * (Zip w2 m2) = Zip (w1*w2) (m1*m2)
  abs    (Zip w m) = Zip (abs w) (abs m)
  signum (Zip w m) = Zip (signum w) (signum m)
  negate (Zip w m) = Zip (negate w) (negate m)
  fromInteger i = Zip (fromInteger i) (fromInteger i)


{-# LANGUAGE Rank2Types, FlexibleContexts #-}
module Numeric.MaxEnt.Deconvolution.Internal where
import Numeric.MaxEnt
import Data.List.Split
import Data.List
import Debug.Trace
import Numeric.Integration.TanhSinh
import Numeric.AD
import Control.Applicative
import Numeric.MaxEnt.Deconvolution.Convolution2D
import Math.GaussianQuadratureIntegration
import Numeric.MaxEnt.Deconvolution.GaussianPSF
import Debug.Trace
import Data.Word
import Data.Array.Repa hiding ((++), map, transpose, toList)
import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.Eval as R
import qualified Data.Array.Repa.Repr.ForeignPtr as R
import qualified Data.Vector.Unboxed as UV
import Control.Lens hiding (ala)
import System.FilePath.Lens
import qualified Data.Vector.Storable as VS
import Foreign.Marshal.Array
import System.IO.Unsafe
import Foreign.ForeignPtr.Safe
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as BS
import Codec.Picture.Png
import Codec.Picture.Types
import Numeric
import Unsafe.Coerce
import qualified Data.Vector.Storable as S
import Numeric.MaxEnt.Deconvolution.Iterative
import Numeric.LinearAlgebra
import Data.Packed.Matrix
--import Numeric.MaxEnt.Deconvolution.IterativeRepa
--import qualified Data.Array.Repa 			        as A
--import qualified Data.Array.Repa.Repr.Unboxed   as U
--import qualified Data.Array.Repa.Stencil		        as A
--import qualified Data.Array.Repa.Stencil.Dim2	        as A
--import Data.Functor.Identity

traceIt x = trace (show x) x

traceItNote :: Show a => String -> a -> a
traceItNote msg x = trace (msg ++ show x) x

gaussianConvolve2D :: Double -> [[Double]] -> [[Double]]
gaussianConvolve2D var xs =
    convolve2DZeroBoundary (sampleGaussian2D var width height) $ xs where
        width = length . head $ xs
        height = length xs

gaussianConvolve2D' :: Double -> [[Double]] -> [[Double]]
gaussianConvolve2D' var xs =
    convolve2DZeroBoundary (sampleGaussian2D' var width height) $ xs where
        width = length . head $ xs
        height = length xs

--deconvolve2D :: [[Double]] -> [Double] -> Either String [[Double]]
deconvolve2D steps psf image = result where
    convoPSF = toConvolve2DZeroBoundary 0 psf

    total = sum . concat $ image
    width = length . head $ image
    height = length image

    normalizedImage  = normalize . concat $ image

    fromLinearImage = chunksOf width . map (total*)

    -- I think here I can
    result = case linear 0.001 (LC convoPSF normalizedImage) of
        Right x -> Right $ fromLinearImage $ S.toList x
        Left x  -> Left $ show x

    --result = Right $ fromLinearImage $ toList $
    --    update steps (fromLists convoPSF) (fromList normalizedImage)

    --result = Right $ fromLinearImage $ A.toList $
    --    runIdentity $ update' steps psf (fromListUnboxed (Z :. width :. height) normalizedImage) $
    --        (delay $ fromListUnboxed (Z :. width :. height) $ replicate (width * height) 0)

    --result = Right $ fromLinearImage $ last $ linear'' (LC convoPSF normalizedImage) :: Either String [[Double]]

--testPSF :: (RealFloat a)
--       => [[a]]
testPSF = gaussianPSF 0.25 5 5

--          -> a -> Int -> Int -> [[a]]
gsamples var width height = gaussianPSF var width height

testInput :: [[Double]]
testInput = [[0,0,0], [0.0, 1.0, 0.0], [0,0,0]]

testInput1 :: [[Double]]
testInput1 = [[0,  0,   0, 0,   0],
              [0,  1.0/9.0,  1.0/9.0, 1.0/9.0,   0],
              [0,  1.0/9.0,  1.0/9.0, 1.0/9.0,   0],
              [0,  1.0/9.0,  1.0/9.0, 1.0/9.0,   0],
              [0,  0,   0, 0,   0]]

testInput2 :: [[Double]]
testInput2 = [[0,  0, 0,   0,   0],
              [0,  0, 0,   0,   0],
              [0,  0, 1.0, 0,   0],
              [0,  0, 0,   0,   0],
              [0,  0, 0,   0,   0]]

testInput3 :: [[Double]]
testInput3 = map (map (1.0-)) testInput2

testDecon :: Either String [[Double]]
testDecon = deconvolve2D 1000 testPSF
    (gaussianConvolve2D  0.25 testInput2)

roundTrip :: Double -> [[Double]] -> Either String [[Double]]
roundTrip var image = deconvolve2D 1000 (gsamples var width height) $
    (gaussianConvolve2D var image) where
        width = length . head $ image
        height = length image

printRT var = either show (show . FloatList) . roundTrip var

toWord8 :: Double -> Word8
toWord8 x = floor $ x * 256.0

--testId ::  (forall a. RealFloat a => [[a]]) -> [[Double]] -> Either String [[Double]]
testId x y = Right y

getR (PixelRGBA8 x _ _ _) = x

toPixel x = [x, x, x, 255]

--TODO make an image that matches the one above
--make increasingly larger examples until I can find the smallest ont
--that exhibits the problem
deconvolve' :: Int -> Double -> Image Pixel8 -> Image Pixel8
deconvolve' steps var i@(Image width height dat) = result where
    input = [[fromIntegral (pixelAt i w h) / 256.0 |
        w <- [0..(width  - 1)]] |
        h <- [0..(height - 1)]  ]

        --
    result = case deconvolve2D steps (gsamples var width height) input of
        Right z -> Image width height . VS.fromList $ concatMap (map toWord8) z
        Left x  -> error ""  --   $ show x

-- The test that I should have, is
-- open a file and blur it
-- and then unblur it
blurImage :: Double -> Image Pixel8 -> Image Pixel8
blurImage var i@(Image width height dat) = result where
    input = [[fromIntegral (pixelAt i w h) / 256.0 |
        w <- [0..(width  - 1)]] |
        h <- [0..(height - 1)]  ]

    blured = gaussianConvolve2D' var input

    result = Image width height $ VS.fromList $ reverse $ concatMap (map toWord8) $ transpose blured

deconvolve :: Int -> Double -> FilePath -> Int -> IO ()
deconvolve steps var filePath i = do
    let outputPath = basename %~ (++ "_decon" ++ show i) $ filePath
    pngBytes <- BS.readFile filePath
    let image = either error id . decodePng $ pngBytes
        newImage = case image of
            ImageY8 x -> deconvolve' steps var x
            x -> error $ "bad format"
    writePng outputPath newImage

deconvolve2 :: Int -> Double -> FilePath -> IO ()
deconvolve2 steps var filePath = do
    let outputPath = basename %~ (++ "_decon") $ filePath
        blurPath   = basename %~ (++ "_blur") $ filePath
    pngBytes <- BS.readFile filePath
    let image = either error id . decodePng $ pngBytes

    newImage <- case image of
        ImageY8 x -> do
            --blur the image and write out the blured one
            let blurImaged = blurImage var x

            writePng blurPath blurImaged
            return $ deconvolve' steps var blurImaged
        x -> error $ "bad format"
    writePng outputPath newImage


-- |
-- Module      : AOC.Challenge.Day11
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Day 11.  See "AOC.Solver" for the types used in this module!

module AOC.Challenge.Day11 (
    day11a
  , day11b
  ) where

import           AOC.Common                       (Point, meanVar)
import           AOC.Solver                       ((:~>)(..))
import           Control.DeepSeq                  (force)
import           Data.Foldable                    (maximumBy)
import           Data.Ix                          (range)
import           Data.Map                         (Map)
import           Data.Maybe                       (catMaybes)
import           Data.Ord                         (comparing)
import           Data.Profunctor                  (dimap)
import           Linear                           (V2(..))
import           Text.Read                        (readMaybe)
import qualified Control.Foldl                    as F
import qualified Data.Map                         as M
import qualified Data.Set                         as S
import qualified Statistics.Distribution          as D
import qualified Statistics.Distribution.Binomial as D
import qualified Statistics.Distribution.Normal   as D

powerLevel :: Int -> Point -> Int
powerLevel sid (V2 x y) = hun ((rid * y + sid) * rid) - 5
  where
    hun = (`mod` 10) . (`div` 100)
    rid = x + 10

findMaxThree :: Map Point Int -> Point
findMaxThree mp = fst
                . maximumBy (comparing snd)
                . map (\x -> (x, go x))
                $ range (V2 1 1, V2 298 298)
  where
    go p = sum [ mp M.! (p + shift)
               | shift <- range (V2 0 0, V2 2 2)
               ]

mkMap :: Int -> Map Point Int
mkMap i = M.fromSet (powerLevel i) . S.fromList $ range (V2 1 1, V2 300 300)

day11a :: Int :~> Point
day11a = MkSol
    { sParse = readMaybe
    , sShow  = \(V2 x y) -> show x ++ "," ++ show y
    , sSolve = Just . findMaxThree . mkMap
    }

day11b :: Int :~> (Point, Int)
day11b = MkSol
    { sParse = readMaybe
    , sShow  = \(V2 x y, s) -> show x ++ "," ++ show y ++ "," ++ show s
    , sSolve = Just . findMaxAny . mkMap
    }

findMaxAny :: Map Point Int -> (Point, Int)
findMaxAny mp = fst $ go 1
  where
    go n
       | goOn > 0.001 = maximumBy (comparing snd) [((pMax, n), oMax), go (n + 1)]
       | otherwise    = ((pMax, n), oMax)
                      -- & traceShow (n, oMax, goOn)
      where
        (pMax, oMax) = maximumBy (comparing snd)
            [ (p, fromIntegral (fromSAT sat p n))
            | p <- range (V2 1 1, V2 (300 - n + 1) (300 - n + 1))
            ]
        goOn = sum
            [ probGreaterThan oMax n'
            | n' <- [n + 1 .. 300]
            ]
    !sat = summedAreaTable mp
    σ :: Double
    σ = sqrt $ ((4 + 5)**2)/12   -- stdev of uniform distribution between -5 and 4
    probGreaterThan o n
        | prob2 == 0 = 0
        | otherwise  = probIn2
      where
        n' = fromIntegral n
        distr2  = D.normalDistr (-(n' ** 2) * 0.5) (n' * σ)
        prob2   = D.complCumulative distr2 o
        numIn2  = ( (300 `div` n) ^ (2 :: Int)    -- we compensate for dependence
                  + (300 - n + 1)^(2 :: Int)
                  ) `div` 2
        probIn2 = 1 - D.probability (D.binomial numIn2 prob2) 0
                    -- this is technically not a bernoulli process,
                    -- because our items are dependent. but we fudge this
                    -- by tweaking the number of trials

fromSAT :: Map Point Int -> Point -> Int -> Int
fromSAT sat (subtract (V2 1 1)->p) n = sum . catMaybes $
    [            M.lookup p            sat
    ,            M.lookup (p + V2 n n) sat
    , negate <$> M.lookup (p + V2 0 n) sat
    , negate <$> M.lookup (p + V2 n 0) sat
    ]

summedAreaTable :: Map Point Int -> Map Point Int
summedAreaTable mp = force sat
  where
    sat = M.mapWithKey go mp
    go p0 v = (+ v) . sum . catMaybes $
      [ negate <$> M.lookup (p0 - V2 1 1) sat
      ,            M.lookup (p0 - V2 1 0) sat
      ,            M.lookup (p0 - V2 0 1) sat
      ]

-- | Debug: find the variance of the map at every square size
_chunkyVars :: Map Point Int -> Map Int Double
_chunkyVars mp = flip M.fromSet (S.fromList [1..300]) $ \n ->
    F.fold (dimap fromIntegral snd meanVar)
      [ fromSAT sat p n
      | p <- range (V2 1 1, V2 (300 - n + 1) (300 - n + 1))
      ]
  where
    !sat = summedAreaTable mp


{-# OPTIONS_GHC -fglasgow-exts #-}
{-# LANGUAGE UndecidableInstances #-}
{- OPTIONS_GHC -fwarn-incomplete-patterns -}

module GeneralizedSignals where

--import
import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Base as SVB
import qualified Data.StorableVector.ST.Strict as SVST
import Foreign.Storable
--import EvalM
import Foreign.Storable.Tuple
import Data.List
import Control.Monad.ST
import Control.Monad
import TNUtils
import VectorsL
import Nats
--import PlotGnuplot
import Data.Ix
import Data.Complex
import Numeric.FFT
import Math.FFT
import qualified Data.Array.CArray as CA
import System.IO.Unsafe

newtype Time = Time Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)
newtype Freq = Freq Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)
newtype Length = Length Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)
newtype SpaceFreq = SpaceFreq Double
    deriving (Eq, Show, Num, Fractional, Floating, Discretizable, IsDouble, Storable)

class (Eq a,Show a,Num a,Fractional a,Floating a)=> IsDouble a where
    toDouble :: a -> Double
    fromDouble :: Double -> a

instance IsDouble Double where
    toDouble = id
    fromDouble = id

class Discretizable a where
    discreteIndex ::  (a,a) -> a -> a -> Int
    discreteRange :: (a,a) -> a -> [a]
    discreteRngLength :: (a,a) -> a -> Int
    foldRangeM :: Monad m => (a,a) -> a -> b -> (a -> b -> m b) -> m b

instance Discretizable Double where
    discreteIndex (lo,hi) dx x = round $ (x-lo)/dx
    discreteRange (lo,hi) dx = [lo, lo+dx..hi]
    discreteRngLength (lo,hi) dx = round $ (hi-lo)/dx
    foldRangeM (lo,hi) dx init f =
        let nmax = discreteRngLength (lo,hi) dx
            go n acc | n > nmax = return acc
                     | otherwise = do
              next <- f ((realToFrac n)*dx+lo) acc
              go (n+1) next
        in go 0 init

instance Discretizable Int where
    discreteIndex (lo,hi) 1 x = x-lo
    discreteIndex (lo,hi) dx x = (x-lo) `div` dx
    discreteRange (lo,hi) dx = [lo,lo+dx..hi]
    discreteRngLength (lo,hi) dx = (hi-lo) `div` dx
    foldRangeM (lo,hi) dx init f =
        let go n acc | n > hi = return acc
                     | otherwise = do
              next <- f n acc
              go (n+dx) next
        in go lo init

{-instance Discretizable a => Discretizable (Vec Z a) where
    discreteIndex (lo,hi) dx x = 0
    discreteRange (lo,hi) dx = [vnil]
    discreteRngLength (lo,hi) dx = 1
    foldRangeM (lo,hi) dx init f = return init -}

instance Discretizable a => Discretizable (Vec (S Z) a) where
    discreteIndex (lo,hi) dx x = discreteIndex (vcar lo,vcar hi) (vcar dx) (vcar x)
    discreteRange (lo,hi) dx = map (`vcons` vnil) $ discreteRange (vcar lo,vcar hi) (vcar dx)
    discreteRngLength (lo,hi) dx = discreteRngLength (vcar lo,vcar hi) (vcar dx)
    foldRangeM (lo,hi) dx init f = foldRangeM (vcar lo,vcar hi) (vcar dx) init f'
        where f' x acc = f (x `vcons` vnil) acc

instance (Discretizable a, Discretizable (Vec (S n) a)) => Discretizable (Vec (S (S n)) a) where
    discreteIndex (vx, vy) (vd) (vw)
        = discreteIndex (vcar vx,vcar vy) (vcar vd) (vcar vw) +
         (discreteIndex (vcar vx,vcar vy) (vcar vd) (vcar vy) + 1) *
         discreteIndex (vcdr vx,vcdr vy) (vcdr vd) (vcdr vw)
    discreteRange (vx, vy) vd
        = [ vcons xys vxsys | vxsys <- discreteRange (vcdr vx, vcdr vy) (vcdr vd),
                              xys <- discreteRange (vcar vx,vcar vy) (vcar vd)]
    discreteRngLength (vx, vy) (vd) = discreteRngLength (vcar vx,vcar vy) (vcar vd) *
                                      discreteRngLength (vcdr vx,vcdr vy) (vcdr vd)
--    foldRangeM :: Monad m => (a,a) -> a -> b -> (a -> b -> m b) -> m b
    foldRangeM (vlo, vhi) vd  init f =
      let f' x acc =
             foldRangeM (vcdr vlo, vcdr vhi) (vcdr vd) acc $ \ix acc' -> f (x `vcons` ix) acc'
      in foldRangeM (vcar vlo, vcar vhi) (vcar vd) init f'

f v (a, i) = return $ (v:a, i+1)

tf = print =<< foldRangeM (0::Vec Three Int, 2) (1) ([],0) f

tf2 = discreteIndex (0::Vec Three Int, 2) (1) (2 `vcons` 2 `vcons` 2 `vcons` vnil)

{-$ \x-> do
         let thisf ix acc = f (x `vcons` ix) acc
         foldRangeM (vcdr vlo, vcdr vhi) (vcdr vd) init thisf -}

type family Volume a :: *
type instance Volume Time = (Time,Time)
type instance Volume Freq = (Freq,Freq)

--type family FInverses a b :: * -> *

--type instance FInverses Time Freq

--type family FInv a :: *

--type instance FInv Time = Freq
--type instance (FInv t ~ s) => FInv s = t

{-class Inverses a b | a -> b, b-> a where
    invertFwd :: a -> b
    invertBack :: b -> a

instance Inverses Time Freq where
    invertFwd (Time t) = Freq $ recip t
    invertBack (Freq f) = Time $ recip f

instance Inverses Length SpaceFreq where
    invertFwd (Length t) = SpaceFreq $ recip t
    invertBack (SpaceFreq f) = Length $ recip f-}

class HasInverse a where
      type Inv a :: *
      invertValue :: a -> Inv a

instance HasInverse Int where
      type Inv Int = Int
      invertValue = negate -- haha. FIXME

instance HasInverse Time where
      type Inv Time = Freq
      invertValue (Time t) = Freq $ recip t

instance HasInverse Freq where
      type Inv Freq = Time
      invertValue (Freq t) = Time $ recip t

instance HasInverse Length where
      type Inv Length = SpaceFreq
      invertValue (Length t) = SpaceFreq $ recip t

instance HasInverse SpaceFreq where
      type Inv SpaceFreq = Length
      invertValue (SpaceFreq t) = Length $ recip t

instance HasInverse a => HasInverse (Vec n a) where
      type Inv (Vec n a) = Vec n (Inv a)
      invertValue = fmap invertValue

data Signal a b where
    Signal  :: (Discretizable a, {-Ix (Discretized a),-} Storable b, Storable a) =>
               a -> (a,a) -> (Inv a, Inv a) -> SV.Vector b -> Signal a b
    SigFmap :: (b->b') -> Signal a b -> Signal a b'
    SigFun  :: (a->b) -> Signal a b

sigLims :: Signal a b -> (a,a)
sigLims (Signal d1 lims1 invlims1 arr1) = lims1
sigLims (SigFmap _ s) = sigLims s

instance Functor (Signal a) where
    fmap f (SigFmap g s) = SigFmap (f . g) s
    fmap f s = SigFmap f s

sOfVToVofS :: Signal a (Vec n b) -> Vec n (Signal a b)
sOfVToVofS s = undefined

vOfSToSofV :: Vec n (Signal a b) -> Signal a (Vec n b)
vOfSToSofV voss = let sigs = vecToList voss
                  in undefined

type Events a b = [(a,b)]
type Region a b = [(Volume a,b)]

at :: Signal a b -> a -> b
at (Signal delta lims _ arr) x =  arr `SV.index` (discreteIndex lims delta x)
at (SigFmap f s) x = f $ s `at` x
at (SigFun f) x = f x

fill :: (Discretizable a, Storable b, Storable a) =>
        a -> (a,a) -> (Inv a, Inv a) -> (a->b) -> Signal a b
fill delta lims invlims f =
     Signal delta lims invlims $ SV.pack $ map f $ discreteRange lims delta

trans :: Storable c => (a->c) -> Signal a b -> Signal a c
trans f (Signal delta lims ilims arr) = Signal delta lims ilims newarr where
    newarr = SV.pack $ map f $ discreteRange lims delta

trans' :: Storable c => (a->b->c) -> Signal a b -> Signal a c
trans' f (Signal delta lims ilims arr) = Signal delta lims ilims newarr where
    newarr = SV.pack $ map (uncurry f) $ zip (discreteRange lims delta) (SV.unpack arr)
--    newarr = SV.zipWith f (SV.pack $ discreteRange lims delta) (arr)

--DOESNT WORK -- problem may be in foldRangeM implementation
fillIO' :: (Discretizable a, Storable b, Show a, Storable a) =>
          a -> (a,a) -> (Inv a, Inv a) -> (a->IO b) -> IO (Signal a b)
fillIO' delta lims invlims mf = do
  let n = discreteRngLength lims delta
  print n
  stv <- stToIO $ SVST.new_ (n+1800) --FUDGE
  let f ixv nix = do v <- mf ixv
                     --print (ixv, nix)
                     stToIO $ SVST.write stv nix v
                     return $ nix+1
  _ <- foldRangeM lims delta 0 f
  sv <- stToIO $ SVST.unsafeFreeze stv
  return $ Signal delta lims invlims sv

fillIO :: (Discretizable a, Storable b, Show a, Storable a) =>
          a -> (a,a) -> (Inv a, Inv a) -> (a->IO b) -> IO (Signal a b)
fillIO delta lims invlims mf = do
  let rng = discreteRange lims delta
  let n = discreteRngLength lims delta
  xs <- mapM mf rng
  return $ Signal delta lims invlims $ SV.pack xs

zipSignalsWith :: Storable c => (a -> b -> c) -> Signal i a -> Signal i b -> Signal i c
zipSignalsWith f (Signal d1 lims1 invlims1 arr1) (Signal d2 lims2 invlims2 arr2)
               = Signal d1 lims1 invlims2 $ SV.zipWith f arr1 arr2

--toCArray :: Signal a b -> CArray

class X a b where
    fourier :: Signal a b -> Signal (Inv a) b

instance X Time (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ fft $ SV.unpack arr

instance X Freq (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ ifft $ SV.unpack arr

instance X Length (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ fft $ SV.unpack arr

instance X SpaceFreq (Complex Double) where
    fourier (Signal d lims invlims arr) = Signal (invertValue d) invlims lims $ SV.pack $ ifft $ SV.unpack arr

--instance X (Vec Z Length) (Complex Double) where
--    fourier = undefined

instance (Nat n, X (Vec n Length) (Complex Double), Storable (Vec n SpaceFreq),
          Discretizable (Vec n SpaceFreq), Ix (Vec n Int), CA.Shapable (Vec n Int))
    => X (Vec n Length) (Complex Double) where
    fourier s@(Signal d lims invlims arr) =
        let ixdims :: Nat n => Signal (Vec n a) b -> n
            ixdims = undefined
            dims = toInt $ ixdims s
            (oldp, oldn, oldsomething) = SVB.toForeignPtr arr
            cixs = discreteLims s
            carray = unsafePerformIO (CA.unsafeForeignPtrToCArray oldp cixs)
            newcarray = dftN [0..dims] carray
            (newn, newp) = CA.toForeignPtr newcarray
            newarr = SVB.fromForeignPtr newp newn
        in Signal (invertValue d) invlims lims newarr

--instance IsDouble d => Shapable (Vec Z d) where

discreteLims :: (Nat n, IsDouble d) => Signal (Vec n d) a -> (Vec n Int, Vec n Int)
discreteLims s@(Signal d (lo,hi) invlims arr)
    = ( fmap (round . toDouble) $ lo/d,  fmap (round . toDouble) $ hi/d)

instance (Storable (b,c), X a b, X a c) => X a (b,c) where
    fourier s = zipSignalsWith (,) (fourier $ fmap fst s) (fourier $ fmap snd s)

instance (X a b) => X a (Vec n b) where
    fourier sOfVecs =  vOfSToSofV $ fmap fourier $ sOfVToVofS sOfVecs

class XRC a where
    fourierRC :: Signal a Double -> Signal (Inv a) (Complex Double)
    fourierCR :: Signal (Inv a) (Complex Double) -> Signal a Double


{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE ScopedTypeVariables  #-}

module Genetic where

import           Numeric.LinearAlgebra
import           System.Random.Shuffle    (shuffleM)
import           System.Random
import           System.ProgressBar
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Concurrent.Async (mapConcurrently)
import qualified Data.Vector as V
import           Data.List
import           Data.Maybe

class (Ord a, Eq a, Show a) => Individual a where
    fitness :: a -> IO (Maybe Double)
    mutateIndividual :: a -> IO a
    crossover :: (a, a) -> IO a

runGeneticAlgorithm :: Individual a
                    => Double      -- Goal
                    -> Int         -- maxNumIterations
                    -> Int         -- eliteSize
                    -> [a]         -- Initial (random) genetation
                    -> ProgressRef -- Progress Bar
                    -> IO a
runGeneticAlgorithm threshold
                    maxNumIterations
                    eliteSize
                    population
                    progRef   = do
    finalGroup <- foldM (makeGeneration threshold eliteSize progRef)
                        initialGeneration
                        [1..maxNumIterations]
    best finalGroup
    where
        initialGeneration = population

makeGeneration :: Individual a
               => Double      -- Goal
               -> Int         -- eliteSize
               -> ProgressRef
               -> [a]         -- population
               -> Int         -- generation number
               -> IO [a]
makeGeneration goal eliteSize progRef population _ = do
    cleanPop <- filterM (fmap isJust . fitness) population
    mFitness <- fitness =<< best cleanPop
    let bestFitness = fromMaybe (error "Invalid Fitness") mFitness
    if bestFitness < goal
    then return cleanPop
    else do incProgress progRef 1
            nextGeneration cleanPop eliteSize

nextGeneration :: Individual a
               => [a]    -- population
               -> Int    -- elite size
               -> IO [a]
nextGeneration population eliteSize = do
    children    <- makeChildren population eliteSize
    elita <- elite population eliteSize
    return (elita ++ children)

best :: Individual a
     => [a]
     -> IO a
best pop = do
    fitnesses <- mapConcurrently fitness pop
    let combined = filter (isJust . fst) $ zip fitnesses pop
    --             [(fitness, individual)]
        chosen   = minimumBy (\x y -> compare (fst x) (fst y)) combined
    return $ snd chosen

mutate :: Individual a => [a] -> IO [a]
mutate = mapConcurrently mutateIndividual

elite :: Individual a
      => [a]
      -> Int
      -> IO [a]
elite individuals nElit = do
    fitnesses <- mapConcurrently fitness individuals
    let combined   = zip fitnesses individuals
        results    = map snd . take nElit . sortOn fst $ combined
        mFitnesses = map fst . take nElit . sortOn fst $ combined
    return results

selectCouples :: Individual a
           => [a] -- population
           -> Int -- number of couples
           -> IO [(a, a)]
selectCouples population numOfCouples =
    sequence . take numOfCouples $ repeat $ selectCouple population

selectCouple :: Individual a
             => [a]
             -> IO (a, a)
selectCouple population = do
    fitnesses <- mapConcurrently fitness population
    let combined    = filter (isJust . fst) $ zip fitnesses population
        population_ = map snd . sortOn (fromJust . fst) $ combined
        indexList   = [0..(length population_ - 1)]
        duplicated  = duplicateByValue indexList
    [index1, index2] <- take 2
                      . nub
                    <$> shuffleM duplicated
    return ( population_ !! index1
           , population_ !! index2
           )

duplicateByValue :: [Int]
                 -> [Int]
duplicateByValue = foldr (\ x -> (++) (replicate x x)) []

makeChildren :: Individual a
             => [a]    -- population
             -> Int    -- nElit
             -> IO [a]
makeChildren population nElit = do
    parents   <- selectCouples population $ length population - nElit
    children  <- mapConcurrently crossover parents
    fitnesses <- mapConcurrently fitness children
    mutate $ map fst $ filter (isJust . snd) $ zip children fitnesses

bestFit :: Individual a
        => [a]
        -> IO Double
bestFit pop = minimum . catMaybes <$> mapConcurrently fitness pop


{-# LANGUAGE FlexibleContexts #-}
module Statistics.Distribution.Random.Gamma (
   gamma
) where

{-
 -  Shape parameter a >= 1.  Algorithm GD in:
 -
 -  Ahrens, J.H. and Dieter, U. (1982).
 -  Generating gamma variates by a modified
 -  rejection technique.
 -  Comm. ACM, 25, 47-54.
 -
 -
 -  Shape parameter 0 < a < 1. Algorithm GS in:
 -
 -  Ahrens, J.H. and Dieter, U. (1974).
 -  Computer methods for sampling from gamma, beta,
 -  poisson and binomial distributions.
 -  Computing, 12, 223-246.
 -}

import qualified Statistics.Distribution.Random.Exponential as D
import qualified Statistics.Distribution.Random.Normal as D
import qualified Statistics.Distribution.Random.Uniform as D
import           Data.List (foldl')
import           Data.Number.LogFloat (expm1)
import           Data.Word
import           Random.CRI

sqrt32, exp_m1 :: Double
sqrt32 = sqrt 32
{-# NOINLINE sqrt32 #-}
exp_m1 = exp (-1)
{-# NOINLINE exp_m1 #-}

q1, q2, q3, q4, q5, q6, q7 :: Double
q1 =  0.04166669
q2 =  0.02083148
q3 =  0.00801191
q4 =  0.00144121
q5 = -7.388e-5
q6 =  2.4511e-4
q7 =  2.424e-4
qs :: [Double]
qs = [q7, q6, q5, q4, q3, q2, q1]
{-# INLINE qs #-}

a1, a2, a3, a4, a5, a6, a7 :: Double
a1 =  0.3333333
a2 = -0.250003
a3 =  0.2000062
a4 = -0.1662921
a5 =  0.1423657
a6 = -0.1367177
a7 =  0.1233795
as :: [Double]
as = [a7, a6, a5, a4, a3, a2, a1]
{-# INLINE as #-}

horner :: [Double] -> Double -> Double
horner q r = foldl' (\ a b -> (a + b) * r) 0 q
{-# INLINE horner #-}

gamma :: (Source m g Double, Source m g Word32) => Double -> Double -> g m -> m Double
gamma shape scale rng
    | shape == 0    =  return 0
    | shape <  1    =  gammaGS shape scale rng
    | otherwise     =  gammaGD shape scale rng

gammaGS :: (Source m g Double, Source m g Word32) => Double -> Double -> g m -> m Double
gammaGS shape scale rng =
    let e  = 1 + exp_m1 * shape
        go = do
            ru <- D.uniform     rng
            re <- D.exponential rng
            let p                  = e * ru
                x      | p >= 1    = - log ((e - p) / shape)
                       | otherwise = exp (log p / shape)
                accept | p >= 1    = re >= (1 - shape) * log x
                       | otherwise = re >= x
            if accept then return (scale * x) else go
    in go

gammaGD :: (Source m g Double, Source m g Word32) => Double -> Double -> g m -> m Double
gammaGD shape scale rng =

    let s2 = shape - 0.5
        s  = sqrt s2
        d  = sqrt32 - s * 12
        r  = 1 / shape
        q0 = horner qs r

        (b, si, c)
            | shape <=  3.686 = (0.463 + s + 0.178 * s2, 1.235,            0.195 / s - 0.079 + 0.16 * s)
            | shape <= 13.022 = (1.654 + 0.0076 * s2,    1.68 / s + 0.275, 0.062 / s + 0.024)
            | otherwise       = (1.77,                   0.75,             0.1515 / s)

        calc_q t
            | abs v <= 0.25   = q0 + 0.5 * t * t * horner as v
            | otherwise       = q0 - s * t + 0.25 * t * t + (s2 + s2) * log (1 + v)
            where v = t / (s + s)

        choose_t = do
            e  <- D.exponential rng
            u' <- D.uniform rng
            let uu = u' + u' - 1
            let tt = if uu < 0 then b - si * e else b + si * e
            if tt >= -0.71874483771719
                then do
                    let qq = calc_q tt
                    if qq > 0 && c * abs uu <= expm1 qq * exp (e - 0.5 * tt * tt)
                        then return tt
                        else choose_t
                else choose_t

    in do
        t <- D.normal rng
        let x = s + 0.5 * t
            ret_val = scale * x * x
        if t >= 0
            then return ret_val
            else do

                u <- D.uniform rng
                if d * u <= t * t * t
                    then return ret_val
                    else do

                        let q = calc_q t
                        if x > 0 && log (1 - u) <= q
                            then return ret_val
                            else do
                                ttt <- choose_t
                                let xx = s + 0.5 * ttt
                                return $ scale * xx * xx


{-# LANGUAGE RecordWildCards #-}

module Main where

import           Control.DeepSeq                  (force)
import           Control.Exception                (evaluate)
import           Control.Monad                    (replicateM)
import           Data.List                        (minimumBy)
import           Data.Ord                         (comparing)
import qualified Data.Text                        as T
import qualified Data.Vector.Unboxed              as V
import           GHC.Base                         (String)
import           Options.Applicative              (Parser (..), auto,
                                                   execParser, fullDesc, help,
                                                   helper, info, long, metavar,
                                                   option, progDesc,
                                                   showDefault, strOption,
                                                   value, (<**>))
import qualified Statistics.Distribution.StudentT as S
import qualified Statistics.Function              as S
import qualified Statistics.Sample                as S
import qualified Statistics.Test.StudentT         as S
import qualified Statistics.Test.Types            as ST
import           System.IO                        (hGetContents)
import           System.Process                   (CreateProcess (..),
                                                   StdStream (..), shell,
                                                   waitForProcess,
                                                   withCreateProcess)
import           Text.Printf                      (printf)

type Samples = V.Vector Double

data Args = Args
  { old   :: T.Text
  , new   :: T.Text
  , forks :: Int
  }

data Diff = Diff
  { meanPct :: Double
  , oldMean :: Double
  , oldMin  :: Double
  , oldMax  :: Double
  , newMin  :: Double
  , newMax  :: Double
  , newMean :: Double
  , test    :: Maybe (ST.Test S.StudentT) }
  deriving (Show)

args :: Parser Args
args = Args
    <$> strOption
        ( long "old"
        <> metavar "OLD"
        <> help "Command that returns samples for \"old\"" )
    <*> strOption
        ( long "new"
        <> metavar "NEW"
        <> help "Command that returns samples for \"new\"" )
    <*> option auto
        ( long "forks"
        <> help "How often the commands will be invoked to gather samples"
        <> showDefault
        <> value 1
        <> metavar "FORKS" )

execProc :: String -> IO Samples
execProc cmd =
  withCreateProcess createProc $ \_ (Just out) _ ph -> do
    samples <- textToSamples <$> hGetContents out
    evaluate $ force samples
    ex <- waitForProcess ph
    pure samples
  where
    textToSamples = V.fromList . map read . lines
    createProc = (shell cmd) { std_out = CreatePipe }

calcDiff :: Samples -> Samples -> Diff
calcDiff oldSamples newSamples = Diff{..}
  where
    (oldMin, oldMax) = S.minMax oldSamples
    (newMin, newMax) = S.minMax newSamples
    oldMean = S.mean oldSamples
    newMean = S.mean newSamples
    oldStdDev = S.stdDev oldSamples
    newStdDev = S.stdDev newSamples
    meanPct = percDiff oldMean newMean
    test = S.welchTTest S.BGreater oldSamples newSamples

percDiff :: Double -> Double -> Double
percDiff x y = (abs (x - y) / ((x + y) / 2)) * 100

cmpPerf :: Args -> IO ()
cmpPerf (Args old new forks) = do
  listOfOldSamples <- replicateM forks $ do
    putStrLn ("Gathering metrics for: " <> T.unpack old)
    execProc $ T.unpack old
  listOfNewSamples <- replicateM forks $ do
    putStrLn ("Gathering metrics for: " <> T.unpack new)
    execProc $ T.unpack new
  let
    bestOldSamples = minimumBy (comparing S.mean) listOfOldSamples
    bestNewSamples = minimumBy (comparing S.mean) listOfNewSamples
    diff = calcDiff bestOldSamples bestNewSamples
  printf   "%.2f%% mean difference\n" (meanPct diff)
  putStrLn " Cmd |       Mean |        Min |        Max |"
  printf   " Old | %10.3f | %10.3f | %10.3f |\n" (oldMean diff) (oldMin diff) (oldMax diff)
  printf   " New | %10.3f | %10.3f | %10.3f |\n" (newMean diff) (newMin diff) (newMax diff)

main :: IO ()
main = cmpPerf =<< execParser opts
  where
    opts = info (args <**> helper)
      ( fullDesc
      <> progDesc "Compare the (runtime) values returned by two different programs" )


-- |
-- Module      : Occlusion.Behaviours
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 25 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------
{-# LANGUAGE Rank2Types #-}

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Occlusion.Behaviours where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import Data.Maybe   (fromMaybe)

import Control.Lens

import Southpaw.Math.Constants
import Southpaw.Math.Trigonometry

import           Occlusion.Types
import qualified Occlusion.Lenses as L
import qualified Occlusion.Core   as Core

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------

-- |
player :: Behaviour Character AppState Character
player self appstate = self & foldr (.) id [L.position %~ (+(dt*v)), L.velocity .~ currentVelocity appstate]
  where
    dt = (1.0/appstate^.L.animation.L.fps):+0
    v  = self^.L.velocity
  -- modifyIORef stateref (\appstate -> appstate & scene.player.velocity .~ currentVelocity appstate) -- TODO: This is ugly

-- |
npc :: [Complex Double] -> Behaviour Character AppState Character
npc path self appstate = self & L.position %~ flip fromMaybe (Core.walkalong path (n*v*dt))
  where
    dt = (1.0/appstate^.L.animation.L.fps)
    v  = realPart . abs $ self^.L.velocity
    n  = fromIntegral $ appstate^.L.animation.L.frame

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

-- |
-- TODO: Make polymorphic
run :: Lens AppState AppState Character Character -> AppState -> AppState
run char appstate = appstate & char .~ (appstate^.char.L.behaviour) (appstate^.char) appstate

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

-- |
-- TODO: Move
currentVelocity :: AppState -> Complex Double
currentVelocity appstate
  | appstate^.L.input.L.click == Nothing = 0.0:+0.0
  | abs delta^.L.real < 12.0             = 0.0
  | otherwise                            = mkPolar 64.0 (phase delta)
  where
    delta = appstate^.L.input.L.mouse - appstate^.L.scene.L.player.L.position


module TestRotation where

import Cube (Cube)
import Numeric.LinearAlgebra ((<>))
import Rotation
import qualified Cube as Cube
import qualified Data.List as List

testRotations :: MatrixRotation -> Bool
testRotations m =
  let ms = take 4 $ iterate (<> m) m
      areIdentity = fmap (== identity) ms
      testAreIdentity = areIdentity == [False, False, False, True]
      testAllDifferent = (4 ==) . length $ List.nub ms
  in testAreIdentity && testAllDifferent

testTopToFront :: Bool
testTopToFront =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, b), (top, 0, 1, b), (top, 0, 2, b), (top, 1, 0, b), (top, 1, 1, b), (top, 1, 2, b), (top, 2, 0, b), (top, 2, 1, b), (top, 2, 2, b)
                 , (frn, 0, 0, w), (frn, 0, 1, w), (frn, 0, 2, w), (frn, 1, 0, w), (frn, 1, 1, w), (frn, 1, 2, w), (frn, 2, 0, w), (frn, 2, 1, w), (frn, 2, 2, w)
                 , (rgt, 0, 0, r), (rgt, 0, 1, r), (rgt, 0, 2, r), (rgt, 1, 0, r), (rgt, 1, 1, r), (rgt, 1, 2, r), (rgt, 2, 0, r), (rgt, 2, 1, r), (rgt, 2, 2, r)
                 , (bck, 0, 0, y), (bck, 0, 1, y), (bck, 0, 2, y), (bck, 1, 0, y), (bck, 1, 1, y), (bck, 1, 2, y), (bck, 2, 0, y), (bck, 2, 1, y), (bck, 2, 2, y)
                 , (lft, 0, 0, o), (lft, 0, 1, o), (lft, 0, 2, o), (lft, 1, 0, o), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, o), (lft, 2, 1, o), (lft, 2, 2, o)
                 , (bot, 0, 0, g), (bot, 0, 1, g), (bot, 0, 2, g), (bot, 1, 0, g), (bot, 1, 1, g), (bot, 1, 2, g), (bot, 2, 0, g), (bot, 2, 1, g), (bot, 2, 2, g) ]
      rotated = rotate topToFront Cube.solvedCube
  in expected == rotated
  && testRotations topToFront

getSideColors :: Cube -> [(Cube.Side, Cube.Color)]
getSideColors = fmap head . filter (not . null) . List.group . fmap (\(s,_,_,c) -> (s,c))

testTopToBack :: Bool
testTopToBack = testRotations topToBack

testTopToRight :: Bool
testTopToRight =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, o), (top, 0, 1, o), (top, 0, 2, o), (top, 1, 0, o), (top, 1, 1, o), (top, 1, 2, o), (top, 2, 0, o), (top, 2, 1, o), (top, 2, 2, o)
                 , (frn, 0, 0, g), (frn, 0, 1, g), (frn, 0, 2, g), (frn, 1, 0, g), (frn, 1, 1, g), (frn, 1, 2, g), (frn, 2, 0, g), (frn, 2, 1, g), (frn, 2, 2, g)
                 , (rgt, 0, 0, w), (rgt, 0, 1, w), (rgt, 0, 2, w), (rgt, 1, 0, w), (rgt, 1, 1, w), (rgt, 1, 2, w), (rgt, 2, 0, w), (rgt, 2, 1, w), (rgt, 2, 2, w)
                 , (bck, 0, 0, b), (bck, 0, 1, b), (bck, 0, 2, b), (bck, 1, 0, b), (bck, 1, 1, b), (bck, 1, 2, b), (bck, 2, 0, b), (bck, 2, 1, b), (bck, 2, 2, b)
                 , (lft, 0, 0, y), (lft, 0, 1, y), (lft, 0, 2, y), (lft, 1, 0, y), (lft, 1, 1, y), (lft, 1, 2, y), (lft, 2, 0, y), (lft, 2, 1, y), (lft, 2, 2, y)
                 , (bot, 0, 0, r), (bot, 0, 1, r), (bot, 0, 2, r), (bot, 1, 0, r), (bot, 1, 1, r), (bot, 1, 2, r), (bot, 2, 0, r), (bot, 2, 1, r), (bot, 2, 2, r) ]
      rotated = rotate topToRight Cube.solvedCube
  in expected == rotated
  && testRotations topToRight

testTopToLeft :: Bool
testTopToLeft = testRotations topToLeft

testTopTwistRight :: Bool
testTopTwistRight =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, w), (top, 0, 1, w), (top, 0, 2, w), (top, 1, 0, w), (top, 1, 1, w), (top, 1, 2, w), (top, 2, 0, w), (top, 2, 1, w), (top, 2, 2, w)
                 , (frn, 0, 0, r), (frn, 0, 1, r), (frn, 0, 2, r), (frn, 1, 0, r), (frn, 1, 1, r), (frn, 1, 2, r), (frn, 2, 0, r), (frn, 2, 1, r), (frn, 2, 2, r)
                 , (rgt, 0, 0, b), (rgt, 0, 1, b), (rgt, 0, 2, b), (rgt, 1, 0, b), (rgt, 1, 1, b), (rgt, 1, 2, b), (rgt, 2, 0, b), (rgt, 2, 1, b), (rgt, 2, 2, b)
                 , (bck, 0, 0, o), (bck, 0, 1, o), (bck, 0, 2, o), (bck, 1, 0, o), (bck, 1, 1, o), (bck, 1, 2, o), (bck, 2, 0, o), (bck, 2, 1, o), (bck, 2, 2, o)
                 , (lft, 0, 0, g), (lft, 0, 1, g), (lft, 0, 2, g), (lft, 1, 0, g), (lft, 1, 1, g), (lft, 1, 2, g), (lft, 2, 0, g), (lft, 2, 1, g), (lft, 2, 2, g)
                 , (bot, 0, 0, y), (bot, 0, 1, y), (bot, 0, 2, y), (bot, 1, 0, y), (bot, 1, 1, y), (bot, 1, 2, y), (bot, 2, 0, y), (bot, 2, 1, y), (bot, 2, 2, y) ]
      rotated = rotate topTwistRight Cube.solvedCube
  in expected == rotated
  && testRotations topTwistRight

testTopTwistLeft :: Bool
testTopTwistLeft = testRotations topTwistLeft

testRightL :: Bool
testRightL =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, w), (top, 0, 1, w), (top, 0, 2, b), (top, 1, 0, w), (top, 1, 1, w), (top, 1, 2, b), (top, 2, 0, w), (top, 2, 1, w), (top, 2, 2, b)
                 , (frn, 0, 0, g), (frn, 0, 1, g), (frn, 0, 2, w), (frn, 1, 0, g), (frn, 1, 1, g), (frn, 1, 2, w), (frn, 2, 0, g), (frn, 2, 1, g), (frn, 2, 2, w)
                 , (rgt, 0, 0, r), (rgt, 0, 1, r), (rgt, 0, 2, r), (rgt, 1, 0, r), (rgt, 1, 1, r), (rgt, 1, 2, r), (rgt, 2, 0, r), (rgt, 2, 1, r), (rgt, 2, 2, r)
                 , (bck, 0, 0, y), (bck, 0, 1, b), (bck, 0, 2, b), (bck, 1, 0, y), (bck, 1, 1, b), (bck, 1, 2, b), (bck, 2, 0, y), (bck, 2, 1, b), (bck, 2, 2, b)
                 , (lft, 0, 0, o), (lft, 0, 1, o), (lft, 0, 2, o), (lft, 1, 0, o), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, o), (lft, 2, 1, o), (lft, 2, 2, o)
                 , (bot, 0, 0, y), (bot, 0, 1, y), (bot, 0, 2, g), (bot, 1, 0, y), (bot, 1, 1, y), (bot, 1, 2, g), (bot, 2, 0, y), (bot, 2, 1, y), (bot, 2, 2, g) ]
      rotated = rotate rightL Cube.solvedCube
  in expected == rotated
  && testRotations rightL

testRightR :: Bool
testRightR = testRotations rightR

testLeftL :: Bool
testLeftL =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, g), (top, 0, 1, w), (top, 0, 2, w), (top, 1, 0, g), (top, 1, 1, w), (top, 1, 2, w), (top, 2, 0, g), (top, 2, 1, w), (top, 2, 2, w)
                 , (frn, 0, 0, y), (frn, 0, 1, g), (frn, 0, 2, g), (frn, 1, 0, y), (frn, 1, 1, g), (frn, 1, 2, g), (frn, 2, 0, y), (frn, 2, 1, g), (frn, 2, 2, g)
                 , (rgt, 0, 0, r), (rgt, 0, 1, r), (rgt, 0, 2, r), (rgt, 1, 0, r), (rgt, 1, 1, r), (rgt, 1, 2, r), (rgt, 2, 0, r), (rgt, 2, 1, r), (rgt, 2, 2, r)
                 , (bck, 0, 0, b), (bck, 0, 1, b), (bck, 0, 2, w), (bck, 1, 0, b), (bck, 1, 1, b), (bck, 1, 2, w), (bck, 2, 0, b), (bck, 2, 1, b), (bck, 2, 2, w)
                 , (lft, 0, 0, o), (lft, 0, 1, o), (lft, 0, 2, o), (lft, 1, 0, o), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, o), (lft, 2, 1, o), (lft, 2, 2, o)
                 , (bot, 0, 0, b), (bot, 0, 1, y), (bot, 0, 2, y), (bot, 1, 0, b), (bot, 1, 1, y), (bot, 1, 2, y), (bot, 2, 0, b), (bot, 2, 1, y), (bot, 2, 2, y) ]
      rotated = rotate leftL Cube.solvedCube
  in expected == rotated
  && testRotations leftL

testLeftR :: Bool
testLeftR = testRotations leftR

testTopL :: Bool
testTopL =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, w), (top, 0, 1, w), (top, 0, 2, w), (top, 1, 0, w), (top, 1, 1, w), (top, 1, 2, w), (top, 2, 0, w), (top, 2, 1, w), (top, 2, 2, w)
                 , (frn, 0, 0, o), (frn, 0, 1, o), (frn, 0, 2, o), (frn, 1, 0, g), (frn, 1, 1, g), (frn, 1, 2, g), (frn, 2, 0, g), (frn, 2, 1, g), (frn, 2, 2, g)
                 , (rgt, 0, 0, g), (rgt, 0, 1, g), (rgt, 0, 2, g), (rgt, 1, 0, r), (rgt, 1, 1, r), (rgt, 1, 2, r), (rgt, 2, 0, r), (rgt, 2, 1, r), (rgt, 2, 2, r)
                 , (bck, 0, 0, r), (bck, 0, 1, r), (bck, 0, 2, r), (bck, 1, 0, b), (bck, 1, 1, b), (bck, 1, 2, b), (bck, 2, 0, b), (bck, 2, 1, b), (bck, 2, 2, b)
                 , (lft, 0, 0, b), (lft, 0, 1, b), (lft, 0, 2, b), (lft, 1, 0, o), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, o), (lft, 2, 1, o), (lft, 2, 2, o)
                 , (bot, 0, 0, y), (bot, 0, 1, y), (bot, 0, 2, y), (bot, 1, 0, y), (bot, 1, 1, y), (bot, 1, 2, y), (bot, 2, 0, y), (bot, 2, 1, y), (bot, 2, 2, y) ]
      rotated = rotate topL Cube.solvedCube
  in expected == rotated
  && testRotations topL

testTopR :: Bool
testTopR = testRotations topR

testBottomL :: Bool
testBottomL =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, w), (top, 0, 1, w), (top, 0, 2, w), (top, 1, 0, w), (top, 1, 1, w), (top, 1, 2, w), (top, 2, 0, w), (top, 2, 1, w), (top, 2, 2, w)
                 , (frn, 0, 0, g), (frn, 0, 1, g), (frn, 0, 2, g), (frn, 1, 0, g), (frn, 1, 1, g), (frn, 1, 2, g), (frn, 2, 0, r), (frn, 2, 1, r), (frn, 2, 2, r)
                 , (rgt, 0, 0, r), (rgt, 0, 1, r), (rgt, 0, 2, r), (rgt, 1, 0, r), (rgt, 1, 1, r), (rgt, 1, 2, r), (rgt, 2, 0, b), (rgt, 2, 1, b), (rgt, 2, 2, b)
                 , (bck, 0, 0, b), (bck, 0, 1, b), (bck, 0, 2, b), (bck, 1, 0, b), (bck, 1, 1, b), (bck, 1, 2, b), (bck, 2, 0, o), (bck, 2, 1, o), (bck, 2, 2, o)
                 , (lft, 0, 0, o), (lft, 0, 1, o), (lft, 0, 2, o), (lft, 1, 0, o), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, g), (lft, 2, 1, g), (lft, 2, 2, g)
                 , (bot, 0, 0, y), (bot, 0, 1, y), (bot, 0, 2, y), (bot, 1, 0, y), (bot, 1, 1, y), (bot, 1, 2, y), (bot, 2, 0, y), (bot, 2, 1, y), (bot, 2, 2, y) ]
      rotated = rotate bottomL Cube.solvedCube
  in expected == rotated
  && testRotations bottomL

testBottomR :: Bool
testBottomR = testRotations bottomR

testFrontL :: Bool
testFrontL =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, w), (top, 0, 1, w), (top, 0, 2, w), (top, 1, 0, w), (top, 1, 1, w), (top, 1, 2, w), (top, 2, 0, r), (top, 2, 1, r), (top, 2, 2, r)
                 , (frn, 0, 0, g), (frn, 0, 1, g), (frn, 0, 2, g), (frn, 1, 0, g), (frn, 1, 1, g), (frn, 1, 2, g), (frn, 2, 0, g), (frn, 2, 1, g), (frn, 2, 2, g)
                 , (rgt, 0, 0, y), (rgt, 0, 1, r), (rgt, 0, 2, r), (rgt, 1, 0, y), (rgt, 1, 1, r), (rgt, 1, 2, r), (rgt, 2, 0, y), (rgt, 2, 1, r), (rgt, 2, 2, r)
                 , (bck, 0, 0, b), (bck, 0, 1, b), (bck, 0, 2, b), (bck, 1, 0, b), (bck, 1, 1, b), (bck, 1, 2, b), (bck, 2, 0, b), (bck, 2, 1, b), (bck, 2, 2, b)
                 , (lft, 0, 0, o), (lft, 0, 1, o), (lft, 0, 2, w), (lft, 1, 0, o), (lft, 1, 1, o), (lft, 1, 2, w), (lft, 2, 0, o), (lft, 2, 1, o), (lft, 2, 2, w)
                 , (bot, 0, 0, o), (bot, 0, 1, o), (bot, 0, 2, o), (bot, 1, 0, y), (bot, 1, 1, y), (bot, 1, 2, y), (bot, 2, 0, y), (bot, 2, 1, y), (bot, 2, 2, y) ]
      rotated = rotate frontL Cube.solvedCube
  in expected == rotated
  && testRotations frontL

testFrontR :: Bool
testFrontR = testRotations frontR

testBackL :: Bool
testBackL =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, o), (top, 0, 1, o), (top, 0, 2, o), (top, 1, 0, w), (top, 1, 1, w), (top, 1, 2, w), (top, 2, 0, w), (top, 2, 1, w), (top, 2, 2, w)
                 , (frn, 0, 0, g), (frn, 0, 1, g), (frn, 0, 2, g), (frn, 1, 0, g), (frn, 1, 1, g), (frn, 1, 2, g), (frn, 2, 0, g), (frn, 2, 1, g), (frn, 2, 2, g)
                 , (rgt, 0, 0, r), (rgt, 0, 1, r), (rgt, 0, 2, w), (rgt, 1, 0, r), (rgt, 1, 1, r), (rgt, 1, 2, w), (rgt, 2, 0, r), (rgt, 2, 1, r), (rgt, 2, 2, w)
                 , (bck, 0, 0, b), (bck, 0, 1, b), (bck, 0, 2, b), (bck, 1, 0, b), (bck, 1, 1, b), (bck, 1, 2, b), (bck, 2, 0, b), (bck, 2, 1, b), (bck, 2, 2, b)
                 , (lft, 0, 0, y), (lft, 0, 1, o), (lft, 0, 2, o), (lft, 1, 0, y), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, y), (lft, 2, 1, o), (lft, 2, 2, o)
                 , (bot, 0, 0, y), (bot, 0, 1, y), (bot, 0, 2, y), (bot, 1, 0, y), (bot, 1, 1, y), (bot, 1, 2, y), (bot, 2, 0, r), (bot, 2, 1, r), (bot, 2, 2, r) ]
      rotated = rotate backL Cube.solvedCube
  in expected == rotated
  && testRotations backL

testBackR :: Bool
testBackR = testRotations backR

testCombined :: Bool
testCombined =
  let [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expected = [ (top, 0, 0, o), (top, 0, 1, o), (top, 0, 2, b), (top, 1, 0, o), (top, 1, 1, w), (top, 1, 2, b), (top, 2, 0, o), (top, 2, 1, r), (top, 2, 2, r)
                 , (frn, 0, 0, b), (frn, 0, 1, g), (frn, 0, 2, w), (frn, 1, 0, y), (frn, 1, 1, g), (frn, 1, 2, w), (frn, 2, 0, y), (frn, 2, 1, y), (frn, 2, 2, b)
                 , (rgt, 0, 0, g), (rgt, 0, 1, r), (rgt, 0, 2, w), (rgt, 1, 0, g), (rgt, 1, 1, r), (rgt, 1, 2, w), (rgt, 2, 0, y), (rgt, 2, 1, y), (rgt, 2, 2, o)
                 , (bck, 0, 0, r), (bck, 0, 1, b), (bck, 0, 2, g), (bck, 1, 0, r), (bck, 1, 1, b), (bck, 1, 2, w), (bck, 2, 0, y), (bck, 2, 1, y), (bck, 2, 2, g)
                 , (lft, 0, 0, w), (lft, 0, 1, w), (lft, 0, 2, w), (lft, 1, 0, b), (lft, 1, 1, o), (lft, 1, 2, o), (lft, 2, 0, y), (lft, 2, 1, g), (lft, 2, 2, r)
                 , (bot, 0, 0, g), (bot, 0, 1, g), (bot, 0, 2, r), (bot, 1, 0, o), (bot, 1, 1, y), (bot, 1, 2, r), (bot, 2, 0, o), (bot, 2, 1, b), (bot, 2, 2, b) ]
      rotation = bottomL <> leftL <> backL <> rightL <> frontL <> topL
      rotated = rotate rotation Cube.solvedCube
  in expected == rotated
  && testRotations backL

testNameMatrixRotation :: Bool
testNameMatrixRotation = and $ do
  (r, n) <- rotations
  let n' = nameMatrixRotation r
  return $ n == n'

testInverse :: Bool
testInverse = and $ do
  (r, _) <- rotations
  let (Just r') = inverse r
  let (Just r'') = inverse r'
  return $ r == r''

tests = [ ("Rotation.topToFront", testTopToFront)
        , ("Rotation.topToBack", testTopToBack)
        , ("Rotation.topToRight", testTopToRight)
        , ("Rotation.topToLeft", testTopToLeft)
        , ("Rotation.topTwistRight", testTopTwistRight)
        , ("Rotation.topTwistLeft", testTopTwistLeft)
        , ("Rotation.rightL", testRightL)
        , ("Rotation.rightR", testRightR)
        , ("Rotation.leftL", testLeftL)
        , ("Rotation.leftR", testLeftR)
        , ("Rotation.topL", testTopL)
        , ("Rotation.topR", testTopR)
        , ("Rotation.bottomL", testBottomL)
        , ("Rotation.bottomR", testBottomR)
        , ("Rotation.frontL", testFrontL)
        , ("Rotation.frontR", testFrontR)
        , ("Rotation.backL", testBackL)
        , ("Rotation.backR", testBackR)
        , ("combined", testCombined)
        , ("Rotation.nameMatrixRotation", testNameMatrixRotation)
        , ("Rotation.inverse", testInverse)
        ]