{-# LANGUAGE TupleSections #-}

module Board where

import           Control.Arrow
import           Control.Monad.Reader
import           Control.Monad.State
import           Data.Array                     ( Array )
import qualified Data.Array                    as A
import           Data.IntSet                    ( IntSet )
import qualified Data.IntSet                   as IntSet
import           Data.List
import           Data.Maybe
import qualified Data.Vector                   as V
import           Math.Geometry.Grid
import           Math.Geometry.Grid.Octagonal
import           Numeric.LinearAlgebra
import           System.Random

import           Types

-- | Generate a unique 'Int' in a range.
-- The state variable is a random number seed and an 'IntSet' of used
-- 'Int's.
randRUniq :: RandomGen g => (Int, Int) -> State (g, IntSet) Int
randRUniq (l, h) = state $ \(g, ns) ->
        let
                vacant = IntSet.toList $ IntSet.fromList [l .. h] IntSet.\\ ns
                (idx, g') = randomR (0, length vacant - 1) g
                n = vacant !! idx
        in
                (n, (g', IntSet.insert n ns))

-- | Generate 'n' random unique numbers in an inclusive range.
randRUniq1 :: RandomGen g => Int -> (Int, Int) -> State g [Int]
randRUniq1 n bounds =
        state
                $ second fst
                . runState (replicateM n $ randRUniq bounds)
                . (, IntSet.empty)

-- | Generate a subset of a list.
randSubList :: RandomGen g => Int -> [a] -> State g [a]
randSubList n xs = map (xs !!) <$> randRUniq1 n (0, length xs - 1)

-- | Generate 'n' mines on a board that is 'w' columns by 'h' rows
-- such that none are adjacent to the specified location ('x', 'y').
genMinesSafe
        :: RandomGen g
        => Int
        -> (Int, Int)
        -> (Int, Int)
        -> State g [(Int, Int)]
genMinesSafe n (h, w) (y, x) =
        randSubList n
                $  ((,) <$> [0 .. h - 1] <*> [0 .. w - 1])
                \\ ((,) <$> [y - 1 .. y + 1] <*> [x - 1 .. x + 1])

-- | Generate a logical matrix of dimensions ('w', 'h') representing
-- the locations of the mines in 'xs'.
minesMask :: (Int, Int) -> [(Int, Int)] -> Matrix Double
minesMask (h, w) xs =
        let
                blank     = V.replicate (w * h) 0
                withMines = blank V.// zip
                        (map (uncurry $ (+) . (* w)) xs)
                        (repeat 1)
        in
                subMatrix (1, 1) (h, w)
                . conv2 (konst 1 (3, 3))
                . (h >< w)
                $ V.toList withMines

-- | Generate the minesweeper answer key of dimensions 'd' and mines at 'xs'.
minesKey :: (Int, Int) -> [(Int, Int)] -> Array (Int, Int) Tile
minesKey d@(h, w) xs =
        A.listArray ((0, 0), (h - 1, w - 1))
                    (pure . round <$> toList (flatten $ minesMask d xs))
                A.// zip xs (repeat Nothing)

-- | Generate a random minesweeper key with 'n' mines and dimensions 'd'.
minesBoard
        :: RandomGen g
        => Int -- ^ The amount of mines to generate.
        -> (Int, Int) -- ^ The dimensions of the board.
        -> (Int, Int) -- ^ The initial safe location. No mines will be
    -- adjacent to this cell.
        -> State g (Array (Int, Int) Tile)
minesBoard n d = fmap (minesKey d) . genMinesSafe n d

gameBoard :: RandomGen g => Int -> (Int, Int) -> (Int, Int) -> State g Board
gameBoard = ((fmap ((, Covered) <$>) .) .) . minesBoard

uncoverMines :: Board -> Board
uncoverMines = fmap (\b@(t, _) -> if isJust t then b else (Nothing, Uncovered))

addMark :: Mark -> (Bool, Board) -> (Int, Int) -> (Bool, Board)
addMark m s@(b, board) x
        | b         = (True, board A.// [(x, second (const m) $ board A.! x)])
        | otherwise = s

mark, flag :: (Bool, Board) -> (Int, Int) -> (Bool, Board)
mark = addMark Marked
flag = addMark Flagged

surroundingTiles :: Board -> (Int, Int) -> [(Int, Int)]
surroundingTiles board idx =
        let ((y0, x0), (y1, x1)) = A.bounds board
            boardW               = x1 - x0 + 1
            boardH               = y1 - y0 + 1
        in  neighbours (rectOctGrid boardW boardH) idx

uncover :: (Bool, Board) -> (Int, Int) -> (Bool, Board)
uncover s@(inPlay, board) idx
        |
  -- Click on a mine
          inPlay && isNothing tile
        = (False, uncoverMines board)
        |

  -- Click on a covered or marked positive square
          inPlay && (tileMark == Covered || tileMark == Marked) && not isZero
        = (inPlay, board A.// [(idx, (tile, Uncovered))])
        |

  -- Click on a square that is 0.
          inPlay && (tileMark == Marked || tileMark == Covered)
        = foldr (flip uncover)
                (True, board A.// [(idx, (tile, Uncovered))])
                surrounding
        |

  -- Click on a square that is completely flagged
          inPlay && not isZero && tileMark == Uncovered && tile == Just
                (length flaggedSurround)
        = foldr (flip uncover)
                (True, board A.// [(idx, (tile, Uncovered))])
                unFlaggedSurround
        | otherwise
        = s
    where
        (tile, tileMark)  = board A.! idx
        isZero            = tile == Just 0
        surrounding       = surroundingTiles board idx
        flaggedSurround = filter ((== Flagged) . snd . (board A.!)) surrounding
        unFlaggedSurround = surrounding \\ flaggedSurround


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses   #-}
{-# LANGUAGE ImpredicativeTypes, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Count-Min Sketch
-- | Creator: Xiao Ling
-- | Created: 12/17/2015
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module CountMinSketch where

import Control.Monad.Random.Class
import Control.Monad.Random
import Control.Monad.State

import Data.Matrix
import Data.Conduit
import qualified Data.Conduit.List as Cl
import qualified Data.Matrix as M

import Core
import Utils
import Statistics

{-----------------------------------------------------------------------------
  Types
------------------------------------------------------------------------------}

type D        = Int              -- * num rows
type W        = Int              -- * num cols
type P        = Int              -- * some prime that is 2^n - 1
type Event    = Int              -- * some input event from universe:
                                 -- * {1, .., i, .., n}

type Sketch   = Matrix Int       -- * D x W matrix  of counters
type ABs      = Matrix Int       -- * 2 x D matrix  of a_j and b_j for j = 1..d
type ColIdx   = [Int]            -- * indices of counters to increment
type Hash     = Event -> ColIdx  -- * Maps event to d x 1 list of column indices

-- * Algorithm State
type AState    = (Sketch,Hash,(D,W,P))

-- * the algorithm uses `some m` equipped with state and randomness
type Some m    = (MonadRandom m, MonadState AState m)

-- * Algorithm outputes an `Estimator` that
-- * estimates the fequence of `event` i from  universe {1 ... i... n},
type Estimator = Event -> Int

{-----------------------------------------------------------------------------
    Count Min Sketch Conduit
------------------------------------------------------------------------------}

-- * Batch mode
sketch' :: EpsDelta -> Event -> Batch Event IO Estimator
sketch' t n = sketch t n `using` eval

-- * Streaming mode
sketch :: Some m => EpsDelta -> Event -> Streaming Event m ()
sketch t n = inits False t n $= awaitWith update'

-- * initalize algorithm given confidence/error `t` and maximum event `n`
inits :: Some m => Bool -> EpsDelta -> Event -> Conduit Event m Event
inits False t n = lift (inits' t n) >> inits True t n
inits _     _ _ = pass

-- * a concrete `eval`uation of Some `m`
-- * run with vacuous inital condition to be thrown away
eval :: StateT AState (Rand StdGen) () -> IO Estimator
eval m = fmap (\(s,h,_) -> \i -> freq i s h)
        . evalRandIO $ execStateT m $ (ones 1 1, const [], (0,0,0))

{-------------------------------------------------------------------------------
   CountMinSketch Primitives
--------------------------------------------------------------------------------}

-- * initalize algorithm with (w,d,p), zero counters, and hash function
inits' :: Some m => EpsDelta -> Event -> m ()
inits' ed n = do
    let t = newDwp ed n
    setParam t
    hash    <- newHash t
    let sketch = newSketch t
    setHash   hash
    setSketch sketch

-- * on event `i`, update the `sketch`
update' :: Some m => Event -> m ()
update' i = do
    pm      <- askParam
    sketch  <- getSketch
    hash    <- askHash
    let incr = toMask (hash i) pm
    setSketch $ incr .+. sketch

{-------------------------------------------------------------------------------
   Effectful Subroutines
--------------------------------------------------------------------------------}

-- * h_j (i) = (a_j x i + b_j mod p) mod w
-- * sumRow-wise  [ a_1 ... a_j ... a_d ]  .*.  [ i_1 ...  i_d ]     mod p  mod w
-- *              [ b_1 ... b_j ... b_d ]       [ 1   ...  1   ]
-- * result:      [ idx_1, ..., idx_d   ]
-- *               where the indices start at 1
newHash :: MonadRandom m => (D,W,P) -> m Hash
newHash t = do
    ab  <- newAbs t
    return $ toHash t ab

-- * intialize random cofficients a_j, b_j, for d x 2 matrix:
-- * [ a_1 ... a_j ... a_d ]
-- * [ b_1 ... b_j ... b_d ]
newAbs :: MonadRandom m => (D,W,P) -> m ABs
newAbs (d,_,p) = (M.fromList 2 d . take (2*d)) <$> getRandomRs (1,p)

{-----------------------------------------------------------------------------
    Pure Subroutines
------------------------------------------------------------------------------}

-- * given some event `i` from  universe {1 ... i... n}, a table of
-- * counters Sketch `s` and hash function `h`, compute the
-- * frequence of event `i`
freq :: Event -> Sketch -> Hash -> Int
freq i s h = minimum $ go s (h i) 1
    where go s []     _ = []
          go s (c:cs) d = (getElem d c s) : (go s cs $ succ d)

-- * compute d,w, and p given accuracy/confidence level and
-- * max value of event `n`
newDwp :: EpsDelta -> Event -> (D,W,P)
newDwp (ED e d) n = (round $ log2(1/d), round $ 2/e, round $ 2^n -1)

-- * initalize the sketch
newSketch :: (D,W,P) -> Sketch
newSketch (d,w,_) = zeros d w

-- * Note `is` maps some `i` from event universe onto d x 2 matrix:
-- * [ i_1 ... i_2 ]
-- * [ 1   ...  1  ]
toHash :: (D, W, P) -> ABs -> Hash
toHash (d, w, p) ab i = go $ sumR (ab .*. is) ... (\v -> v `mod` p `mod` w)
    where go = fmap (+1) . toList
          is = (i .+ (zeros 1 d)) <-> ones 1 d

-- * given hash function, construct a binary mask of size `w` x `d`
-- * to increment sketch
toMask :: ColIdx -> (D,W,P) -> Sketch
toMask ks (d,w,p) = go w ks
    where
        go w ks = let vec c w = zeros 1 (c-1) <|> single 1 <|> zeros 1 (w - c) in
            case ks of
                c:[] -> vec c w
                c:cs -> vec c w <-> go w cs

{-----------------------------------------------------------------------------
    Managing State
    TODO : redesign so this is not necessary
------------------------------------------------------------------------------}

askParam :: Some m => m (D,W,P)
askParam    = get >>= \(_,_,p) -> return p

setParam :: Some m => (D,W,P) -> m ()
setParam t  = get >>= \(s,h,_) -> put (s,h,t)

getSketch :: Some m => m Sketch
getSketch   = get >>= \(s,_,_) -> return s

setSketch :: Some m => Sketch -> m ()
setSketch s = get >>= \(_,h,t) -> put (s,h,t)

askHash :: Some m => m Hash
askHash     = get >>= \(_,h,_) -> return h

setHash :: Some m => Hash -> m ()
setHash h   = get >>= \(s,_,t) -> put (s,h,t)

{-----------------------------------------------------------------------------
    Find conduit lib function for these:
------------------------------------------------------------------------------}

-- * identity conduit, why can't I find this in the library?
pass :: Monad m => Conduit i m i
pass  = do
    mx <- await
    case mx of
        Nothing -> return ()
        Just x  -> yield x >> pass

awaitWith :: Monad m => (i -> m a) -> ConduitM i o m ()
awaitWith g = do
    mi <- await
    case mi of
        Nothing -> return ()
        Just i  -> lift (g i) >> awaitWith g


{-# LANGUAGE ScopedTypeVariables, TypeFamilies, TupleSections #-}

module Fuml.Core where

import qualified Data.Vector.Storable as VS
import Numeric.LinearAlgebra
import Data.List (nub)
import Lens.Micro
import Control.Monad.Identity

data Weighted o = Weighted Double o

-- |The result of running a predictive model
data Predict p a = Predict
  { model :: p -- ^ the internal state of the model
  , predict :: Vector Double -> a -- ^ the prediction function
  }

instance Functor (Predict p) where
  fmap f (Predict m predf) = Predict m (f . predf)

newtype Supervisor m o p a
  = Supervisor { runSupervisor :: Maybe p
                               -> [(Vector Double, o)]
                               -> m (Predict p a) }

instance Functor m => Functor (Supervisor m o p) where
  fmap f (Supervisor sf) = Supervisor $ \mp d -> fmap (fmap f) $ sf mp d

-- | Helper function for running simple Supervisors in the Identity monad
runSupervisor' :: Supervisor Identity o p a -- ^ the 'Supervisor' value
               -> [(Vector Double, o)] -- ^ the dataset
               -> Predict p a -- ^ the 'Predict' value
runSupervisor' (Supervisor sf) = runIdentity . sf Nothing

oneVsRest :: (Monad m, Eq a, Functor m) => Supervisor m Bool p b -> Supervisor m a [(a,p)] [(a,b)]
oneVsRest subsuper  = Supervisor $ \_ theData -> do
  let classes = nub $ map snd theData
      boolize c (v, c1) = (v, c == c1)
      train c = fmap (c,) $ runSupervisor subsuper Nothing $ map (boolize c) theData
  models <-  mapM train classes
  return $ Predict (map (over _2 model) models) $ \v -> map (\(c,pr) -> (c, predict pr v)) models

(~~) :: (a -> o) -> [a -> Double] -> [a] -> [(Vector Double, o)]
fy ~~ fxs = let f w = (VS.fromList (map ($w) fxs) , fy w)
            in map f

--prepare :: (a -> o) -> [a -> Double] -> [a] -> [(Vector Double, o)]
--prepare fo fxs d = fo ~~ fxs $ d


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DefaultSignatures #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Operations on free vector spaces.
-----------------------------------------------------------------------------
module Linear.Vector
  ( Additive(..)
  , E(..)
  , negated
  , (^*)
  , (*^)
  , (^/)
  , sumV
  , basis
  , basisFor
  , scaled
  , outer
  , unit
  ) where

import Control.Applicative
import Control.Lens
import Data.Complex
import Data.Foldable as Foldable (forM_, foldl')
import Data.Functor.Compose
import Data.Functor.Product
import Data.HashMap.Lazy as HashMap
import Data.Hashable
import Data.IntMap as IntMap
import Data.Map as Map
import qualified Data.Vector as Vector
import Data.Vector (Vector)
import qualified Data.Vector.Mutable as Mutable
import GHC.Generics
import Linear.Instances ()

-- $setup
-- >>> import Linear.V2

-- | Basis element
newtype E t = E { el :: forall x. Lens' (t x) x }

infixl 6 ^+^, ^-^
infixl 7 ^*, *^, ^/

class GAdditive f where
  gzero :: Num a => f a
  gliftU2 :: (a -> a -> a) -> f a -> f a -> f a
  gliftI2 :: (a -> b -> c) -> f a -> f b -> f c

instance GAdditive U1 where
  gzero = U1
  {-# INLINE gzero #-}
  gliftU2 _ U1 U1 = U1
  {-# INLINE gliftU2 #-}
  gliftI2 _ U1 U1 = U1
  {-# INLINE gliftI2 #-}

instance (GAdditive f, GAdditive g) => GAdditive (f :*: g) where
  gzero = gzero :*: gzero
  {-# INLINE gzero #-}
  gliftU2 f (a :*: b) (c :*: d) = gliftU2 f a c :*: gliftU2 f b d
  {-# INLINE gliftU2 #-}
  gliftI2 f (a :*: b) (c :*: d) = gliftI2 f a c :*: gliftI2 f b d
  {-# INLINE gliftI2 #-}

instance (Additive f, GAdditive g) => GAdditive (f :.: g) where
  gzero = Comp1 $ gzero <$ (zero :: f Int)
  {-# INLINE gzero #-}
  gliftU2 f (Comp1 a) (Comp1 b) = Comp1 $ liftU2 (gliftU2 f) a b
  {-# INLINE gliftU2 #-}
  gliftI2 f (Comp1 a) (Comp1 b) = Comp1 $ liftI2 (gliftI2 f) a b
  {-# INLINE gliftI2 #-}

instance Additive f => GAdditive (Rec1 f) where
  gzero = Rec1 zero
  {-# INLINE gzero #-}
  gliftU2 f (Rec1 g) (Rec1 h) = Rec1 (liftU2 f g h)
  {-# INLINE gliftU2 #-}
  gliftI2 f (Rec1 g) (Rec1 h) = Rec1 (liftI2 f g h)
  {-# INLINE gliftI2 #-}

instance GAdditive f => GAdditive (M1 i c f) where
  gzero = M1 gzero
  {-# INLINE gzero #-}
  gliftU2 f (M1 g) (M1 h) = M1 (gliftU2 f g h)
  {-# INLINE gliftU2 #-}
  gliftI2 f (M1 g) (M1 h) = M1 (gliftI2 f g h)
  {-# INLINE gliftI2 #-}

instance GAdditive Par1 where
  gzero = Par1 0
  gliftU2 f (Par1 a) (Par1 b) = Par1 (f a b)
  {-# INLINE gliftU2 #-}
  gliftI2 f (Par1 a) (Par1 b) = Par1 (f a b)
  {-# INLINE gliftI2 #-}

-- | A vector is an additive group with additional structure.
class Functor f => Additive f where
  -- | The zero vector
  zero :: Num a => f a
#ifndef HLINT
  default zero :: (GAdditive (Rep1 f), Generic1 f, Num a) => f a
  zero = to1 gzero
#endif

  -- | Compute the sum of two vectors
  --
  -- >>> V2 1 2 ^+^ V2 3 4
  -- V2 4 6
  (^+^) :: Num a => f a -> f a -> f a
  (^+^) = liftU2 (+)
  {-# INLINE (^+^) #-}

  -- | Compute the difference between two vectors
  --
  -- >>> V2 4 5 ^-^ V2 3 1
  -- V2 1 4
  (^-^) :: Num a => f a -> f a -> f a
  x ^-^ y = x ^+^ negated y

  -- | Linearly interpolate between two vectors.
  lerp :: Num a => a -> f a -> f a -> f a
  lerp alpha u v = alpha *^ u ^+^ (1 - alpha) *^ v
  {-# INLINE lerp #-}

  -- | Apply a function to merge the 'non-zero' components of two vectors, unioning the rest of the values.
  --
  -- * For a dense vector this is equivalent to 'liftA2'.
  --
  -- * For a sparse vector this is equivalent to 'unionWith'.
  liftU2 :: (a -> a -> a) -> f a -> f a -> f a
#ifndef HLINT
  default liftU2 :: Applicative f => (a -> a -> a) -> f a -> f a -> f a
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
#endif

  -- | Apply a function to the components of two vectors.
  --
  -- * For a dense vector this is equivalent to 'liftA2'.
  --
  -- * For a sparse vector this is equivalent to 'intersectionWith'.
  liftI2 :: (a -> b -> c) -> f a -> f b -> f c
#ifndef HLINT
  default liftI2 :: Applicative f => (a -> b -> c) -> f a -> f b -> f c
  liftI2 = liftA2
  {-# INLINE liftI2 #-}
#endif

instance (Additive f, Additive g) => Additive (Product f g) where
  zero = Pair zero zero
  liftU2 f (Pair a b) (Pair c d) = Pair (liftU2 f a c) (liftU2 f b d)
  liftI2 f (Pair a b) (Pair c d) = Pair (liftI2 f a c) (liftI2 f b d)
  Pair a b ^+^ Pair c d = Pair (a ^+^ c) (b ^+^ d)
  Pair a b ^-^ Pair c d = Pair (a ^-^ c) (b ^-^ d)
  lerp alpha (Pair a b) (Pair c d) = Pair (lerp alpha a c) (lerp alpha b d)

instance (Additive f, Additive g) => Additive (Compose f g) where
  zero = Compose $ zero <$ (zero :: f Int)
  {-# INLINE zero #-}
  Compose a ^+^ Compose b = Compose $ liftU2 (^+^) a b
  {-# INLINE (^+^) #-}
  Compose a ^-^ Compose b = Compose $ liftU2 (^-^) a b
  {-# INLINE (^-^) #-}
  liftU2 f (Compose a) (Compose b) = Compose $ liftU2 (liftU2 f) a b
  {-# INLINE liftU2 #-}
  liftI2 f (Compose a) (Compose b) = Compose $ liftI2 (liftI2 f) a b
  {-# INLINE liftI2 #-}

instance Additive ZipList where
  zero = ZipList []
  {-# INLINE zero #-}
  liftU2 f (ZipList xs) (ZipList ys) = ZipList (liftU2 f xs ys)
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Additive Vector where
  zero = mempty
  {-# INLINE zero #-}
  liftU2 f u v = case compare lu lv of
    LT | lu == 0   -> v
       | otherwise -> Vector.modify (\ w -> Foldable.forM_ [0..lu-1] $ \i -> Mutable.unsafeWrite w i $ f (Vector.unsafeIndex u i) (Vector.unsafeIndex v i)) v
    EQ -> Vector.zipWith f u v
    GT | lv == 0   -> u
       | otherwise -> Vector.modify (\ w -> Foldable.forM_ [0..lv-1] $ \i -> Mutable.unsafeWrite w i $ f (Vector.unsafeIndex u i) (Vector.unsafeIndex v i)) u
    where
      lu = Vector.length u
      lv = Vector.length v
  {-# INLINE liftU2 #-}
  liftI2 = Vector.zipWith
  {-# INLINE liftI2 #-}

instance Additive Maybe where
  zero = Nothing
  {-# INLINE zero #-}
  liftU2 f (Just a) (Just b) = Just (f a b)
  liftU2 _ Nothing ys = ys
  liftU2 _ xs Nothing = xs
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Additive [] where
  zero = []
  {-# INLINE zero #-}
  liftU2 f = go where
    go (x:xs) (y:ys) = f x y : go xs ys
    go [] ys = ys
    go xs [] = xs
  {-# INLINE liftU2 #-}
  liftI2 = Prelude.zipWith
  {-# INLINE liftI2 #-}

instance Additive IntMap where
  zero = IntMap.empty
  {-# INLINE zero #-}
  liftU2 = IntMap.unionWith
  {-# INLINE liftU2 #-}
  liftI2 = IntMap.intersectionWith
  {-# INLINE liftI2 #-}

instance Ord k => Additive (Map k) where
  zero = Map.empty
  {-# INLINE zero #-}
  liftU2 = Map.unionWith
  {-# INLINE liftU2 #-}
  liftI2 = Map.intersectionWith
  {-# INLINE liftI2 #-}

instance (Eq k, Hashable k) => Additive (HashMap k) where
  zero = HashMap.empty
  {-# INLINE zero #-}
  liftU2 = HashMap.unionWith
  {-# INLINE liftU2 #-}
  liftI2 = HashMap.intersectionWith
  {-# INLINE liftI2 #-}

instance Additive ((->) b) where
  zero   = const 0
  {-# INLINE zero #-}
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

instance Additive Complex where
  zero = 0 :+ 0
  {-# INLINE zero #-}
  liftU2 f (a :+ b) (c :+ d) = f a c :+ f b d
  {-# INLINE liftU2 #-}
  liftI2 f (a :+ b) (c :+ d) = f a c :+ f b d
  {-# INLINE liftI2 #-}

instance Additive Identity where
  zero = Identity 0
  {-# INLINE zero #-}
  liftU2 = liftA2
  {-# INLINE liftU2 #-}
  liftI2 = liftA2
  {-# INLINE liftI2 #-}

-- | Compute the negation of a vector
--
-- >>> negated (V2 2 4)
-- V2 (-2) (-4)
negated :: (Functor f, Num a) => f a -> f a
negated = fmap negate
{-# INLINE negated #-}

-- | Sum over multiple vectors
--
-- >>> sumV [V2 1 1, V2 3 4]
-- V2 4 5
sumV :: (Foldable f, Additive v, Num a) => f (v a) -> v a
sumV = Foldable.foldl' (^+^) zero
{-# INLINE sumV #-}

-- | Compute the left scalar product
--
-- >>> 2 *^ V2 3 4
-- V2 6 8
(*^) :: (Functor f, Num a) => a -> f a -> f a
(*^) a = fmap (a*)
{-# INLINE (*^) #-}

-- | Compute the right scalar product
--
-- >>> V2 3 4 ^* 2
-- V2 6 8
(^*) :: (Functor f, Num a) => f a -> a -> f a
f ^* a = fmap (*a) f
{-# INLINE (^*) #-}

-- | Compute division by a scalar on the right.
(^/) :: (Functor f, Fractional a) => f a -> a -> f a
f ^/ a = fmap (/a) f
{-# INLINE (^/) #-}

-- | Produce a default basis for a vector space. If the dimensionality
-- of the vector space is not statically known, see 'basisFor'.
basis :: (Additive t, Traversable t, Num a) => [t a]
basis = basisFor (zero :: Additive v => v Int)

-- | Produce a default basis for a vector space from which the
-- argument is drawn.
basisFor :: (Traversable t, Num a) => t b -> [t a]
basisFor = \t ->
   ifoldMapOf traversed ?? t $ \i _ ->
     return                  $
       iover  traversed ?? t $ \j _ ->
         if i == j then 1 else 0
{-# INLINABLE basisFor #-}

-- | Produce a diagonal (scale) matrix from a vector.
--
-- >>> scaled (V2 2 3)
-- V2 (V2 2 0) (V2 0 3)
scaled :: (Traversable t, Num a) => t a -> t (t a)
scaled = \t -> iter t (\i x -> iter t (\j _ -> if i == j then x else 0))
  where
  iter :: Traversable t => t a -> (Int -> a -> b) -> t b
  iter x f = iover traversed f x
{-# INLINE scaled #-}

-- | Create a unit vector.
--
-- >>> unit _x :: V2 Int
-- V2 1 0
unit :: (Additive t, Num a) => ASetter' (t a) a -> t a
unit l = set' l 1 zero

-- | Outer (tensor) product of two vectors
outer :: (Functor f, Functor g, Num a) => f a -> g a -> f (g a)
outer a b = fmap (\x->fmap (*x) b) a


--Utility functions
--Copyright Laurence Emms 2018

module Util (loadIntToRegisters,
             addRegisters,
             ifC,
             whileC,
             doWhileC) where

import Data.Complex
import Data.Bits
import Data.List
import Register
import Instruction
import ClassicalCircuit
import QuantumCircuit

loadIntToRegisters :: (Floating a, Show a, Ord a) => Int -> Int -> [Instruction a]
loadIntToRegisters x r = loadIntBit x r 0 []

loadIntBit :: (Floating a, Show a, Ord a) => Int -> Int -> Int -> [Instruction a] -> [Instruction a]
loadIntBit x r b commands
    | b < 32 = if (testBit x b)
               then (loadIntBit x r (b + 1) (commands ++ [(ITrue (Register (r + b)))]))
               else (loadIntBit x r (b + 1) (commands ++ [(IFalse (Register (r + b)))]))
    | otherwise = commands

addRegisters :: (Floating a, Show a, Ord a) => Int -> Int -> Int -> Int -> Int -> [Instruction a]
addRegisters carry temp a b r
    = (IFalse (Register carry)) :
      (IFalse (Register temp)) :
      (fmap (\(x, y, z) -> (CallCircuit adder
                                       [Right (Register carry),
                                        Right (Register temp),
                                        Right (Register x),
                                        Right (Register y),
                                        Right (Register z)]))
                                       (zip3 [a..(a + 31)] [b..(b + 31)] [r..(r + 31)]))

ifC :: (Floating a, Show a, Ord a) => String -> String -> Classical a -> [Instruction a] -> [Instruction a] -> [Instruction a]
ifC thenLabel endLabel c x y = (JumpWhen thenLabel c) :
                               y ++
                               [(Jump endLabel),
                                (Label thenLabel)] ++
                               x ++
                               [(Label endLabel)]

whileC :: (Floating a, Show a, Ord a) => String -> String -> Classical a -> [Instruction a] -> [Instruction a]
whileC conditionLabel loopLabel c x = [(Jump conditionLabel),
                                       (Label loopLabel)] ++
                                      x ++
                                      [(Label conditionLabel),
                                       (JumpWhen loopLabel c)]

doWhileC :: (Floating a, Show a, Ord a) => String -> Classical a -> [Instruction a] -> [Instruction a]
doWhileC whileLabel c x = (Label whileLabel) :
                          x ++
                          [(JumpWhen whileLabel c)]


{-# LANGUAGE CPP                 #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Layers.FullyConnected where

import           System.Random.MWC             (create)

import           Data.Constraint               (Dict (..))
import           Data.Proxy
import           Data.Singletons               ()
import           Data.Kind                     (Type)
import           Data.Serialize
import           Data.Either

import           GHC.TypeLits
import           Unsafe.Coerce                 (unsafeCoerce)

import           Hedgehog
import qualified Hedgehog.Range              as Range

import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix
import           Test.Utils.Rnf
import           Test.Grenade.Layers.Internal.Reference

import qualified Numeric.LinearAlgebra.Data   as D
import           Numeric.LinearAlgebra.Static (R)
import qualified Numeric.LinearAlgebra.Static as H

import           Grenade.Core
import           Grenade.Layers.FullyConnected
import           Grenade.Utils.ListStore

data OpaqueFullyConnected :: Type where
     OpaqueFullyConnected :: (KnownNat i, KnownNat o, KnownNat (i * o)) => FullyConnected i o -> OpaqueFullyConnected

instance Show OpaqueFullyConnected where
    show (OpaqueFullyConnected n) = show n

genOpaqueFullyConnected :: Gen OpaqueFullyConnected
genOpaqueFullyConnected = do
  input   :: Integer  <- choose 2 100
  output  :: Integer  <- choose 1 100
  let Just input'      = someNatVal input
  let Just output'     = someNatVal output
  case (input', output') of
    (SomeNat (Proxy :: Proxy i'), SomeNat (Proxy :: Proxy o')) ->
      case (unsafeCoerce (Dict :: Dict ()) :: Dict (KnownNat (i' * o'))) of
        Dict -> do
          wB    <- randomVector
          bM    <- randomVector
          wN    <- uniformSample
          kM    <- uniformSample
          return . OpaqueFullyConnected $ (FullyConnected (FullyConnected' wB wN) (ListStore 0 [Just $ FullyConnected' bM kM]) :: FullyConnected i' o')

prop_fully_connected_run_forwards_backwards :: Property
prop_fully_connected_run_forwards_backwards = property $ do
  OpaqueFullyConnected (fc :: FullyConnected i o) <- blindForAll genOpaqueFullyConnected
  let (FullyConnected (FullyConnected' b w) _) = fc
  input <- forAll randomVector
  let (ntape, nout)  = naiveFullyConnectedRunForwards w b input
      (nnw, nnb, nd) = naiveFullyConnectedBackprop w (ntape, nout)
      (tape, S1D out :: S ('D1 o)) = runForwards fc (S1D input)
      (grad, S1D d   :: S ('D1 i)) = runBackwards fc tape (S1D out)
      FullyConnected' nb nw = grad

  H.extract out  `isSimilarVectorTo` H.extract nout
  H.extract d    `isSimilarVectorTo` H.extract nd
  H.extract nw   `isSimilarMatrixTo` H.extract nnw
  H.extract nb   `isSimilarVectorTo` H.extract nnb
  H.extract tape `isSimilarVectorTo` H.extract ntape

prop_fully_connected_rnf :: Property
prop_fully_connected_rnf = property $ do
  gen <- evalIO create
  fc :: FullyConnected 3 5 <- evalIO $ createRandomWith UniformInit gen
  let (FullyConnected fc'@(FullyConnected' b w) _) = fc

  -- Evaluates the weights and biases
  err0 <- evalIO $ tryEvalRnf (FullyConnected (error expectedErrStr) mkListStore)
  assert $ rnfRaisedErr err0

  -- Evaluates the list store
  err1 <- evalIO $ tryEvalRnf (FullyConnected fc' (error expectedErrStr))
  assert $ rnfRaisedErr err1

  -- Evaluates the biases
  err2 <- evalIO $ tryEvalRnf (FullyConnected (FullyConnected' (error expectedErrStr) w) mkListStore)
  assert $ rnfRaisedErr err2

  -- Evaluates the weights
  err3 <- evalIO $ tryEvalRnf (FullyConnected (FullyConnected' b (error expectedErrStr)) mkListStore)
  assert $ rnfRaisedErr err3

  OpaqueFullyConnected (fcValid :: FullyConnected i o) <- blindForAll genOpaqueFullyConnected
  noErrStr <- evalIO $ tryEvalRnf fcValid
  assert $ rnfNoError noErrStr

prop_fully_connected_is_randomizable :: Property
prop_fully_connected_is_randomizable = withTests 1 $ property $ do
  gen <- evalIO create
  _ :: FullyConnected 3 5 <- evalIO $ createRandomWith UniformInit gen
  success

prop_fully_connected_serialized :: Property
prop_fully_connected_serialized = property $ do
  OpaqueFullyConnected (fc :: FullyConnected i o) <- blindForAll genOpaqueFullyConnected
  let bs = encode fc
  let dec = decode bs :: Either String (FullyConnected i o)
  assert $ isRight dec
  let Right fc' = dec
      (FullyConnected (FullyConnected' b w)   _) = fc
      (FullyConnected (FullyConnected' b' w') _) = fc'
  H.extract b === H.extract b'
  H.extract w === H.extract w'

prop_fully_connected_show :: Property
prop_fully_connected_show = withTests 1 $ property $ do
  OpaqueFullyConnected fc <- blindForAll genOpaqueFullyConnected
  (show fc) `seq` success

prop_fully_connected_update :: Property
prop_fully_connected_update = withTests 1 $ property $ do
  let fc :: FullyConnected 3 5 = createFixed
      input = createFixedInput :: R 3
      (tape, out :: S ('D1 5)) = runForwards fc (S1D input)
      (grad, _   :: S ('D1 3)) = runBackwards fc tape out
      FullyConnected (FullyConnected' b w)   _ = runUpdate defSGD  fc grad
      FullyConnected (FullyConnected' b' w') _ = runUpdate defAdam fc grad
      expW :: H.L 5 3  = H.fromList [ 0.1214, 0.2429, 0.3643
                                    , 0.5229, 0.6171, 0.7114
                                    , 0.9243, 0.9914, 1.0586
                                    , 1.3257, 1.3657, 1.4057
                                    , 1.7271, 1.7400, 1.7529 ]
      expW' :: H.L 5 3 = H.fromList [ 0.1419, 0.2847, 0.4276
                                    , 0.5704, 0.7133, 0.8561
                                    , 0.9990, 1.1419, 1.2847
                                    , 1.4276, 1.5704, 1.7133
                                    , 1.8561, 1.9990, 2.1419 ]
  assert $ allCloseV b  (H.fromList [0.1214,0.2371,0.3529,0.4686,0.5843])
  assert $ allCloseV b' (H.fromList [0.1419,0.2847,0.4276,0.5704,0.7133] )
  assert $ allCloseP (S2D w)  (S2D expW) 0.001
  assert $ allCloseP (S2D w') (S2D expW') 0.001
  where
    createFixed :: forall i o. (KnownNat i, KnownNat o) => FullyConnected i o
    createFixed = fc
      where
        i' = fromIntegral $ natVal (Proxy :: Proxy i)
        o' = fromIntegral $ natVal (Proxy :: Proxy o)
        bias :: H.R o   = H.fromList (map (/7) [1..o'])
        acts :: H.L o i = H.fromList (map (/7) [1..(o' * i')])
        fc :: FullyConnected i o = FullyConnected (FullyConnected' bias acts) mkListStore

    createFixedInput :: forall i. (KnownNat i) => R i
    createFixedInput = H.fromList [1..i']
      where
          i' = fromIntegral $ natVal (Proxy :: Proxy i)

prop_fully_connected_foldable_gradient :: Property
prop_fully_connected_foldable_gradient = property $ do
  OpaqueFullyConnected (fc :: FullyConnected i o) <- blindForAll genOpaqueFullyConnected
  m <- forAll $ genRealNum (Range.constant 0 0.9)
  let FullyConnected fc'@(FullyConnected' b w) _ = fc
      FullyConnected' b' w' = mapGradient (*m) fc'
      bl  = D.toList $ H.extract b
      wl  = concat $ D.toLists $ H.extract w
      bl' = D.toList $ H.extract b'
      wl' = concat $ D.toLists $ H.extract w'
      [sb, sw] = squaredSums fc'
      ssum =  sum . map (** 2)

  (map (*m) bl) `isSimilarListTo` bl'
  (map (*m) wl) `isSimilarListTo` wl'

  ssum bl `isWithinPrecisionOf` sb
  ssum wl `isWithinPrecisionOf` sw
  where
    isWithinPrecisionOf = isWithinOf 0.001

tests :: IO Bool
tests = checkParallel $$(discover)


{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module DFT.Plan
  ( module HM
  , DFTPlanID(..)
  , DFTType(..)
  , DFTPlan
  , FFTWLock
  , getFFTWLock
  , importFFTWWisdom
  , exportFFTWWisdom
  , getDFTPlan
  , emptyPlan
  , checkPlanID
  , dftExecuteWithPlan
  , dftExecuteBatch
  , dftExecuteBatchP
  , dftExecute
  , dft2dPlan
  , idft2dPlan
  , dft1dGPlan
  , idft1dGPlan
  ) where

import           Control.Concurrent.Async
import           Control.Concurrent.MVar
import           Control.Monad                as M
-- import           Control.Monad.Parallel       as MP
import           Data.Bits                    (Bits, complement, (.&.), (.|.))
import           Data.Complex
import           Data.Hashable
import           Data.HashMap.Strict          as HM
import           Data.List                    as L
import           Data.Vector.Storable         as VS
import           Data.Vector.Storable.Mutable as VSM
import           DFT.FFI
import           Foreign.C
import           Foreign.Marshal.Array
import           GHC.Generics                 (Generic)

-- The type definitions of Flag and Sign are borrowed from the fft package.
-- This module allows shared plan execuation so that dft can run parallelly.
newtype Flag = Flag
  { unFlag :: FFTWFlag
  } deriving (Eq, Show, Num, Bits)

data Sign
  = DFTForward
  | DFTBackward
  deriving (Eq, Show)

measure :: Flag
measure = Flag c_measure

unSign :: Sign -> FFTWSign
unSign DFTForward  = c_forward
unSign DFTBackward = c_backward

data DFTType
  = DFT1DG
  | IDFT1DG
  | DFT2D
  | IDFT2D
  deriving (Show, Eq, Generic)

instance Hashable DFTType

data DFTPlanID = DFTPlanID
  { dftType :: DFTType
  , dftDims :: [Int]
  , dftIdx  :: [Int]
  } deriving (Show, Eq, Generic)

instance Hashable DFTPlanID

type DFTPlan = HashMap DFTPlanID Plan

type FFTWLock = MVar ()

{-# INLINE getFFTWLock #-}
getFFTWLock :: IO FFTWLock
getFFTWLock = newMVar ()

{-# INLINE importFFTWWisdom #-}
importFFTWWisdom :: FilePath -> IO ()
importFFTWWisdom filePath = newCString filePath >>= c_import_wisdom_filename >> return ()

{-# INLINE exportFFTWWisdom #-}
exportFFTWWisdom :: FilePath -> IO ()
exportFFTWWisdom filePath = newCString filePath >>= c_export_wisdom_filename >> return ()

{-# INLINE emptyPlan #-}
emptyPlan :: DFTPlan
emptyPlan = HM.empty

{-# INLINE checkPlanID #-}
checkPlanID :: DFTPlan -> DFTPlanID -> Bool
checkPlanID plan planID =
  case HM.lookup planID plan of
    Nothing -> False
    Just p -> True

{-# INLINE getDFTPlan #-}
getDFTPlan :: DFTPlan -> DFTPlanID -> Plan
getDFTPlan plan planID =
  case HM.lookup planID plan of
    Nothing -> error $ "getDFTPlan: couldn't find plan for ID " L.++ show planID
    Just p -> p

{-# INLINE dftExecuteWithPlan #-}
dftExecuteWithPlan ::
     DFTPlanID
  -> Plan
  -> VS.Vector (Complex Double)
  -> IO (VS.Vector (Complex Double))
dftExecuteWithPlan planID@(DFTPlanID t d i) plan vec = do
  v <- VSM.new (L.product d)
  VS.unsafeWith vec $ \ip -> VSM.unsafeWith v $ \op -> c_execute_dft plan ip op
  u <- VS.freeze v
  if t == IDFT1DG || t == IDFT2D
    then if L.null i
           then return $ VS.map (/ (fromIntegral . L.product $ d)) u
           else if size == 0
                  then error $
                       "dftExecuteWithPlan: dimension error.\n" L.++ "dims: " L.++
                       show d L.++
                       "\nIdx: " L.++
                       show i
                  else return $ VS.map (/ (fromIntegral size)) u
    else return u
  where
    size = L.product . L.take (L.last i - L.head i + 1) . L.drop (L.head i) $ d

{-# INLINE dftExecuteBatch #-}
dftExecuteBatch ::
     DFTPlan
  -> DFTPlanID
  -> [VS.Vector (Complex Double)]
  -> IO [VS.Vector (Complex Double)]
dftExecuteBatch hashMap planID@(DFTPlanID t d i) vecs =
  case HM.lookup planID hashMap of
    Nothing ->
      error $
      "dftExecuteBatch: couldn't find plan for ID " L.++ show planID L.++ "\nin\n" L.++
      (show hashMap)
    Just plan -> M.mapM (dftExecuteWithPlan planID plan) vecs

{-# INLINE dftExecuteBatchP #-}
dftExecuteBatchP ::
     DFTPlan
  -> DFTPlanID
  -> [VS.Vector (Complex Double)]
  -> IO [VS.Vector (Complex Double)]
dftExecuteBatchP hashMap planID@(DFTPlanID t d i) vecs =
  case HM.lookup planID hashMap of
    Nothing ->
      error $ "dftExecuteBatch: couldn't find plan for ID " L.++ show planID
    Just plan -> mapConcurrently (dftExecuteWithPlan planID plan) vecs

{-# INLINE dftExecute #-}
dftExecute ::
     DFTPlan
  -> DFTPlanID
  -> VS.Vector (Complex Double)
  -> IO (VS.Vector (Complex Double))
dftExecute hashMap planID@(DFTPlanID t d i) vec =
  case HM.lookup planID hashMap of
    Nothing -> error $ "dftExecute: couldn't find plan for ID " L.++ show planID L.++ "\n" L.++ show (keys hashMap)
    Just plan -> dftExecuteWithPlan planID plan vec

{-# INLINE dft2dG #-}
dft2dG ::
     MVar ()
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> Sign
  -> Flag
  -> IO (Plan, VS.Vector (Complex Double))
dft2dG lock' rows cols vec sign flag = do
  v <- VSM.new (rows * cols)
  x <- takeMVar lock'
  p <-
    VS.unsafeWith vec $ \ip ->
      VSM.unsafeWith v $ \op -> do
        c_plan_dft_2d
          (fromIntegral cols)
          (fromIntegral rows)
          ip
          op
          (unSign sign)
          (unFlag flag)
  putMVar lock' x
  c_execute p
  u <- VS.freeze v
  return (p, u)

{-# INLINE dft2dPlan #-}
dft2dPlan ::
     FFTWLock
  -> DFTPlan
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> IO (DFTPlan, VS.Vector (Complex Double))
dft2dPlan lock' hashMap rows cols vec =
  case HM.lookup planID hashMap of
    Nothing -> do
      (p, v) <- dft2dG lock' rows cols vec DFTForward measure
      return (HM.insert planID p hashMap, v)
    Just p -> do
      v <- dftExecuteWithPlan planID p vec
      return (hashMap, v)
  where
    planID = DFTPlanID DFT2D [cols, rows] []

{-# INLINE idft2dPlan #-}
idft2dPlan ::
     FFTWLock
  -> DFTPlan
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> IO (DFTPlan, VS.Vector (Complex Double))
idft2dPlan lock' hashMap rows cols vec =
  case HM.lookup planID hashMap of
    Nothing -> do
      (p, v) <- dft2dG lock' rows cols vec DFTBackward measure
      return
        (HM.insert planID p hashMap, VS.map (/ (fromIntegral $ rows * cols)) v)
    Just p -> do
      v <- dftExecuteWithPlan planID p vec
      return (hashMap, v)
  where
    planID = DFTPlanID IDFT2D [cols, rows] []

-- This is a generalied 1d dft, the 1 dimension can be many dimensions
-- which are ascending ordered and continued. For example, given a N
-- dimension array, the generalized 1d dft dimension is
-- either [0,1..M] or [M,M+1.. N-1], where 0 <= M and M <= N-1
-- the dimension corresponding to the largest index spins the fastest.
{-# INLINE dft1dGGeneric #-}
dft1dGGeneric ::
     MVar ()
  -> [Int]
  -> [Int]
  -> VS.Vector (Complex Double)
  -> Sign
  -> Flag
  -> IO (Plan, VS.Vector (Complex Double))
dft1dGGeneric lock' dims dftIndex vec sign flag
  | L.and (L.zipWith (\a b -> a + 1 == b) dftIndex . L.tail $ dftIndex) &&
      (not . L.null $ dftIndex) &&
      (L.head dftIndex == 0 || L.last dftIndex == rank - 1) = do
    v <- VSM.new . L.product $ dims
    x <- takeMVar lock'
    p <-
      VS.unsafeWith vec $ \ip ->
        VSM.unsafeWith v $ \op ->
          withArray (L.map fromIntegral dftDims) $ \n -> do
            let totalNum = L.product dims
                dftNum = L.product dftDims
                stride =
                  if L.last dftIndex == rank - 1
                    then 1
                    else L.product . L.drop (1 + L.last dftIndex) $ dims
                dist =
                  if L.last dftIndex == rank - 1
                    then dftNum
                    else 1
            c_plan_many_dft
              (fromIntegral dftRank)
              n
              (fromIntegral $ div totalNum dftNum)
              ip
              n
              (fromIntegral stride)
              (fromIntegral dist)
              op
              n
              (fromIntegral stride)
              (fromIntegral dist)
              (unSign sign)
              (unFlag flag)
    c_execute p
    putMVar lock' x
    u <- VS.freeze v
    return (p, u)
  | otherwise =
    error $
    "dft1dG: dimension list doesn't satisify the restriction of generalized 1d dft.\n" L.++
    "dims: " L.++
    show dims L.++
    "\ndftIndex: " L.++
    show dftIndex
  where
    rank = L.length dims
    dftRank = L.length dftIndex
    dftDims =
      L.take (L.last dftIndex - L.head dftIndex + 1) . L.drop (L.head dftIndex) $
      dims

{-# INLINE dft1dGPlan #-}
dft1dGPlan ::
     FFTWLock
  -> DFTPlan
  -> [Int]
  -> [Int]
  -> VS.Vector (Complex Double)
  -> IO (DFTPlan, VS.Vector (Complex Double))
dft1dGPlan lock' hashMap dims dftIndex vec = do
  when
    (L.product dims /= VS.length vec)
    (error "dft1dGPlan: input vector length and dims don't match.")
  case HM.lookup planID hashMap of
    Nothing -> do
      (p, v) <- dft1dGGeneric lock' dims dftIndex vec DFTForward measure
      return (HM.insert planID p hashMap, v)
    Just p -> do
      v <- dftExecuteWithPlan planID p vec
      return (hashMap, v)
  where
    planID = DFTPlanID DFT1DG dims dftIndex

{-# INLINE idft1dGPlan #-}
idft1dGPlan ::
     FFTWLock
  -> DFTPlan
  -> [Int]
  -> [Int]
  -> VS.Vector (Complex Double)
  -> IO (DFTPlan, VS.Vector (Complex Double))
idft1dGPlan lock' hashMap dims dftIndex vec = do
  when
    (L.product dims /= VS.length vec)
    (error "idft1dGPlan: input vector length and dims don't match.")
  case HM.lookup planID hashMap of
    Nothing -> do
      (p, v) <- dft1dGGeneric lock' dims dftIndex vec DFTBackward measure
      return (HM.insert planID p hashMap, VS.map (/ size) v)
    Just p -> do
      v <- dftExecuteWithPlan planID p vec
      return (hashMap, v)
  where
    planID = DFTPlanID IDFT1DG dims dftIndex
    size =
      fromIntegral .
      L.product .
      L.take (L.last dftIndex - L.head dftIndex + 1) . L.drop (L.head dftIndex) $
      dims


{-# language
        ScopedTypeVariables
  #-}

module Main (main) where

import Test.QuickCheck.Classes
import Test.QuickCheck (Arbitrary(..))

import Data.Complex (Complex(..))
import GHC.Fingerprint (Fingerprint(..))
import GHC.Real (Ratio(..))
import Data.Proxy (Proxy(..))

import Data.Primitive.Instances ()

main :: IO ()
main = lawsCheckMany
  [ ("Complex Double", [primLaws (Proxy :: Proxy (Complex Double))])
  , ("Ratio Double", [primLaws (Proxy :: Proxy (Ratio Int))])
  , ("Fingerprint", [primLaws (Proxy :: Proxy Fingerprint)])
  ]

instance Arbitrary Fingerprint where
  arbitrary = Fingerprint <$> arbitrary <*> arbitrary


module Main where
import System.Environment (getArgs)
import Control.Concurrent.Async
import Control.Concurrent hiding (yield)
import Control.Monad
import System.Posix.Signals
import Pipes
import Pipes.Concurrent as PC
import qualified Pipes.Group as PG
import qualified Pipes.Prelude as P
import Sound.Pulse.Simple
import Data.WAVE
import qualified Data.List.Extra as E
import qualified Numeric.FFT as F
import Data.Complex
import Data.Monoid
import Data.Int
import Data.Maybe
import qualified UI.HSCurses.CursesHelper as CH
import qualified UI.HSCurses.Curses as C

takeStep :: Int -> [a] -> [a]
takeStep _ [] = []
takeStep n (x:xs)
  | n >= 0 = x : takeStep n (drop (n-1) xs)
  | otherwise = takeStep (-n) (reverse xs)

rjust :: Int -> a -> [a] -> [a]
rjust size filler inp
  | length inp >= size = inp
  | otherwise = inp ++ replicate (size - length inp) filler

fft :: Int -> Pipe [WAVESample] [Double] IO()
fft chunkSize = forever $ do
    let step = 18
    let fftNum = 128
    a <- await
    let vec = rjust fftNum 0 $takeStep step $ takeStep 2 a
    let spec = map magnitude $ F.dft(map fromIntegral $ vec)
    let ratio = 3e-9
    yield $ map (\x -> x * ratio) spec

convSpectrumMat :: [Double] -> Int -> Int -> [(Int, String)]
convSpectrumMat input offsetY height =
    serializer 0
    where
        arr = rjust 1 0 $ map (\x -> floor x) input
        serializer y
            | height - offsetY <= y= []
            | otherwise  = [(abs(height - y), map (\i -> if i>y then '#' else ' ') arr)]++serializer (y+1)

spectrogram :: Int -> Int -> String-> Consumer [Double] IO ()
spectrogram chunkNo chunkNum filename = forever $ do
    (height, width) <-lift $ C.scrSize
    a <- await
    let arr = rjust 1 0 a
    let offsetX = 0
    let offsetY = 4
    let mat = convSpectrumMat arr offsetY (height-2)
    let cprintLn y str = C.mvWAddStr C.stdScr y offsetX str
    lift $ mapM (\(y, s) -> cprintLn y $take width s) mat
    lift $ cprintLn 0 ("     chunkNo:"++(show chunkNo) ++ "/" ++ show chunkNum)
    lift $ cprintLn 1 ("max spectrum:"++(show $ maximum arr ))
    lift $ cprintLn 2 ("num spectrum:"++(show $ length arr ))
    lift $ cprintLn 3 ("    filename:"++filename)
    lift $ C.refresh
    spectrogram (chunkNo+1) chunkNum filename

concat' :: [[a]] -> [a]
concat' xss = foldr (++) [] xss

main = do
    tid <- myThreadId
    let handler = do
            putStrLn "keybord interupted"
            killThread tid
    installHandler keyboardSignal (Catch handler) Nothing
    (waveFileName:_) <- getArgs
    input <- getWAVEFile waveFileName
    let header = waveHeader input
    let resolution = 2^64 :: Double
    let channels = waveNumChannels header
    let samples = waveSamples input
    let sampleLength = fromMaybe 0 $waveFrames header
    let chunkSize = 2^11
    let chunkNum = (sampleLength `div` chunkSize) + 1

    s<-simpleNew Nothing "ninja" Play Nothing "youjo"
        (SampleSpec (S32 LittleEndian) (waveFrameRate header) 2) Nothing Nothing
        -- WAVESample is 32 bit, so you must set S32 in pulseaudio or casting.

    --runEffect $ each (E.chunksOf 2048 samples) >-> P.map concat' >-> P.mapM_ (simpleWrite s)

    (plotterOutput, plotterInput) <- spawn unbounded
    player <- async $ do
        runEffect $ each (((E.chunksOf (chunkSize)).(map$ map (`div` 3))) samples) >-> P.map concat' >-> P.chain(simpleWrite s) >-> toOutput plotterOutput
        performGC
    plotter <- async $ do
        CH.start
        C.refresh
        runEffect $ fromInput plotterInput >-> fft chunkSize>-> spectrogram 0 chunkNum waveFileName
        C.endWin
        performGC
    mapM_ wait [player, plotter]

    simpleDrain s
    simpleFree s


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.Image (
  Image(..)
, Image8
, map
, (!)
, (!?)
, createImg
, createImgI
, splitColor
, mergeColor
, toDouble
, fromDouble
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C
import Data.Vector.Image.Color.RGB

data Image a =
  Image {
    img_width :: Int
  , img_height :: Int
  , img_dat :: Vector a
  } deriving (Show,Eq)

type Image8 = Image (RGB Word8)

map :: (Storable a, Storable b) => (a->b) -> Image a -> Image b
map func (Image w h d) = (Image w h (V.map func d))

index :: Storable a => Image a -> (Int,Int) -> a
index (Image w h d) (x,y) = d V.! (w*y+x)

(!?) :: Storable a => Image a -> (Int,Int) -> Maybe a
(!?) img@(Image w h d) (x,y) | x < 0 || x >= w || y <0 || y>= h = Nothing
                             | otherwise = Just (index img (x,y))

(!) :: (Storable a,Integral a) => Image a -> (Int,Int) -> a
(!) img idx = fromMaybe 0 (img !? idx)

-- | Generate Image
createImg :: Int -- ^ Width of image
            -> Int -- ^ Height of image
            -> (Int -> Int -> RGB8) -- ^ Pixel Generator
            -> Image8
createImg w h func = Image w h $ V.fromList $ do
  y <- [0..(h-1)]
  x <- [0..(w-1)]
  return $ func x y

-- | Generate Image
createImgI :: Int -- ^ Width of image
           -> Int -- ^ Height of image
           -> (Int -> Int -> RGB Int) -- ^ Pixel Generator
           -> Image8
createImgI w h func = Image w h $ V.fromList $ do
  y <- [0..(h-1)]
  x <- [0..(w-1)]
  return $ fmap fromIntegral $ fmap (\v -> if v < 0 then 0 else if v > 255 then 255 else v) $ func x y

splitColor :: (ColorSpace c a,Storable (c a),Storable a)
           => Image (c a)
           -> [Image a]
splitColor img@(Image w h d) = flip P.map [0..(num-1)] $ (\idx -> Image w h $ V.map (\v -> toList v !! idx) d)
  where
    num :: Int
    num = dim (dummyColor img)
    dummyColor :: Image (c a) -> c a
    dummyColor = error "do not eval this"

mergeColor :: (ColorSpace c a,Storable (c a),Storable a,Element a)
              => [Image a]
              -> Image (c a)
mergeColor img@((Image w h d):_) = Image w h $ V.fromList $ P.map fromList $ L.transpose $ P.map (V.toList.img_dat) img
mergeColor [] = error "image-list is null."

toDouble :: (Integral a,Storable a) => Image a -> Image Double
toDouble img= map fromIntegral img

fromDouble :: (Integral a,Storable a) => Image Double -> Image a
fromDouble img = map truncate img


{-# LANGUAGE TemplateHaskell, RankNTypes #-}

module Geometry.LRS (
    -- * Types
    Dictionary,
    -- * Functions
    getDictionary,
    selectPivot,
    reverseRS,
    lrs,
    polytope2LiftedCone,
    liftedCone2polyotpe,
    lrsFacetEnumeration,

    colFromList
    )
 where

import Geometry.Vertex
import Geometry.Matrix

import Data.Matrix as M hiding (trace)
import qualified Data.Vector as V
import Data.Maybe
import Data.List
import Util
import Control.Lens
import Debug.Trace
import Numeric.LinearProgramming hiding (simplex)
import Numeric.LinearAlgebra.HMatrix as HM
import Data.List

type Row = M.Matrix Rational
type Col = M.Matrix Rational

data Dictionary = Dict {
                            __B :: [Int],
                            __N :: [Int],
                            _dict :: M.Matrix Rational
                        }
    deriving (Show, Eq)

makeLenses ''Dictionary

numRows :: Dictionary -> Int
numRows dictionary = nrows $ dictionary^.dict

numCols :: Dictionary -> Int
numCols dictionary = ncols $ dictionary^.dict

(|*|) = multStrassen

colFromList :: [a] -> M.Matrix a
colFromList = colVector . V.fromList

rowFromList :: [a] -> M.Matrix a
rowFromList = rowVector . V.fromList

-- | Wrapper of submatrix function. This function works for indices starting from 0 rather than 1 in the submatrix function
submatrix' ::
    (Int, Int)      -- ^ Rows indices
    -> (Int, Int)   -- ^ Cols indices
    -> M.Matrix a
    -> M.Matrix a
submatrix' (ro,rk) (co,ck) m = submatrix (ro+1) (rk+1) (co+1) (ck+1) m

-- | Lens for columns
colAt :: Int -> Lens' (M.Matrix Rational) Col
colAt j = lens (getCol' j) (\m c -> setCol' j c m)

getCol' :: Int -> M.Matrix a -> M.Matrix a
getCol' idx = colVector . (getCol (idx+1))

setCol' :: Int -> Col -> M.Matrix Rational -> M.Matrix Rational
setCol' idx col mat
    | idx == 0 = col <|> right
    | idx == (ncols mat) - 1 = left <|> col
    | otherwise = left <|> col <|> right
    where
        left = submatrix' (0,(nrows mat)-1) (0,idx-1) mat
        right = submatrix' (0,(nrows mat)-1) (idx+1, (ncols mat)-1) mat

-- | Lens for rows
rowAt :: Int -> Lens' (M.Matrix Rational) Row
rowAt i = lens (getRow' i) (\m r -> setRow' i r m)

getRow' :: Int -> M.Matrix a -> M.Matrix a
getRow' idx = rowVector . (getRow (idx+1))

setRow' :: Int -> Col -> M.Matrix Rational -> M.Matrix Rational
setRow' idx row mat
    | idx == 0 = row <-> down
    | idx == (nrows mat) - 1 = up <-> row
    | otherwise = up <-> row <-> down

    where
        up = submatrix' (0,idx-1) (0,(ncols mat)-1)  mat
        down = submatrix' (idx+1, (nrows mat)-1) (0,(ncols mat)-1) mat

-- | Lens for accessing elements

elemAt :: (Int, Int) -> Lens' (M.Matrix a) a
elemAt (i, j) = lens (getElem' (i,j)) (\m x -> setElem' x (i,j) m)

getElem' :: (Int, Int) -> M.Matrix a -> a
getElem' (row, col) = getElem (row+1) (col+1)

setElem' :: a -> (Int, Int) -> M.Matrix a -> M.Matrix a
setElem' elem (row, col) mat
    | row < 0 || col < 0 || row >= nrows mat || col >= ncols mat = error ("setElem': Trying to set at (" ++ show row ++ "," ++ show col ++ ") in a " ++ show (nrows mat) ++ "x" ++ show (ncols mat) ++ " matrix.")
    | otherwise = setElem elem (row+1, col+1) mat

mapCol' :: (Int -> a -> a) -> Int -> M.Matrix a -> M.Matrix a
mapCol' f col m
    | col < 0 || col > (ncols m)-1 = error "mapCol': index out of bounds"
    | otherwise = mapCol f (col+1) m

mapRow' :: (Int -> a -> a) -> Int -> M.Matrix a -> M.Matrix a
mapRow' f row m
    | row < 0 || row > (nrows m)-1 = error "mapCol': index out of bounds"
    | otherwise = mapRow f (row+1) m

getOptimumVertex :: M.Matrix Rational -> Col -> Maybe Vertex
getOptimumVertex mat col = fmap (map toRational) $ feasNOpt result
    where
        problem = Maximize $ replicate (ncols mat) (-1)
        constraints = Dense $ safeZipWith (:<=:) (map (map fromRational) $ M.toLists mat) (map fromRational $ concat $ M.toLists col)
        result = exact problem constraints (map Free [1..ncols mat])
        feasNOpt (Optimal (_, vertex)) = Just vertex
        feasNOpt _ = Nothing

{-
    Dictionary form

                p21     p31
    invA_b
                p22     p32

 -}

toHMatrix :: M.Matrix Rational -> HM.Matrix Double
toHMatrix matrix = HM.matrix cols (map (fromRational) (concat $ M.toLists matrix))
    where
        cols = ncols matrix

findBasis :: (Maybe (M.Matrix Rational), Maybe (M.Matrix Rational)) -> M.Matrix Rational -> Int -> (Maybe (M.Matrix Rational), Maybe (M.Matrix Rational))
findBasis (basic, cobasic) original colNum
    | currRank == rows && colNum < cols = findBasis (basic, Just newCobasic) original (colNum + 1)-- Just ((fromJust cobasic) <|> (submatrix' (0,rows-1) (colNum, (ncols original)-1) original)) )
    | currRank == rows = (basic, cobasic)-- Just ((fromJust cobasic) <|> (submatrix' (0,rows-1) (colNum, (ncols original)-1) original)) )
    | currRank == candRank = findBasis (basic, Just newCobasic) original (colNum + 1)
    | currRank < candRank = findBasis (Just newBasic, cobasic) original (colNum + 1)
    where
        rows = nrows original
        cols = ncols original
        currRank = if isNothing basic then 0 else rank (toHMatrix $ fromJust basic)
        candVec = original ^. colAt colNum
        candMat = if isNothing basic then candVec else (fromJust basic) <|> candVec
        candRank = rank $ toHMatrix candMat

        newBasic = if isNothing basic then candVec else (fromJust basic) <|> candVec
        newCobasic = if isNothing cobasic then candVec else (fromJust cobasic) <|> candVec

sortSystem :: M.Matrix Rational -> (M.Matrix Rational, M.Matrix Rational)
sortSystem matrix = trace ("Matrix : "++ show matrix ++ "\n\nSorted " ++ show (findBasis (Nothing, Nothing) matrix 0) ) (((,) <$> (fromJust . fst) <*> (fromJust . snd)) (findBasis (Nothing, Nothing) matrix 0))
-- sortSystem :: M.Matrix Rational -> Col -> Vertex -> (M.Matrix Rational, Col)

-- sortSystem mat col vertex
--     | cols /=  (rows `div` 2) = (,) mat col
--     | otherwise = (trace $ show $ M.fromLists (upper ++ lower)) (,) (M.fromLists (upper ++ lower)) col
--     where
--         matLists = M.toLists mat

--         rows = nrows mat
--         cols = ncols mat
--         slackMatr = submatrix' (0,rows-1) (0, rows - cols -2) $ identity rows
--         lower  = getIndependent matLists [] 0 rows slackMatr
--         upper = matLists \\ lower

-- getDictionary :: Matrix Rational -> Col -> Vertex -> Dictionary
-- getDictionary _A b vertex = trace ("In dictionary" ++ show (rows,cols, nrows dictionary, ncols dictionary) ++ show newDict) Dict [0..rows] [rows+1..rows+cols] newDict
--     -- trace ("In dictionary" ++ show _A_B)
--     where
--         (newA, newb) = sortSystem _A b vertex
--         rows = nrows newA
--         cols = ncols newA
--         slack = identity rows
--        -- dictionary = newA <|> slack
--         topRow = rowFromList $ 1 : replicate rows 0 ++  replicate cols (1)
--         leftCol = colFromList $ replicate rows 0
--         dictionary = topRow <-> (leftCol <|> slack <|> newA )

--         -- c_B = submatrix' (0,0) (1,rows) topRow
--         -- c_N = submatrix' (0,0) (rows+1, rows+cols) topRow
--         -- _A_B = submatrix' (0,rows-1) (0,rows-1) dictionary
--         -- _A_N = submatrix' (0,rows-1) (rows, rows+cols-1) dictionary
--         _A_B = submatrix' (0,rows) (0,rows) dictionary
--         _A_N = submatrix' (0,rows) (rows+1, rows+cols) dictionary
--         Right invA_B = inverse _A_B

--         p2 = - invA_B |*| _A_N
--         p3 = invA_B |*| newb
--         -- p21 = (c_N - c_B |*| invA_B |*| _A_N)
--         -- p22 = invA_B |*| _A_N
--         -- p31 = -c_B |*| invA_B |*| newb
--         -- p32 = invA_B |*| newb
--         newDict = ((identity (rows+1)) <|> p2 <|> p3)

-- | Constructs initial dictionary given a set of inequalities and an optimum initial vertex
-- getDictionary :: M.Matrix Rational -> Col -> Vertex -> Dictionary
-- getDictionary _A b vertex = Dict [0..rows] [rows+1..rows+cols] newDict
--     where
--         -- (newA, newb) = sortSystem _A b vertex
--         initialDict = _A <|> slack
--         (basic, cobasic) = sortSystem initialDict
--         (newA, newb) = (,) _A b
--         rows = nrows newA
--         cols = ncols newA
--         slack = identity rows
--         dictionary = newA <|> slack
--         topRow = rowFromList $ 1 :  replicate rows 0 ++  replicate cols 1
--         c_B = submatrix' (0,0) (1,rows) topRow
--         c_N = submatrix' (0,0) (rows+1, rows+cols) topRow
--         _A_B = submatrix' (0,rows-1) (0,rows-1) dictionary
--         _A_N = submatrix' (0,rows-1) (rows, rows+cols-1) dictionary
--         Right invA_B = inverse _A_B
--         p21 = (c_N - c_B |*| invA_B |*| _A_N)
--         p22 = invA_B |*| _A_N
--         p31 = -c_B |*| invA_B |*| newb
--         p32 = invA_B |*| newb

--         newDict = ((identity (rows+1)) <|> (p21 <-> p22) <|> (p31 <-> p32))

getDictionary :: M.Matrix Rational -> Col -> Vertex -> Dictionary
getDictionary _A b vertex = Dict [0..rows] [rows+1..rows+cols] newDict
    where

        -- (newA, newb) = sortSystem _A b vertex
        initialDict = _A <|> slack
        (basic, cobasic) = sortSystem initialDict
        (newA, newb) = (,) _A b
        rows = nrows newA
        cols = ncols newA
        slack = identity rows
        -- dictionary = newA <|> slack
        topRow = rowFromList $ 1 :  replicate rows 0 ++  replicate cols 1
        c_B = submatrix' (0,0) (1,rows) topRow
        c_N = submatrix' (0,0) (rows+1, rows+cols) topRow
        _A_B = basic --submatrix' (0,rows-1) (0,rows-1) dictionary
        _A_N = cobasic --submatrix' (0,rows-1) (rows, rows+cols-1) dictionary
        Right invA_B = inverse _A_B
        p21 = (c_N - c_B |*| invA_B |*| _A_N)
        p22 = invA_B |*| _A_N
        p31 = -c_B |*| invA_B |*| newb
        p32 = invA_B |*| newb

        newDict = ((identity (rows+1)) <|> (p21 <-> p22) <|> (p31 <-> p32))

enteringVariable :: Dictionary -> Maybe Int
enteringVariable dictionary
    | null negs = Nothing
    | otherwise = trace ("ENTERINGVAR\nBASE: " ++ show (dictionary^._B) ++ "\nCOBASE: " ++ show (dictionary^._N))  (Just $ (fst.head.sort) negs)
    where
        rows = nrows $ dictionary^.dict
        cobasic_0 = zip (dictionary^._N) (map (\j -> dictionary^.dict.elemAt (0,j)) [rows..])
        negs = filter (\(_,value) -> value < 0) cobasic_0

lexMinRatio :: Dictionary -> Int -> Int
lexMinRatio dictionary s
    | dim+1 == rows = 0
    | null indexed_s = 0
    | otherwise = trace ("LEXMINRATIO\nBASE: " ++ show (dictionary^._B) ++ "\nCOBASE: " ++ show (dictionary^._N))  (dictionary ^. _B) !! (fst $ indexed_s !! (fromJust $ elemIndex (minimum ratios) ratios))
    where
        rows = numRows dictionary
        cols = numCols dictionary
        dim = cols - rows - 1
        dictMatrix = dictionary^.dict
        col_s = dictMatrix ^. colAt s
        _D = (dictMatrix ^. colAt (cols-1)) <|> submatrix' (0,rows-1) (0, rows-1) dictMatrix
        slice_s = (concat.M.toLists) $ submatrix' (dim+1, rows-1) (0,0) col_s
        indexed_s = filter (snd.(fmap (>0))) $ safeZipWith (,) [dim+1..rows-1] slice_s
        sub_D = map (head . M.toLists . (\i -> _D ^. rowAt i) . fst) indexed_s
        ratios = safeZipWith (map $) (map ((flip (/)) . snd) indexed_s) sub_D

-- | Selects entering and leaving variables for the next pivot in the dictionary
selectPivot :: Dictionary -> Maybe (Int, Int)
selectPivot dictionary
    | s == Nothing = Nothing
    | otherwise = Just (r, fromJust s)
    where
        s = enteringVariable dictionary
        r = lexMinRatio dictionary (fromJust s)

pivot :: Int ->Int -> Dictionary -> Dictionary
pivot r s dictionary = Dict new_B new_N (_E |*| dictMatrix)
    where
        dictMatrix = dictionary ^. dict
        col_s = dictMatrix ^. colAt (s)
        t = fromJust $ elemIndex r (dictionary ^. _B) -- idxR
        idxS = fromJust $ elemIndex s (dictionary ^. _N)
        a_t = col_s ^. elemAt (t,0)
        eta = (mapCol' (\_ x -> -x/a_t) 0 col_s) & elemAt (t,0) .~ (1/a_t)
        _E = (identity (nrows dictMatrix)) & colAt t .~ eta
        new_B = dictionary^._B & element t .~ s
        new_N = dictionary^._N & element idxS .~ r

simplex :: Dictionary -> Dictionary
simplex dictionary
    | idxsToPivot == Nothing = dictionary
    | otherwise = simplex $ uncurry pivot (fromJust idxsToPivot) dictionary
    where
        idxsToPivot = selectPivot dictionary

-- | Reverse search.
reverseRS ::
        Dictionary
    ->  Int          -- ^ element in N
    ->  Maybe Int
reverseRS dictionary v
    | conditions == False = trace ("REVERSE RS: " ++ show [firstCondition, sndCondition, lastCondition]) Nothing
    | conditions == True = trace ("REVERSE RS: " ++ show v) Just u
    where
        dictMatrix = dictionary^.dict
        rows = nrows dictMatrix
        v_col = dictMatrix ^. colAt (v)
        w_row_0 = dictMatrix ^. rowAt 0
        u = lexMinRatio dictionary v
        i = fromJust $ elemIndex u (dictionary ^. _B)
        w_row_i = mapRow' (\_ x -> (v_col ^. elemAt (0,0))/(v_col ^. elemAt (i,0)) * x) 0 $ dictMatrix ^. rowAt (i)
        diff_ws = (head.M.toLists) $ w_row_0 - w_row_i
        firstCondition = (w_row_0 ^. elemAt (0,v)) > 0
        sndCondition = u /= 0
        lastCondition = all (>=0) [(diff_ws!!(snd idx))|idx <- (zip (dictionary^._N) [rows..]), (fst idx) < u ]
        conditions = firstCondition && sndCondition && lastCondition -- Proposition 6.1

-- reverseRS :: -- reverse with pivot and selectPivot
--     Dictionary ->
--     Int ->
--     Maybe Int
-- reverseRS dictionary v
--     | newPivots == Nothing = Nothing
--     | condition == False = Nothing
--     | condition == True = Just u
--     where
--         u = lexMinRatio dictionary v
--         prev_B = pivot u v dictionary
--         newPivots = selectPivot prev_B
--         condition = fst (fromJust newPivots) == v && (snd (fromJust newPivots)) == u

getVertex :: Dictionary -> Vertex
getVertex dictionary = concat $ M.toLists $ submatrix' (1,dim) (cols-1, cols-1) (dictionary^.dict)
    where
        rows = numRows dictionary
        cols = numCols dictionary
        dim = cols - rows - 1

revSearch :: Dictionary -> [Vertex]
revSearch dictionary@(Dict _B _N dictMatrix) -- = getVertex dictionary : concatMap revSearch pivoted
    | (not.null) possibleRay = possibleRay ++ (concatMap revSearch pivoted)
    | otherwise = getVertex dictionary : concatMap revSearch pivoted
    where
        rows = numRows dictionary
        cols = numCols dictionary
        valid_N = [i | i <- _N, (reverseRS dictionary i) /= Nothing]
        valid_B = map (lexMinRatio dictionary) valid_N
        pivoted = map (\(r, s) ->  pivot r s dictionary) $ zip valid_B valid_N
        possibleRay = hasRay dictionary

-- lexMin :: Dictionary -> Maybe Extremal
-- lexMin dictionary

hasRay :: Dictionary -> [Vertex]
hasRay dictionary = trace ("RAYS: " ++ show (map snd rays) ++ "\n\n" ++ show (map fst rays) ++ "\n" ++ show (dictionary ^. dict) ++ "\n\n") (map snd rays)
    -- idxsPivot == Nothing = Nothing
    -- r == 0 = Just ray
    -- otherwise = Nothing
    where
        dictMatrix = dictionary ^. dict
        rows = numRows dictionary
        dim = cols - rows - 1
        cols = numCols dictionary
        nonPositive column = all (<=0) ((concat.M.toLists) column)
        colsWithRays = if dim+1 == rows then
                            [(j, dictMatrix^.colAt j) | j <- [rows..cols-2]]
                        else [(j, dictMatrix^.colAt j) | j <- [rows..cols-2], nonPositive (submatrix' (dim+1,rows-1) (j,j) dictMatrix )]
        rays =  zip (map fst colsWithRays) (map (concat . M.toLists . (submatrix' (1,dim) (0,0)) . snd ) colsWithRays)

-- | Implementation of Lexicographic Reverse Search
lrs :: M.Matrix Rational -> Col -> Vertex-> [Vertex]
lrs matrix b vertex = (sort.nub) $ revSearch dictionary
    where
        dictionary = simplex $ getDictionary matrix b vertex

-- | Transforms a V polytope into a H pointed cone
polytope2LiftedCone :: [Vertex] -> (M.Matrix Rational, Col)
polytope2LiftedCone points = (colOnes <|> hMatrix , colZeros)
    where
        nrows = length points
        colOnes = colFromList $ replicate nrows 1
        hMatrix = M.fromLists points
        colZeros = colFromList $ replicate nrows 0

-- | Transforms a V pointed cone into a H polyotpe
liftedCone2polyotpe :: [Vertex] -> (M.Matrix Rational, Col)
liftedCone2polyotpe points = (hMatrix, b)
    where
        dim = length (points !! 0)
        rays  = points \\ [(replicate dim 0)] -- remove vertex at the origin
        b = colFromList $ map (negate.head) rays
        hMatrix = M.fromLists $ map tail rays

-- | Implementation of Lexicographic Reverse Search for Facet Enumeration
lrsFacetEnumeration :: [IVertex] -> (M.Matrix Rational, Col)
lrsFacetEnumeration ipoints = hPolytope
    where
        points =  map (map toRational) ipoints
        cols = length (points!!0)
        initialVertex = replicate (cols + 1) 0
        hLiftedCone@(matrix, b) = polytope2LiftedCone points
        vLiftedCone = lrs matrix b initialVertex
        hPolytope = liftedCone2polyotpe vLiftedCone


module AutoEncoder (
  preTrains,
  preTrain
  ) where

import Numeric.LinearAlgebra
import Data.List
import Common
import Forward
import BackProp
import ActivFunc
import Other

-- TODO: refactor
preTrains :: R -> Int -> (Matrix R -> Matrix R, Matrix R -> Matrix R) -> Matrix R -> [Matrix R] -> [Matrix R]
preTrains rate iter fDf x ws = snd (mapAccumL (preTrain rate iter fDf) x ms) `mappend` [m]
  where (ms, m) = (init ws, last ws)

preTrain :: R -> Int -> (Matrix R -> Matrix R, Matrix R -> Matrix R) -> Matrix R -> Matrix R -> (Matrix R, Matrix R)
preTrain rate iter fDf@(f, df) x w = (nx, nw)
  where
    nx = forward f nw x
    [nw, _] = last . take iter $ iterate (backProp f rate fDf (x, x)) [w, mirrorWeight w]

mirrorWeight :: Matrix R -> Matrix R
mirrorWeight w = w' ||| konst 0 (rows w', 1) -- initial bias: 0
  where w' = tr $ weightWithoutBias w


{-# LANGUAGE Strict
           , OverloadedStrings#-}
module Lib where

import qualified Data.HashTable.IO as H
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Numeric.LinearAlgebra
import qualified Data.Vector.Storable as VS
import System.IO
import Data.List
import Control.Monad
import System.Process
import System.Random
import Text.Read (readMaybe)

someFunc :: IO ()
someFunc = putStrLn "someFunc"

nanCheck :: Double -> Double
nanCheck x = if show x == "NaN" || show x == "Infinity" then 0 else x

st :: Double
st = 10

delta :: Double
delta = 1

chars :: T.Text
chars = "qwertyuiopasdfghjklzxcvbnmęóąśłżźćń\t1234567890 "

el :: Char -> T.Text -> Bool
el = T.any . (==)

preprocess :: T.Text -> [T.Text]
preprocess = map (T.take 20) . T.words . T.map check . T.toLower
             where check x = if x `el` chars then x else ' '

type Probs = H.CuckooHashTable T.Text Double

readWords :: Handle -> IO Probs
readWords handl = do
  dict <- H.newSized 2500000
  loop handl dict
  randomIO >>= H.insert dict "<bias>"
  return dict
  where
    loop :: Handle -> Probs -> IO ()
    loop hdl dic =
        hIsEOF hdl >>= \x -> if x
            then return ()
            else do
              line <- preprocess <$> TIO.hGetLine hdl
              let ins word =
                    if (T.length word > 4)
                    then H.lookup dic word >>=
                                    \mprob -> case mprob of
                                       Just _ -> do
                                         return ()
                                       Nothing -> randomIO >>= H.insert dic word
                    else return ()
              case line of
                _:dat -> mapM_ ins dat
                _     -> hPutStrLn stderr "Can't parse line:" >> hPrint stderr line
              loop hdl dic

saveProbs :: Probs -> IO ()
saveProbs dic = do
  H.mapM_ (\(word, v) -> putStrLn . unwords $ [T.unpack . T.strip $ word, show v]) dic

readProbs :: String -> IO Probs
readProbs filename = do
  dict <- H.newSized 2500000
  withFile filename ReadMode $ readFromFile dict
  return dict
  where
    readFromFile :: Probs -> Handle -> IO ()
    readFromFile dic hdl =
      hIsEOF hdl >>= \x -> if x
                              then return ()
                              else do
                                w <- hGetLine hdl
                                case words w of
                                  [p, pr] ->
                                    H.insert dic (T.pack p) (read pr)
                                  _ -> return ()
                                readFromFile dic hdl

linear :: Vector Double -> Vector Double -> Double
linear = (<.>)

rmse :: Vector Double -> Vector Double -> Double -> Double
rmse weights inp ex = let
  prob = linear weights inp
  in (prob - ex)^2

gradRMSE :: Vector Double -> Vector Double -> Double -> Double -> Vector Double
gradRMSE weights inp ex m = VS.map (\x -> x * (linear weights inp - ex)/m) inp

counts :: Eq a => a -> [a] -> Double
counts = ((fromIntegral. length).).filter.(==)

evalLine :: Probs -> [T.Text] -> IO Double
evalLine dict line = do
  let wek@(_:words) = "<bias>" : (filter ((>4). T.length) . nub $ line)
      inp   = VS.cons 1 (foldr (\c acc -> VS.cons (log . (+1) . counts c $ line) acc) VS.empty words)
  weights <- foldM (\vec c -> do
                       mval <-  H.lookup dict c
                       case mval of
                         Just v -> return $ VS.cons v vec
                         Nothing-> return $ VS.cons 0 vec) VS.empty wek
  return ((1900+).(100*).linear weights $ inp)

updateWeights :: Probs -> [T.Text] -> Double -> Double -> Double -> IO Double
updateWeights dict line ex step m = do
  let wek@(_:words) = "<bias>":(filter ((>4).T.length) . nub $ line)
      inp   = VS.cons 1 (foldr (\c acc -> VS.cons (log . (+1) . counts c $ line) acc) VS.empty words)
  weights <- foldM (\vec c -> do
                       mval <- H.lookup dict c
                       case mval of
                         Just v -> return $ VS.snoc vec v
                         Nothing-> return $ VS.snoc vec 0) VS.empty wek
  let grad       = gradRMSE weights inp ex m
      newWeights = weights - (VS.map (clip.nanCheck.(step *)) grad)
      newError   = rmse newWeights inp ex
  foldM_ (\acc x -> do
            let newW = VS.head acc
            unless (newW == 0) $ H.insert dict x newW
            return $ VS.tail acc) newWeights words
  return $ nanCheck newError
  where
    clip x = max x 3

checkError :: Probs -> [T.Text] -> Double -> IO Double
checkError dict line ex = do
  let wek@(_:words) = "<bias>":nub line
      inp   = VS.cons 1 . foldr (\c acc -> VS.cons (log . (+1) . counts c $ line) acc) VS.empty $ words
  weights <- foldM (\vec c -> do
                       mval <- H.lookup dict c
                       case mval of
                         Just v -> return $ VS.snoc vec v
                         Nothing-> return $ VS.snoc vec 0) VS.empty wek
  return $ rmse weights inp ex

train :: [(Double,[T.Text])] -> Probs -> Int -> IO ()
train corpus dict maxEp = do
  let len     = length corpus
      testset = take (len`div`10) corpus
      trainset= drop (len`div`10) corpus
      trlen   = length trainset
      testlen = length testset
      trainHelp :: Double -> Int -> Double -> IO ()
      trainHelp err epoch step = do
          let trainOnLine (year, dat) = updateWeights dict dat year step (fromIntegral trlen)
          mapM_ trainOnLine trainset
          newErr <- sqrt <$> foldM (\er (ex, line) -> (+(er/(fromIntegral testlen))) <$> checkError dict line ex) 0 testset
          let
            lossRate   = (err - newErr)
            newSteptmp = if lossRate > 0 then step*2 else step/100
            newStep    = if newSteptmp > 20 then 20 else newSteptmp
          hPutStrLn stderr $ "***********************"
          hPutStrLn stderr $ "  Old Error: " ++ show err
          hPutStrLn stderr $ "  New Error: " ++ show newErr
          hPutStrLn stderr $ "  LossRate: " ++ show lossRate
          unless ((abs lossRate <= delta)||(epoch == maxEp)) $ trainHelp newErr (1+epoch) newStep
  startErr <- sqrt  <$> foldM (\er (ex, line) -> (+(er/fromIntegral testlen)) <$> checkError dict line ex) 0 testset
  trainHelp startErr 1 st

loadCorpus :: String -> IO [(Double,[T.Text])]
loadCorpus filename = do
  withFile filename ReadMode (loop [])
  where loop corp hdl =
          hIsEOF hdl >>= \x -> if x
            then return corp
            else do
              line <- preprocess <$> TIO.hGetLine hdl
              newEntry <- case line of
                cl:dat -> case readMaybe . T.unpack $ cl of
                  Just year -> return $ Just ((year-1900)/100, dat)
                  Nothing   -> hPutStrLn stderr "Can't read as a number:" >> hPrint stderr cl >> return Nothing
                _      -> hPutStrLn stderr "Can't parse line:" >> hPrint stderr line >> return Nothing
              case newEntry of
                Nothing -> loop corp hdl
                Just x  -> loop (x:corp) hdl

eval :: Handle -> Probs -> IO ()
eval hdl dict =
  hIsEOF hdl >>= \x ->
    if x
    then return ()
    else do
      line <- preprocess <$> TIO.hGetLine hdl
      pred <- evalLine dict line
      print pred
      eval hdl dict


{-# LANGUAGE QuasiQuotes, MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
module Lib where

import qualified Language.C.Inline as C
import Foreign.C.Types
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearAlgebra.Data as D
import qualified Numeric.LinearAlgebra.Devel as D
import Control.Arrow ((***))
import Data.Tuple (swap)
C.context (C.baseCtx <> C.vecCtx)

C.include "<math.h>"
C.include "<osqp/osqp.h>"

testsp = D.mkSparse [((0,0), 3.0), ((1,1), 2.0)]
testcsr = D.mkCSR [((0,0), 3.0), ((1,1), 2.0)]

-- getCSR (D.SparseR csr ncols nrows) =

assocVec :: (Num a, Ord a) => [a] -> [(Int,a)]
assocVec v = filter (\(_,x) -> x /= 0) (zipWith (,) [0..] v)

-- conveniently, will assume that if your list stops, it is implied the rest of it is zero. Or maybe inconvenitly from a type safety persepctive
assocMat :: (Num a, Ord a) => [[a]] -> [((Int,Int),a)]
assocMat m = let m' = map assocVec m in addRows m' where addRows m'' = concat $ zipWith (\i xs -> map (\(j,x) -> ((i,j),x)) xs) [0 ..] m''
-- AssocMatrix = [((Int, Int), Double)]

transposeAssoc = map (swap *** id)

mkGCSC = D.tr . D.mkSparse . transposeAssoc
mkGCSC' = mkGCSC . assocMat
mkCSC :: D.AssocMatrix -> CSC'
mkCSC = transpose'' . D.mkCSR . transposeAssoc
mkCSC' :: [[Double]] -> CSC'
mkCSC' = mkCSC . assocMat
-- The CSC type is internal to hmatrix.
data CSC' = CSC'
        { cscVals  :: V.Vector Double
        , cscRows  :: V.Vector CInt
        , cscCols  :: V.Vector CInt
        , cscNRows :: Int
        , cscNCols :: Int
        } deriving Show

-- data OSQPCSC = CSC  -- {nzmax :: CInt, m :: CInt, n :: CInt, col_p :: V.Vector Ptr, col  }
{-
instance LA.Transposable D.CSR CSC'
  where
  tr (D.CSR vs cs rs n m) = CSC' vs cs rs m n
  tr' = tr

instance LA.Transposable CSC' D.CSR
  where
  tr (CSC' vs rs cs n m) = D.CSR vs rs cs m n
  tr' = tr
-}
transpose'' (D.CSR vs cs rs n m) = CSC' vs cs rs m n
transpose''' (CSC' vs rs cs n m) = D.CSR vs rs cs m n

funboy :: IO ()
funboy = do
  x <- [C.exp| double{ cos(1) } |]
  print x
someFunc :: IO ()
someFunc = putStrLn "someFunc"

myfun :: CDouble -> CDouble
myfun x = [C.pure| double{ cos( $(double x)) } |]

-- exampleOSQPData =

-- Actually CDoubles
--data OSQPData = OSQPData {dn :: CInt, dm :: CInt, p :: CSC ,
--     a :: CSC, q :: V.Vector CDouble, l :: V.Vector CDouble, u :: V.Vector CDouble}

-- OSQPWorkspace *osqp_setup(const OSQPData *data, OSQPSettings *settings)
-- c_int osqp_solve(OSQPWorkspace *work)
-- c_int osqp_cleanup(OSQPWorkspace *work)

-- Maybe the easiest thing is to just inject default settings and marshal only what we need
-- Get a demo code in full c and start pulling pieces out.
-- l and u are the easiest

demo :: IO CInt
demo = [C.block|
int {
    // Load problem data
    c_float P_x[4] = {4.00, 1.00, 1.00, 2.00, };
    c_int P_nnz = 4;
    c_int P_i[4] = {0, 1, 0, 1, };
    c_int P_p[3] = {0, 2, 4, };
    c_float q[2] = {1.00, 1.00, };
    c_float A_x[4] = {1.00, 1.00, 1.00, 1.00, };
    c_int A_nnz = 4;
    c_int A_i[4] = {0, 1, 0, 2, };
    c_int A_p[3] = {0, 2, 4, };
    c_float l[3] = {1.00, 0.00, 0.00, };
    c_float u[3] = {1.00, 0.70, 0.70, };
    c_int n = 2;
    c_int m = 3;

    // Problem settings
    OSQPSettings * settings = (OSQPSettings *)c_malloc(sizeof(OSQPSettings));

    // Structures
    OSQPWorkspace * work;  // Workspace
    OSQPData * data;  // OSQPData

    // Populate data
    data = (OSQPData *)c_malloc(sizeof(OSQPData));
    data->n = n;
    data->m = m;
    data->P = csc_matrix(data->n, data->n, P_nnz, P_x, P_i, P_p);
    data->q = q;
    data->A = csc_matrix(data->m, data->n, A_nnz, A_x, A_i, A_p);
    data->l = l;
    data->u = u;

    // Define Solver settings as default
    osqp_set_default_settings(settings);
    settings->alpha = 1.0; // Change alpha parameter

    // Setup workspace
    work = osqp_setup(data, settings);

    // Solve Problem
    osqp_solve(work);

    // Cleanup
    osqp_cleanup(work);
    c_free(data->A);
    c_free(data->P);
    c_free(data);
    c_free(settings);

    return 0;
} |]


{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE DerivingVia #-}
#endif
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}

import Control.Monad ((>=>),forM)
import Control.Applicative (liftA2)
import Data.Complex
import Data.Either
import Data.Euclidean hiding (gcd)
import Data.Fixed
import Data.Functor.Const.Compat
import Data.Functor.Identity.Compat
import Data.HashMap.Strict
import Data.HashSet
import Data.Hashable
import Data.Int
import Data.IntMap.Strict
import Data.IntSet
import Data.Map
import Data.Monoid
import Data.Ord (Down(..))
import Data.Orphans ()
import Data.Proxy (Proxy(..))
import Data.Ratio
import Data.Semigroup.Compat
import Data.Semiring
import Data.Semiring.Generic
import Data.Sequence
import Data.Set
import Data.Word
import GHC.Generics
import GHC.Natural
import qualified Data.Foldable as F
import Prelude hiding (Num(..),(^))
import System.IO.Unsafe
import Test.QuickCheck hiding (Fixed)
import Test.QuickCheck.Gen (suchThat)
import Test.QuickCheck.Instances ()
import Test.Tasty (defaultMain, testGroup, TestTree)
import Text.Show.Functions ()
import qualified Data.Map as Map
import qualified Data.Sequence as Seq
import qualified GHC.Num as Num
import qualified Test.QuickCheck.Classes as QCC
import qualified Control.Exception as E

main :: IO ()
main = do
  QCC.lawsCheckMany namedTests
  F.fold pow_prop
  F.fold gcd_prop

forceError :: a -> IO (Either E.ErrorCall a)
forceError = E.try . E.evaluate

pow_prop :: [IO ()]
pow_prop = [quickCheck pow_prop1, quickCheck pow_prop2, quickCheck pow_prop3]

pow_prop1 :: Int -> Negative Int -> Property
pow_prop1 x y = ioProperty $ do
  p <- forceError (x ^ getNegative y)
  pure $ isLeft p

pow_prop2 :: Int -> Positive Int -> Property
pow_prop2 x y = ioProperty $ do
  p <- forceError (x ^ getPositive y)
  pure $ isRight p

pow_prop3 :: Int -> Property
pow_prop3 x = ioProperty $ do
  p <- forceError (x ^ 0)
  case p of
    Left e -> pure False
    Right x -> pure $ x == 1

gcd_prop :: [IO ()]
gcd_prop = [quickCheck gcd_prop1, quickCheck gcd_prop2]

gcd_prop1 :: Integer -> Integer -> Property
gcd_prop1 x y = Num.abs (fst (x `gcdExt` y)) === x `gcd` y

gcd_prop2 :: Integer -> Integer -> Property
gcd_prop2 x y = y /= 0 ==> (x * s) `mod` y === g `mod` y
  where
    (g, s) = x `gcdExt` y

type Laws = QCC.Laws

semiringLaws :: (Arbitrary a, Show a, Eq a, Semiring a) => Proxy a -> [Laws]
semiringLaws p = [QCC.semiringLaws p]

ringLaws :: (Arbitrary a, Show a, Eq a, Ring a) => Proxy a -> [Laws]
ringLaws p = [QCC.semiringLaws p, QCC.ringLaws p]

euclideanLaws :: (Arbitrary a, Show a, Eq a, Euclidean a) => Proxy a -> [Laws]
#if MIN_VERSION_quickcheck_classes(0,6,3)
euclideanLaws p = [QCC.gcdDomainLaws p, QCC.euclideanLaws p]
#else
euclideanLaws _ = []
#endif

namedTests :: [(String, [Laws])]
namedTests =
  [ ("Bool", semiringLaws pBool)
--  , ("Double", ringLaws pDouble) -- needs to be within some epsilon
--  , ("Float", ringLaws pFloat)   -- needs to be within some epsilon
--  , ("Complex", ringLaws pComplex) -- needs to be within some epsilon
  , ("Mod2", semiringLaws pMod2 ++ ringLaws pMod2 ++ euclideanLaws pMod2)
  , ("Int", semiringLaws pInt ++ ringLaws pInt ++ euclideanLaws pInt)
  , ("Word", semiringLaws pWord ++ ringLaws pWord ++ euclideanLaws pWord)
  , ("Natural", semiringLaws pNatural ++ euclideanLaws pNatural)
  , ("Rational", semiringLaws pRational ++ ringLaws pRational ++ euclideanLaws pRational)
  , ("()", semiringLaws pInt8 ++ ringLaws pUnit)
  , ("Maybe", semiringLaws pMaybe)
  , ("PosRatio", semiringLaws pPosRatio)
  , ("IO", semiringLaws pIO ++ ringLaws pIO)
  , ("Fixed", ringLaws pFixed)
  , ("Identity", ringLaws pIdentity)
  , ("Dual", ringLaws pDual)
  , ("(->)", ringLaws pFunction)
  , ("Down", ringLaws pDown)
  , ("Const", ringLaws pConst)
--  , ("IntMap", semiringLaws pIntMap) -- needs newtypes
  , ("Set", semiringLaws pSet)
--  , ("IntSet", semiringLaws pIntSet) -- needs newtypes
  , ("HashSet", semiringLaws pHashSet)
  , ("HashMap", semiringLaws pHashMap)
  , ("Map", semiringLaws pMap)
  , ("Predicate", semiringLaws pPredicate)
  , ("Equivalence", semiringLaws pEquivalence)
  , ("Op", semiringLaws pOp ++ ringLaws pOp)
  , ("Ap", semiringLaws pAp ++ ringLaws pAp)

  , ("IntSet Sum", semiringLaws pIntSetSum)
  , ("IntSet Product", semiringLaws pIntSetProduct)
  , ("IntSet Min", semiringLaws pIntSetMin)
  , ("IntSet Max", semiringLaws pIntSetMax)

  , ("IntMap Sum", semiringLaws pIntMapSum)
  , ("IntMap Product", semiringLaws pIntMapProduct)
  , ("IntMap Min", semiringLaws pIntMapMin)
  , ("IntMap Max", semiringLaws pIntMapMax)

  , ("AB", semiringLaws pAB)
  ]

#if !(MIN_VERSION_base(4,12,0))
newtype Ap f a = Ap { getAp :: f a }
  deriving (Eq, Functor, Applicative, Show)
instance (Semiring a, Applicative f) => Semiring (Ap f a) where
  zero = pure zero
  one = pure one
  plus = liftA2 plus
  times = liftA2 times
  fromNatural = pure . fromNatural
instance (Ring a, Applicative f) => Ring (Ap f a) where
  negate = fmap negate
#endif
deriving instance (Arbitrary (f a)) => Arbitrary (Ap f a)

newtype Predicate a = Predicate (a -> Bool)
  deriving (Eq, Show)
deriving instance Semiring (Predicate a)
deriving instance (Arbitrary a, CoArbitrary a) => Arbitrary (Predicate a)

newtype Equivalence a = Equivalence (a -> a -> Bool)
  deriving (Eq, Show)
deriving instance Semiring a => Semiring (Equivalence a)
deriving instance (Arbitrary a, CoArbitrary a) => Arbitrary (Equivalence a)

newtype Op a b = Op (b -> a)
  deriving (Eq, Show)
deriving instance Semiring a => Semiring (Op a b)
deriving instance Ring a => Ring (Op a b)
deriving instance (Arbitrary a, CoArbitrary a, CoArbitrary b) => Arbitrary (Op a b)

pAp = Proxy :: Proxy (Ap Identity Int)
pPredicate = Proxy :: Proxy (Predicate Int)
pEquivalence = Proxy :: Proxy (Equivalence Int)
pOp = Proxy :: Proxy (Op Int Int)

deriving instance Arbitrary a => Arbitrary (Down a)

-- | Rationals constrained to being positive
newtype PosRatio a = PosRatio { getPosRatio :: Ratio a }
  deriving (Eq, Show, Semiring)

instance (Arbitrary a, Integral a, Ord a) => Arbitrary (PosRatio a) where
  arbitrary =
    let nm = suchThat (arbitrary :: Gen a) (> 0)
        dm = suchThat (arbitrary :: Gen a) (> 0)
    in PosRatio <$> liftA2 (%) nm dm

instance (Arbitrary a, Eq b) => Eq (a -> b) where
  (==) f g =
    let x = generate arbitrary :: IO a
        x' = unsafePerformIO x
    in f x' == g x'

instance Show (Endo a) where
  show _ = "<<<Endo>>>"

deriving instance (Arbitrary a, Eq a) => Eq (Endo a)

-- no
instance Show (IO a) where
  show _ = "<<<IO action>>>"

-- no no
instance Eq a => Eq (IO a) where
  x == y = unsafePerformIO $ liftA2 (==) x y

-- this is fine
instance Arbitrary a => Arbitrary (IO a) where
  arbitrary = fmap pure (arbitrary :: Gen a)

deriving instance Hashable a => Hashable (Sum a)

deriving instance Arbitrary (IntSetOf a)
deriving instance Arbitrary v => Arbitrary (IntMapOf k v)

deriving instance Arbitrary Mod2

pBool = Proxy :: Proxy Bool
pDouble = Proxy :: Proxy Double
pFloat = Proxy :: Proxy Float
pInt = Proxy :: Proxy Int
pInt8 = Proxy :: Proxy Int8
pInt16 = Proxy :: Proxy Int16
pInt32 = Proxy :: Proxy Int32
pInt64 = Proxy :: Proxy Int64
pInteger = Proxy :: Proxy Integer
pNatural = Proxy :: Proxy Natural
pWord = Proxy :: Proxy Word
pWord8 = Proxy :: Proxy Word8
pWord16 = Proxy :: Proxy Word16
pWord32 = Proxy :: Proxy Word32
pWord64 = Proxy :: Proxy Word64
pUnit = Proxy :: Proxy ()
pMaybe = Proxy :: Proxy (Maybe Int)
pPosRatio = Proxy :: Proxy (PosRatio Int)
pIO = Proxy :: Proxy (IO Int)
pComplex = Proxy :: Proxy (Complex Double)
pFixed = Proxy :: Proxy (Fixed E0)
pMin = Proxy :: Proxy (Min Int)
pMax = Proxy :: Proxy (Max Int)
pIdentity = Proxy :: Proxy (Identity Int)
pDual = Proxy :: Proxy (Dual Int)
pEndo = Proxy :: Proxy (Endo (Sum Int))
pSum  = Proxy :: Proxy (Sum Int)
pProduct = Proxy :: Proxy (Product Int)
pDown = Proxy :: Proxy (Down Int)
pSeq = Proxy :: Proxy (Seq Int)
pSet = Proxy :: Proxy (Set (Sum Int))
pHashSet = Proxy :: Proxy (HashSet (Sum Int))
pFunction = Proxy :: Proxy (Int -> Int)
pMap = Proxy :: Proxy (Map (Sum Int) Int)
pHashMap = Proxy :: Proxy (HashMap (Sum Int) Int)
pConst = Proxy :: Proxy (Const Int Int)
pAlt = Proxy :: Proxy (Alt Maybe Int)
pRational = Proxy :: Proxy Rational
pMod2 = Proxy :: Proxy Mod2
pAB = Proxy :: Proxy AB

pIntSetSum = Proxy :: Proxy (IntSetOf (Sum Int))
pIntSetProduct = Proxy :: Proxy (IntSetOf (Product Int))
pIntSetMin = Proxy :: Proxy (IntSetOf (Min Int))
pIntSetMax = Proxy :: Proxy (IntSetOf (Max Int))

pIntMapSum = Proxy :: Proxy (IntMapOf (Sum Int) Int)
pIntMapProduct = Proxy :: Proxy (IntMapOf (Product Int) Int)
pIntMapMin = Proxy :: Proxy (IntMapOf (Min Int) Int)
pIntMapMax = Proxy :: Proxy (IntMapOf (Max Int) Int)

newtype A = A Integer
  deriving (Show, Eq, Semiring, Arbitrary)
newtype B = B Integer
  deriving (Show, Eq, Semiring, Arbitrary)

data AB = AB A B
  deriving (Show, Eq, Generic)
#if MIN_VERSION_base(4,12,0)
  deriving (Semiring) via (GenericSemiring AB)
#else
instance Semiring AB where
  zero = gzero
  one = gone
  plus = gplus
  times = gtimes
  fromNatural = gfromNatural
#endif

instance Arbitrary AB where
  arbitrary = AB <$> arbitrary <*> arbitrary
  shrink = genericShrink


module Grenade.Dynamic.Recurrent.Layers.LSTM where

import           GHC.TypeLits

import           Numeric.LinearAlgebra.Static

import           Grenade.Dynamic.Specification
import           Grenade.Recurrent.Layers.LSTM

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

instance (KnownNat i, KnownNat o) => GNum (LSTM i o) where
  n |* (LSTM w m) = LSTM (n |* w) (n |* m)
  (LSTM w1 m1) |+ (LSTM w2 m2) = LSTM (w1 |+ w2) (m1 |+ m2)
  gFromRational r = LSTM (gFromRational r) (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)
    where
      v0 = konst 0
      w0 = konst 0
      u0 = konst 0

instance (KnownNat i, KnownNat o) => GNum (LSTMWeights i o) where
  n |* (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    LSTMWeights
      (fromRational n * wf)
      (fromRational n * uf)
      (fromRational n * bf)
      (fromRational n * wi)
      (fromRational n * ui)
      (fromRational n * bi)
      (fromRational n * wo)
      (fromRational n * uo)
      (fromRational n * bo)
      (fromRational n * wc)
      (fromRational n * bc)
  (LSTMWeights wf1 uf1 bf1 wi1 ui1 bi1 wo1 uo1 bo1 wc1 bc1) |+ (LSTMWeights wf2 uf2 bf2 wi2 ui2 bi2 wo2 uo2 bo2 wc2 bc2) =
    LSTMWeights (wf1 + wf2) (uf1 + uf2) (bf1 + bf2) (wi1 + wi2) (ui1 + ui2) (bi1 + bi2) (wo1 + wo2) (uo1 + uo2) (bo1 + bo2) (wc1 + wc2) (bc1 + bc2)
  gFromRational r = LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0
    where
      v0 = fromRational r
      w0 = fromRational r
      u0 = fromRational r