-- |
-- Module:      Math.NumberTheory.DirichletCharacters
-- Copyright:   (c) 2018 Bhavik Mehta
-- Licence:     MIT
-- Maintainer:  Bhavik Mehta <bhavikmehta8@gmail.com>
--
-- Implementation and enumeration of Dirichlet characters.
--

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE ViewPatterns               #-}

module Math.NumberTheory.DirichletCharacters
  (
  -- * Roots of unity
    RootOfUnity
  -- ** Conversions
  , toRootOfUnity
  , fromRootOfUnity
  , toComplex
  -- * An absorbing semigroup
  , OrZero, pattern Zero, pattern NonZero
  , orZeroToNum
  -- * Dirichlet characters
  , DirichletCharacter
  -- ** Construction
  , indexToChar
  , indicesToChars
  , characterNumber
  , allChars
  , fromTable
  -- ** Evaluation
  , eval
  , evalGeneral
  , evalAll
  -- ** Special Dirichlet characters
  , principalChar
  , isPrincipal
  , orderChar
  -- ** Real Dirichlet characters
  , RealCharacter
  , isRealCharacter
  , getRealChar
  , toRealFunction
  , jacobiCharacter
  -- ** Primitive characters
  , PrimitiveCharacter
  , isPrimitive
  , getPrimitiveChar
  , induced
  , makePrimitive
  , WithNat(..)
  -- * Debugging
  , validChar
  ) where

#if !MIN_VERSION_base(4,12,0)
import Control.Applicative                                 (liftA2)
#endif
import Data.Bits                                           (Bits(..))
import Data.Complex                                        (Complex(..), cis)
import Data.Foldable                                       (for_)
import Data.Functor.Identity                               (Identity(..))
import Data.List                                           (mapAccumL, foldl', sort, find, unfoldr)
import Data.Maybe                                          (mapMaybe, fromJust, fromMaybe)
#if MIN_VERSION_base(4,12,0)
import Data.Monoid                                         (Ap(..))
#endif
import Data.Proxy                                          (Proxy(..))
import Data.Ratio                                          ((%), numerator, denominator)
import Data.Semigroup                                      (Semigroup(..), Product(..))
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Vector                                         (Vector, (!))
import GHC.TypeNats.Compat                                 (Nat, SomeNat(..), natVal, someNatVal)
import Numeric.Natural                                     (Natural)

import Math.NumberTheory.ArithmeticFunctions               (totient)
import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.Class                      (KnownNat, Mod, getVal)
import Math.NumberTheory.Moduli.Internal                   (isPrimitiveRoot', discreteLogarithmPP)
import Math.NumberTheory.Moduli.Multiplicative             (MultMod(..), isMultElement)
import Math.NumberTheory.Moduli.Singleton                  (Some(..), cyclicGroupFromFactors)
import Math.NumberTheory.Powers.Modular                    (powMod)
import Math.NumberTheory.Primes                            (Prime(..), UniqueFactorisation, factorise, nextPrime)
import Math.NumberTheory.Utils.FromIntegral                (wordToInt)
import Math.NumberTheory.Utils

-- | A Dirichlet character mod \(n\) is a group homomorphism from \((\mathbb{Z}/n\mathbb{Z})^*\)
-- to \(\mathbb{C}^*\), represented abstractly by `DirichletCharacter`. In particular, they take
-- values at roots of unity and can be evaluated using `eval`.
-- A Dirichlet character can be extended to a completely multiplicative function on \(\mathbb{Z}\)
-- by assigning the value 0 for \(a\) sharing a common factor with \(n\), using `evalGeneral`.
--
-- There are finitely many possible Dirichlet characters for a given modulus, in particular there
-- are \(\phi(n)\) characters modulo \(n\), where \(\phi\) refers to Euler's `totient` function.
-- This gives rise to `Enum` and `Bounded` instances.
newtype DirichletCharacter (n :: Nat) = Generated [DirichletFactor]

-- | The group (Z/nZ)^* decomposes to a product (Z/2^k0 Z)^* x (Z/p1^k1 Z)^* x ... x (Z/pi^ki Z)^*
-- where n = 2^k0 p1^k1 ... pi^ki, and the pj are odd primes, k0 possibly 0. Thus, a group
-- homomorphism from (Z/nZ)^* is characterised by group homomorphisms from each of these factor
-- groups. Furthermore, for odd p, we have (Z/p^k Z)^* isomorphic to Z / p^(k-1)*(p-1) Z, an
-- additive group, where an isomorphism is specified by a choice of primitive root.
-- Similarly, for k >= 2, (Z/2^k Z)^* is isomorphic to Z/2Z * (Z / 2^(k-2) Z) (and for k < 2
-- it is trivial).  (See `lambda` for this isomorphism).
-- Thus, to specify a Dirichlet character, it suffices to specify the value of generators
-- of each of these cyclic groups, when primitive roots are given. This data is given by a
-- DirichletFactor.
-- We have the invariant that the factors must be given in strictly increasing order, and the
-- generator is as given by `generator`, and are each non-trivial. These conditions are verified
-- using `validChar`.
data DirichletFactor = OddPrime { _getPrime :: Prime Natural
                                , _getPower :: Word
                                , _getGenerator :: Natural
                                , _getValue :: RootOfUnity
                                }
                     | TwoPower { _getPower2 :: Int -- this ought to be Word, but many applications
                                                    -- needed to use wordToInt, so Int is cleaner
                                                    -- Required to be >= 2
                                , _getFirstValue :: RootOfUnity
                                , _getSecondValue :: RootOfUnity
                                }
                     | Two

instance Eq (DirichletCharacter n) where
  Generated a == Generated b = a == b

instance Eq DirichletFactor where
  TwoPower _ x1 x2 == TwoPower _ y1 y2 = x1 == y1 && x2 == y2
  OddPrime _ _ _ x == OddPrime _ _ _ y = x == y
  Two              == Two              = True
  _ == _ = False

-- | A representation of <https://en.wikipedia.org/wiki/Root_of_unity roots of unity>: complex
-- numbers \(z\) for which there is \(n\) such that \(z^n=1\).
newtype RootOfUnity =
  RootOfUnity { -- | Every root of unity can be expressed as \(e^{2 \pi i q}\) for some
                -- rational \(q\) satisfying \(0 \leq q < 1\), this function extracts it.
                fromRootOfUnity :: Rational }
  deriving (Eq)

instance Show RootOfUnity where
  show (RootOfUnity q)
    | n == 0    = "1"
    | d == 1    = "-1"
    | n == 1    = "e^(πi/" ++ show d ++ ")"
    | otherwise = "e^(" ++ show n ++ "πi/" ++ show d ++ ")"
    where n = numerator (2*q)
          d = denominator (2*q)

-- | Given a rational \(q\), produce the root of unity \(e^{2 \pi i q}\).
toRootOfUnity :: Rational -> RootOfUnity
toRootOfUnity q = RootOfUnity ((n `rem` d) % d)
  where n = numerator q
        d = denominator q
        -- effectively q `mod` 1
  -- This smart constructor ensures that the rational is always in the range 0 <= q < 1.

-- | This Semigroup is in fact a group, so @'stimes'@ can be called with a negative first argument.
instance Semigroup RootOfUnity where
  RootOfUnity q1 <> RootOfUnity q2 = toRootOfUnity (q1 + q2)
  stimes k (RootOfUnity q) = toRootOfUnity (q * fromIntegral k)

instance Monoid RootOfUnity where
  mappend = (<>)
  mempty = RootOfUnity 0

-- | Convert a root of unity into an inexact complex number. Due to floating point inaccuracies,
-- it is recommended to avoid use of this until the end of a calculation. Alternatively, with
-- the [cyclotomic](http://hackage.haskell.org/package/cyclotomic-0.5.1) package, one can use
-- @[polarRat](https://hackage.haskell.org/package/cyclotomic-0.5.1/docs/Data-Complex-Cyclotomic.html#v:polarRat)
-- 1 . @'fromRootOfUnity' to convert to a cyclotomic number.
toComplex :: Floating a => RootOfUnity -> Complex a
toComplex (RootOfUnity t)
  | t == 1/2 = (-1) :+ 0
  | t == 1/4 = 0 :+ 1
  | t == 3/4 = 0 :+ (-1)
  | otherwise = cis . (2*pi*) . fromRational $ t

-- | For primes, define the canonical primitive root as the smallest such. For prime powers \(p^k\),
-- either the smallest primitive root \(g\) mod \(p\) works, or \(g+p\) works.
generator :: (Integral a, UniqueFactorisation a) => Prime a -> Word -> a
generator p k
  | k == 1 = modP
  | otherwise = if powMod modP (p'-1) (p'*p') == 1 then modP + p' else modP
  where p' = unPrime p
        modP = case cyclicGroupFromFactors [(p,k)] of
                 Just (Some cg) -> head $ filter (isPrimitiveRoot' cg) [2..p'-1]
                 _ -> error "illegal"

-- | Implement the function \(\lambda\) from page 5 of
-- https://www2.eecs.berkeley.edu/Pubs/TechRpts/1984/CSD-84-186.pdf
lambda :: Integer -> Int -> Integer
lambda x e = ((powMod x (2*modulus) largeMod - 1) `shiftR` (e+1)) .&. (modulus - 1)
  where modulus = bit (e-2)
        largeMod = bit (2*e - 1)

-- | For elements of the multiplicative group \((\mathbb{Z}/n\mathbb{Z})^*\), a Dirichlet
-- character evaluates to a root of unity.
eval :: DirichletCharacter n -> MultMod n -> RootOfUnity
eval (Generated ds) m = foldMap (evalFactor m') ds
  where m' = getVal $ multElement m

-- | Evaluate each factor of the Dirichlet character.
evalFactor :: Integer -> DirichletFactor -> RootOfUnity
evalFactor m =
  \case
    OddPrime (toInteger . unPrime -> p) k (toInteger -> a) b ->
      discreteLogarithmPP p k a (m `rem` p^k) `stimes` b
    TwoPower k s b -> (if testBit m 1 then s else mempty)
                   <> lambda (thingy k m) k `stimes` b
    Two -> mempty

thingy :: (Bits p, Num p) => Int -> p -> p
thingy k m = if testBit m 1
                then bit k - m'
                else m'
  where m' = m .&. (bit k - 1)

-- | A character can evaluate to a root of unity or zero: represented by @Nothing@.
evalGeneral :: KnownNat n => DirichletCharacter n -> Mod n -> OrZero RootOfUnity
evalGeneral chi t = case isMultElement t of
                      Nothing -> Zero
                      Just x -> NonZero $ eval chi x

-- | Give the principal character for this modulus: a principal character mod \(n\) is 1 for
-- \(a\) coprime to \(n\), and 0 otherwise.
principalChar :: KnownNat n => DirichletCharacter n
principalChar = minBound

mulChars :: DirichletCharacter n -> DirichletCharacter n -> DirichletCharacter n
mulChars (Generated x) (Generated y) = Generated (zipWith combine x y)
  where combine :: DirichletFactor -> DirichletFactor -> DirichletFactor
        combine Two Two = Two
        combine (OddPrime p k g n) (OddPrime _ _ _ m) =
          OddPrime p k g (n <> m)
        combine (TwoPower k a n) (TwoPower _ b m) =
          TwoPower k (a <> b) (n <> m)
        combine _ _ = error "internal error: malformed DirichletCharacter"

-- | This Semigroup is in fact a group, so @stimes@ can be called with a negative first argument.
instance Semigroup (DirichletCharacter n) where
  (<>) = mulChars
  stimes = stimesChar

instance KnownNat n => Monoid (DirichletCharacter n) where
  mempty = principalChar
  mappend = (<>)

stimesChar :: Integral a => a -> DirichletCharacter n -> DirichletCharacter n
stimesChar s (Generated xs) = Generated (map mult xs)
  where mult :: DirichletFactor -> DirichletFactor
        mult (OddPrime p k g n) = OddPrime p k g (s `stimes` n)
        mult (TwoPower k a b) = TwoPower k (s `stimes` a) (s `stimes` b)
        mult Two = Two

-- | We define `succ` and `pred` with more efficient implementations than
-- @`toEnum` . (+1) . `fromEnum`@.
instance KnownNat n => Enum (DirichletCharacter n) where
  toEnum = indexToChar . fromIntegral
  fromEnum = fromIntegral . characterNumber
  succ x = makeChar x (characterNumber x + 1)
  pred x = makeChar x (characterNumber x - 1)

  enumFromTo x y       = bulkMakeChars x [fromEnum x..fromEnum y]
  enumFrom x           = bulkMakeChars x [fromEnum x..]
  enumFromThenTo x y z = bulkMakeChars x [fromEnum x, fromEnum y..fromEnum z]
  enumFromThen x y     = bulkMakeChars x [fromEnum x, fromEnum y..]

instance KnownNat n => Bounded (DirichletCharacter n) where
  minBound = indexToChar 0
  maxBound = indexToChar (totient n - 1)
    where n = natVal (Proxy :: Proxy n)

-- | We have a (non-canonical) enumeration of dirichlet characters.
characterNumber :: DirichletCharacter n -> Integer
characterNumber (Generated y) = foldl' go 0 y
  where go x (OddPrime p k _ a) = x * m + numerator (fromRootOfUnity a * fromIntegral m)
          where p' = fromIntegral (unPrime p)
                m = p'^(k-1)*(p'-1)
        go x (TwoPower k a b)   = x' * 2 + numerator (fromRootOfUnity a * 2)
          where m = bit (k-2) :: Integer
                x' = x `shiftL` (k-2) + numerator (fromRootOfUnity b * fromIntegral m)
        go x Two = x

-- | Give the dirichlet character from its number.
-- Inverse of `characterNumber`.
indexToChar :: forall n. KnownNat n => Natural -> DirichletCharacter n
indexToChar = runIdentity . indicesToChars . Identity

-- | Give a collection of dirichlet characters from their numbers. This may be more efficient than
-- `indexToChar` for multiple characters, as it prevents some internal recalculations.
indicesToChars :: forall n f. (KnownNat n, Functor f) => f Natural -> f (DirichletCharacter n)
indicesToChars = fmap (Generated . unroll t . (`mod` m))
  where n = natVal (Proxy :: Proxy n)
        (Product m, t) = mkTemplate n

-- | List all characters for the modulus. This is preferred to using @[minBound..maxBound]@.
allChars :: forall n. KnownNat n => [DirichletCharacter n]
allChars = indicesToChars [0..m-1]
  where m = totient $ natVal (Proxy :: Proxy n)

-- | The same as `indexToChar`, but if we're given a character we can create others more efficiently.
makeChar :: Integral a => DirichletCharacter n -> a -> DirichletCharacter n
makeChar x = runIdentity . bulkMakeChars x . Identity

-- | Use one character to make many more: better than indicesToChars since it avoids recalculating
-- some primitive roots
bulkMakeChars :: (Integral a, Functor f) => DirichletCharacter n -> f a -> f (DirichletCharacter n)
bulkMakeChars x = fmap (Generated . unroll t . (`mod` m) . fromIntegral)
  where (Product m, t) = templateFromCharacter x

-- We assign each natural a unique Template, which can be decorated (eg in `unroll`) to
-- form a DirichletCharacter. A Template effectively holds the information carried around
-- in a DirichletFactor which depends only on the modulus of the character.
data Template = OddTemplate { _getPrime'     :: Prime Natural
                            , _getPower'     :: Word
                            , _getGenerator' :: !Natural
                            , _getModulus'   :: !Natural
                            }
              | TwoPTemplate { _getPower2'    :: Int
                             , _getModulus'   :: !Natural
                             } -- the modulus is derivable from the other values, but calculation
                               -- may be expensive, so we pre-calculate it
                               -- morally getModulus should be a prefactored but seems to be
                               -- pointless here
              | TwoTemplate

templateFromCharacter :: DirichletCharacter n -> (Product Natural, [Template])
templateFromCharacter (Generated t) = traverse go t
  where go (OddPrime p k g _) = (Product m, OddTemplate p k g m)
          where p' = unPrime p
                m = p'^(k-1)*(p'-1)
        go (TwoPower k _ _) = (Product (2*m), TwoPTemplate k m)
          where m = bit (k-2)
        go Two = (Product 1, TwoTemplate)

mkTemplate :: Natural -> (Product Natural, [Template])
mkTemplate = go . sort . factorise
  where go :: [(Prime Natural, Word)] -> (Product Natural, [Template])
        go ((unPrime -> 2, 1): xs) = (Product 1, [TwoTemplate]) <> traverse odds xs
        go ((unPrime -> 2, wordToInt -> k): xs) = (Product (2*m), [TwoPTemplate k m]) <> traverse odds xs
          where m = bit (k-2)
        go xs = traverse odds xs
        odds :: (Prime Natural, Word) -> (Product Natural, Template)
        odds (p, k) = (Product m, OddTemplate p k (generator p k) m)
          where p' = unPrime p
                m = p'^(k-1)*(p'-1)

-- the validity of the producted dirichletfactor list here requires the template to be valid
unroll :: [Template] -> Natural -> [DirichletFactor]
unroll t m = snd (mapAccumL func m t)
  where func :: Natural -> Template -> (Natural, DirichletFactor)
        func a (OddTemplate p k g n) = (a1, OddPrime p k g (toRootOfUnity $ (toInteger a2) % (toInteger n)))
          where (a1,a2) = quotRem a n
        func a (TwoPTemplate k n) = (b1, TwoPower k (toRootOfUnity $ (toInteger a2) % 2) (toRootOfUnity $ (toInteger b2) % (toInteger n)))
          where (a1,a2) = quotRem a 2
                (b1,b2) = quotRem a1 n
        func a TwoTemplate = (a, Two)

-- | Test if a given Dirichlet character is prinicpal for its modulus: a principal character mod
-- \(n\) is 1 for \(a\) coprime to \(n\), and 0 otherwise.
isPrincipal :: DirichletCharacter n -> Bool
isPrincipal chi = characterNumber chi == 0

-- | Induce a Dirichlet character to a higher modulus. If \(d \mid n\), then \(a \bmod{n}\) can be
-- reduced to \(a \bmod{d}\). Thus, the multiplicative function on \(\mathbb{Z}/d\mathbb{Z}\)
-- induces a multiplicative function on \(\mathbb{Z}/n\mathbb{Z}\).
--
-- >>> :set -XTypeApplications
-- >>> chi = indexToChar 5 :: DirichletCharacter 45
-- >>> chi2 = induced @135 chi
-- >>> :t chi2
-- Maybe (DirichletCharacter 135)
induced :: forall n d. (KnownNat d, KnownNat n) => DirichletCharacter d -> Maybe (DirichletCharacter n)
induced (Generated start) = if n `rem` d == 0
                            then Just (Generated (combine (snd $ mkTemplate n) start))
                            else Nothing
  where n = natVal (Proxy :: Proxy n)
        d = natVal (Proxy :: Proxy d)
        combine :: [Template] -> [DirichletFactor] -> [DirichletFactor]
        combine [] _ = []
        combine ts [] = map newFactor ts
        combine (t:xs) (y:ys) = case (t,y) of
                                  (TwoTemplate, Two) -> Two: combine xs ys
                                  (TwoTemplate, _) -> Two: combine xs (y:ys)
                                  (TwoPTemplate k _, Two) -> TwoPower k mempty mempty: combine xs ys
                                  (TwoPTemplate k _, TwoPower _ a b) -> TwoPower k a b: combine xs ys
                                  (TwoPTemplate k _, _) -> TwoPower k mempty mempty: combine xs (y:ys)
                                  (OddTemplate p k _ _, OddPrime q _ g a) | p == q -> OddPrime p k g a: combine xs ys
                                  (OddTemplate p k g _, OddPrime q _ _ _) | p < q -> OddPrime p k g mempty: combine xs (y:ys)
                                  _ -> error "internal error in induced: please report this as a bug"
        newFactor :: Template -> DirichletFactor
        newFactor TwoTemplate = Two
        newFactor (TwoPTemplate k _) = TwoPower k mempty mempty
        newFactor (OddTemplate p k g _) = OddPrime p k g mempty
        -- rest (p,k) = OddPrime p k (generator p k) mempty

-- | The <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol> gives a real Dirichlet
-- character for odd moduli.
jacobiCharacter :: forall n. KnownNat n => Maybe (RealCharacter n)
jacobiCharacter = if odd n
                     then Just $ RealChar $ Generated $ map go $ snd $ mkTemplate n
                     else Nothing
  where n = natVal (Proxy :: Proxy n)
        go :: Template -> DirichletFactor
        go (OddTemplate p k g _) = OddPrime p k g $ toRootOfUnity ((toInteger k) % 2)
          -- jacobi symbol of a primitive root mod p over p is always -1
        go _ = error "internal error in jacobiCharacter: please report this as a bug"
          -- every factor of n should be odd

-- | A Dirichlet character is real if it is real-valued.
newtype RealCharacter n = RealChar { -- | Extract the character itself from a `RealCharacter`.
                                     getRealChar :: DirichletCharacter n
                                   }
                                   deriving Eq

-- | Test if a given `DirichletCharacter` is real, and if so give a `RealCharacter`.
isRealCharacter :: DirichletCharacter n -> Maybe (RealCharacter n)
isRealCharacter t@(Generated xs) = if all real xs then Just (RealChar t) else Nothing
  where real :: DirichletFactor -> Bool
        real (OddPrime _ _ _ a) = a <> a == mempty
        real (TwoPower _ _ b) = b <> b == mempty
        real Two = True

-- TODO: it should be possible to calculate this without eval/evalGeneral
-- and thus avoid using discrete log calculations: consider the order of m
-- inside each of the factor groups?
-- | Evaluate a real Dirichlet character, which can only take values \(-1,0,1\).
toRealFunction :: KnownNat n => RealCharacter n -> Mod n -> Int
toRealFunction (RealChar chi) m = case evalGeneral chi m of
                                    Zero -> 0
                                    NonZero t | t == mempty -> 1
                                    NonZero t | t == RootOfUnity (1 % 2) -> -1
                                    _ -> error "internal error in toRealFunction: please report this as a bug"
                                      -- A real character should not be able to evaluate to
                                      -- anything other than {-1,0,1}, so should not reach this branch

-- | Test if the internal DirichletCharacter structure is valid.
validChar :: forall n. KnownNat n => DirichletCharacter n -> Bool
validChar (Generated xs) = correctDecomposition && all correctPrimitiveRoot xs && all validValued xs
  where correctDecomposition = sort (factorise n) == map getPP xs
        getPP (TwoPower k _ _) = (two, fromIntegral k)
        getPP (OddPrime p k _ _) = (p, k)
        getPP Two = (two,1)
        correctPrimitiveRoot (OddPrime p k g _) = g == generator p k
        correctPrimitiveRoot _ = True
        validValued (TwoPower k a b) = a <> a == mempty && (bit (k-2) :: Integer) `stimes` b == mempty
        validValued (OddPrime (unPrime -> p) k _ a) = (p^(k-1)*(p-1)) `stimes` a == mempty
        validValued Two = True
        n = natVal (Proxy :: Proxy n)
        two = nextPrime 2

-- | Get the order of the Dirichlet Character.
orderChar :: DirichletCharacter n -> Integer
orderChar (Generated xs) = foldl' lcm 1 $ map orderFactor xs
  where orderFactor (TwoPower _ (RootOfUnity a) (RootOfUnity b)) = denominator a `lcm` denominator b
        orderFactor (OddPrime _ _ _ (RootOfUnity a)) = denominator a
        orderFactor Two = 1

-- | Test if a Dirichlet character is <https://en.wikipedia.org/wiki/Dirichlet_character#Primitive_characters_and_conductor primitive>.
isPrimitive :: DirichletCharacter n -> Maybe (PrimitiveCharacter n)
isPrimitive t@(Generated xs) = if all primitive xs then Just (PrimitiveCharacter t) else Nothing
  where primitive :: DirichletFactor -> Bool
        primitive Two = False
        -- for odd p, we're testing if phi(p^(k-1)) `stimes` a is 1, since this means the
        -- character can come from some the smaller modulus p^(k-1)
        primitive (OddPrime _ 1 _ a) = a /= mempty
        primitive (OddPrime (unPrime -> p) k _ a) = (p^(k-2)*(p-1)) `stimes` a /= mempty
        primitive (TwoPower 2 a _) = a /= mempty
        primitive (TwoPower k _ b) = (bit (k-3) :: Integer) `stimes` b /= mempty

-- | A Dirichlet character is primitive if cannot be 'induced' from any character with
-- strictly smaller modulus.
newtype PrimitiveCharacter n = PrimitiveCharacter { -- | Extract the character itself from a `PrimitiveCharacter`.
                                                    getPrimitiveChar :: DirichletCharacter n
                                                    }
                                                    deriving Eq

data WithNat (a :: Nat -> *) where
  WithNat :: KnownNat m => a m -> WithNat a

-- | This function also provides access to the new modulus on type level, with a KnownNat instance
makePrimitive :: DirichletCharacter n -> WithNat PrimitiveCharacter
makePrimitive (Generated xs) =
  case someNatVal (product mods) of
    SomeNat (Proxy :: Proxy m) -> WithNat (PrimitiveCharacter (Generated ys) :: PrimitiveCharacter m)
  where (mods,ys) = unzip (mapMaybe prim xs)
        prim :: DirichletFactor -> Maybe (Natural, DirichletFactor)
        prim Two = Nothing
        prim (OddPrime p' k g a) = case find works options of
                                     Nothing -> error "invalid character"
                                     Just (0,_) -> Nothing
                                     Just (i,_) -> Just (p^i, OddPrime p' i g a)
          where options = (0,1): [(i,p^(i-1)*(p-1)) | i <- [1..k]]
                works (_,phi) = phi `stimes` a == mempty
                p = unPrime p'
        prim (TwoPower k a b) = case find worksb options of
                                  Nothing -> error "invalid character"
                                  Just (2,_) | a == mempty -> Nothing
                                  Just (i,_) -> Just (bit i :: Natural, TwoPower i a b)
          where options = [(i, bit (i-2) :: Natural) | i <- [2..k]]
                worksb (_,phi) = phi `stimes` b == mempty

#if !MIN_VERSION_base(4,12,0)
newtype Ap f a = Ap { getAp :: f a }
  deriving (Eq, Functor, Applicative, Monad)

instance (Applicative f, Semigroup a) => Semigroup (Ap f a) where
  (<>) = liftA2 (<>)

instance (Applicative f, Semigroup a, Monoid a) => Monoid (Ap f a) where
  mempty = pure mempty
  mappend = (<>)
#endif

-- | Similar to Maybe, but with different Semigroup and Monoid instances.
type OrZero a = Ap Maybe a
pattern Zero :: OrZero a
pattern Zero = Ap Nothing

pattern NonZero :: a -> OrZero a
pattern NonZero x = Ap (Just x)
{-# COMPLETE Zero, NonZero #-}

-- | Interpret an `OrZero` as a number, taking the `Zero` case to be 0.
orZeroToNum :: Num a => (b -> a) -> OrZero b -> a
orZeroToNum _ Zero = 0
orZeroToNum f (NonZero x) = f x

-- | In general, evaluating a DirichletCharacter at a point involves solving the discrete logarithm
-- problem, which can be hard: the implementations here are around O(sqrt n).
-- However, evaluating a dirichlet character at every point amounts to solving the discrete
-- logarithm problem at every point also, which can be done together in O(n) time, better than
-- using a complex algorithm at each point separately. Thus, if a large number of evaluations
-- of a dirichlet character are required, `evalAll` will be better than `evalGeneral`, since
-- computations can be shared.
evalAll :: forall n. KnownNat n => DirichletCharacter n -> Vector (OrZero RootOfUnity)
evalAll (Generated xs) = V.generate (fromIntegral n) func
  where n = natVal (Proxy :: Proxy n)
        vectors = map mkVector xs
        func :: Int -> OrZero RootOfUnity
        func m = foldMap go vectors
          where go :: (Int, Vector (OrZero RootOfUnity)) -> OrZero RootOfUnity
                go (modulus,v) = v ! (m `mod` modulus)
        mkVector :: DirichletFactor -> (Int, Vector (OrZero RootOfUnity))
        mkVector Two = (2, V.fromList [Zero, mempty])
        mkVector (OddPrime p k (fromIntegral -> g) a) = (modulus, w)
          where
            p' = unPrime p
            modulus = fromIntegral (p'^k) :: Int
            w = V.create $ do
              v <- MV.replicate modulus Zero
              -- TODO: we're in the ST monad here anyway, could be better to use STRefs to manage
              -- this loop, the current implementation probably doesn't fuse well
              let powers = iterateMaybe go (1,mempty)
                  go (m,x) = if m' > 1
                                then Just (m', x<>a)
                                else Nothing
                    where m' = m*g `mod` modulus
              for_ powers $ \(m,x) -> MV.unsafeWrite v m (NonZero x)
              -- don't bother with bounds check since m was reduced mod p^k
              return v
        -- for powers of two we use lambda directly instead, since the generators of the cyclic
        -- groups aren't obvious; it's possible to get them though:
        -- 5^(lambda(5)^{-1} mod 2^(p-2)) mod 2^p
        mkVector (TwoPower k a b) = (modulus, w)
          where
            modulus = bit k
            w = V.generate modulus f
            f m
              | even m = Zero
              | otherwise = NonZero ((if testBit m 1 then a else mempty) <> lambda (toInteger m'') k `stimes` b)
              where m'' = thingy k m

-- somewhere between unfoldr and iterate
iterateMaybe :: (a -> Maybe a) -> a -> [a]
iterateMaybe f x = unfoldr (fmap (\t -> (t, f t))) (Just x)

-- | Attempt to construct a character from its table of values.
-- An inverse to `evalAll`, defined only on its image.
fromTable :: forall n. KnownNat n => Vector (OrZero RootOfUnity) -> Maybe (DirichletCharacter n)
fromTable v = if length v == fromIntegral n
                 then Generated <$> traverse makeFactor tmpl >>= check
                 else Nothing
  where n = natVal (Proxy :: Proxy n)
        n' = fromIntegral n :: Integer
        tmpl = snd (mkTemplate n)
        check :: DirichletCharacter n -> Maybe (DirichletCharacter n)
        check chi = if evalAll chi == v then Just chi else Nothing
        makeFactor :: Template -> Maybe DirichletFactor
        makeFactor TwoTemplate = Just Two
        makeFactor (TwoPTemplate k _) = TwoPower k <$> getValue (-1,bit k) <*> getValue (exp4 k, bit k)
        makeFactor (OddTemplate p k g _) = OddPrime p k g <$> getValue (toInteger g, toInteger (unPrime p)^k)
        getValue :: (Integer,Integer) -> Maybe RootOfUnity
        getValue (g,m) = getAp (v ! fromInteger (fromJust (chineseCoprime (g,m) (1,n' `quot` m)) `mod` n'))

exp4terms :: [Rational]
exp4terms = [4^k % product [1..k] | k <- [0..]]

-- For reasons that aren't clear to me, `exp4` gives the inverse of 1 under lambda, so it gives the generator
-- This is the same as https://oeis.org/A320814
-- In particular, lambda (exp4 n) n == 1 (for n >= 3)
-- I've verified this for 3 <= n <= 2000, so the reasoning in fromTable should be accurate for moduli below 2^2000
exp4 :: Int -> Integer
exp4 n = (`mod` bit n) $ sum $ map (`mod` bit n) $ map (\q -> numerator q * fromMaybe (error "error in exp4") (recipMod (denominator q) (bit n))) $ take n $ exp4terms


import Debug.Trace
-- import Control.Monad.State.Lazy

import System.Environment (getArgs)
import System.IO (readFile)
import qualified Data.Set as S
import qualified Data.Map.Strict as M
import Data.Set (Set)
import Data.Map.Strict (Map)
import Data.Tuple (swap)
import Data.Complex (Complex((:+)), polar)

import Data.List (sort, nub)

readMap :: [[Char]] -> [(Int, Int)]
readMap rows = [(x, y) | (y, row) <- zip [0 ..] rows
                       , (x, c) <- zip [0 ..] row
                       , c == '#']

analyzeMap :: [(Int, Int)] -> (Int, (Int, Int))
analyzeMap points = maximum [(length $ nub angles, p)
                            | p <- points
                            , let angles = map (fst . polarDirection p) points]

polarDirection :: (Int, Int) -> (Int, Int) -> (Float, Float)
polarDirection origin@(x0,y0) pt@(x1,y1) = -- (angle_from_vertical, distance)
  swap -- angle first
  $ positiveAngle -- polar's angle is from -pi to pi, convert to 0 to 2pi so that we can sort by angle
  $ polar -- convert to (magnitude, phase)
  $ fromIntegral (-1 * y) :+ fromIntegral x -- make complex, normally
           -- the positive x axis is phase 0, but here we want to go
           -- clockwise from vertical, so swap x and y and change y's
           -- sign to correct for the fact that the coordinates
           -- increase in the down direction.
  where (x, y) = (x1-x0,y1-y0) -- point from perspective of origin
        positiveAngle (m,p) | p < 0 = (m, 2 * pi + p)
                            | otherwise = (m, p)

main = do
  [mapFile] <- getArgs
  m <- readFile mapFile
  let points = readMap $ lines m
  let (reachable, station) = analyzeMap points
  putStrLn $ "Part 1: " ++ show (reachable, station)
  let polarPoints = -- traceShowId $
        sort -- [(round, angle, xy)] in order
        $ foldl (\l ((angle, distance), xy) ->
                   (if null l then 0
                    else let (round0, angle0, pt0) = head l
                         in if angle == angle0 then round0 + 1
                            else 0,
                     angle, xy):l) ([]::[(Int, Float, Int)])
        -- $ traceShowId
        $ sort -- [((angle, distance), xy)] in order
        $ map (\(x,y) -> (polarDirection station (x,y),100*x+y))
        $ filter (/= station) points

  putStrLn $ "Part 2: " ++ show (polarPoints!!199)
  --putStrLn $ "Part 2: " ++ (show $ head $ drop 199 polarPoints)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE CPP #-}

module Main where

import qualified Data.Vector.Unboxed         as U
import           Data.Vector.Unboxed         (Vector)
import qualified Data.Vector as V

-- DLA
import qualified Statistics.Matrix as M
import qualified Statistics.Matrix.Fast as MF
import qualified Statistics.Matrix.Fast.Algorithms as A

-- hmatrix
import qualified Numeric.LinearAlgebra as H

-- numhask
import qualified NumHask.Array as NH
import qualified NumHask.Prelude as NP

-- data.matrix from matrix
import qualified Data.Matrix as DMX

-- massiv
import qualified Data.Massiv.Array as MA

import qualified System.Random.MWC as Mwc

import qualified Criterion.Main as C

#define N 10
#define N2 100

n :: Int
n = N

vectorGen :: IO (Vector Double)
vectorGen =  do
    gen <- Mwc.create
    Mwc.uniformVector gen (n*n)

matrixDLA :: IO M.Matrix
matrixDLA = do
    vec <- vectorGen
    return $ M.Matrix n n vec

matrixH :: IO (H.Matrix H.R)
matrixH = do
    vec <- vectorGen
    return $ (n H.>< n) $ U.toList $ vec

identH :: Int -> H.Matrix Double
identH = H.ident

identDMX :: Int -> DMX.Matrix Double
identDMX = DMX.identity

elemZero :: Double -> Double
elemZero = const 0

elemSqr :: Double -> Double
elemSqr x = x*x

mapH :: (Double -> Double) -> H.Matrix Double -> H.Matrix Double
mapH = H.cmap

main :: IO ()
main = do

    vDLA' <- vectorGen
    uDLA' <- vectorGen

    let

    --
      subDLA' = U.take n vDLA'
      aDLA' = M.Matrix n n vDLA'
      bDLA' = M.Matrix n n uDLA'

    --
      vList = U.toList vDLA'
      uList = U.toList uDLA'

    --
      aH' = (n H.>< n) vList
      bH' = (n H.>< n) uList

      subH' = H.fromList . take n $ vList
      vH' = H.fromList vList

    --
      aNH' = NP.fromList vList :: NH.Array V.Vector '[N, N] Double
      bNH' = NP.fromList uList :: NH.Array V.Vector '[N, N] Double

      vNH' = NP.fromList vList :: NH.Array V.Vector '[N2] Double

    --
      aDMX' = DMX.fromList n n vList
      bDMX' = DMX.fromList n n uList

    --
      vMA' = MA.fromList MA.Seq vList :: MA.Array MA.P MA.Ix1 Double
      aMA' = MA.resize' (MA.Sz (n MA.:. n)) vMA' :: MA.Array MA.P MA.Ix2 Double
      bMA' = MA.resize' (MA.Sz (n MA.:. n)) $ MA.fromList MA.Seq uList :: MA.Array MA.P MA.Ix2 Double

    C.defaultMain [
        C.env (pure (aDLA', bDLA', subDLA', vDLA')) $ \ ~(aDLA, bDLA, subDLA, vDLA) ->
            C.bgroup "DLA" [
                         C.bench "multiplication" $ C.nf (MF.multiply aDLA) bDLA,
                         C.bench "repeated multiplication" $ C.nf (U.sum . (flip M.row) 1 . MF.multiply bDLA . MF.multiply aDLA . MF.multiply aDLA ) bDLA,
                         C.bench "multiplicationV" $ C.nf (MF.multiplyV aDLA) subDLA,
                         C.bench "qr factorization" $ C.nf A.qr aDLA,
                         C.bench "transpose" $ C.nf MF.transpose aDLA,
                         C.bench "norm" $ C.nf MF.norm vDLA,
                         C.bench "row" $ C.nf (M.row  aDLA) 0,
                         C.bench "column" $ C.nf (M.column  aDLA) 0,
                         C.bench "identity" $ C.nf M.ident n,
                         C.bench "diag" $ C.nf M.diag subDLA,
                         C.bench "map const 0" $ C.nf (M.map elemZero) aDLA,
                         C.bench "map sqr" $ C.nf (M.map elemSqr) aDLA
                       ],

        C.env (pure (aH', bH', subH', vH')) $ \ ~(aH, bH, subH, vH) ->
            C.bgroup "Hmatrix" [
                             C.bench "multiplication" $ C.nf ((<>) aH) bH,
                             C.bench "repeated multiplication" $ C.nf ( H.sumElements . flip (H.?) [1] . (<>) bH . (<>) aH . (<>) aH) bH,
                             C.bench "multiplicationV" $ C.nf ((H.#>) aH) subH,
                             C.bench "qr factorization" $ C.nf H.qr aH,
                             C.bench "transpose" $ C.nf H.tr aH,
                             C.bench "norm" $ C.nf H.norm_2 vH,
                             C.bench "row" $ C.nf ((H.?) aH) [0],
                             C.bench "column" $ C.nf ((H.¿) aH) [0],
                             C.bench "identity" $ C.nf identH n,
                             C.bench "diag" $ C.nf H.diag subH,
                             C.bench "map const 0" $ C.nf (mapH elemZero) aH,
                             C.bench "map sqr" $ C.nf (mapH elemSqr) aH
                           ],

        C.env (pure (aNH', bNH', vNH')) $ \ ~(aNH, bNH, vNH) ->
            C.bgroup "NumHask" [
                             C.bench "multiplication" $ C.nf (NH.mmult aNH) bNH,
                             C.bench "repeated multiplication" $ C.nf ( (\(NH.Array a) -> V.sum a) . NH.row (NP.Proxy :: NP.Proxy 1) . NH.mmult bNH . NH.mmult aNH . NH.mmult aNH ) bNH,
                             C.bench "transpose" $ C.nf NH.transpose aNH,
                             C.bench "norm" $ C.nf (sqrt . (NP.<.> vNH)) vNH,
                             C.bench "row" $ C.nf (NH.row (NP.Proxy :: NP.Proxy 0)) aNH,
                             C.bench "column" $ C.nf (NH.col (NP.Proxy :: NP.Proxy 0)) aNH
                           ],

        C.env (pure (aMA', bMA', vMA')) $ \ ~(aMA, bMA, vMA) ->
            C.bgroup "Massiv" [
                                 C.bench "multiplication" $ C.nf ((MA.|*|) aMA) bMA,
                                 C.bench "multiplication (Par)" $ C.nf ((MA.|*|) (MA.setComp MA.Par aMA)) bMA,
                                 C.bench "repeated multiplication" $ C.nf ( MA.foldlS (+) 0 .  flip (MA.!>) 1 . (MA.|*|) bMA . (MA.|*|) aMA . (MA.|*|) aMA) bMA,
                                 C.bench "repeated multiplication (Par)" $ C.nf ( MA.foldlS (+) 0 .  flip (MA.!>) 1 . (MA.|*|) bMA . (MA.|*|) aMA . (MA.|*|) (MA.setComp MA.Par aMA)) bMA,
                                 C.bench "norm" $ C.nf (sqrt . MA.foldlS (+) 0 . (MA.zipWith (*) vMA)) vMA,
                                 C.bench "transpose" $ C.nf (MA.computeAs MA.P . MA.transpose) aMA,
                                 C.bench "row" $ C.nf (MA.computeAs MA.P . (MA.!>) aMA) 0,
                                 C.bench "column" $ C.nf (MA.computeAs MA.P . (MA.<!) aMA) 0
                                 ],

        C.env (pure (aDMX', bDMX')) $ \ ~(aDMX, bDMX) ->
                C.bgroup "Matrix" [
                                 C.bench "multiplication" $ C.nf (DMX.multStrassenMixed aDMX) bDMX,
                                 C.bench "transpose" $ C.nf DMX.transpose aDMX,
                                 C.bench "row" $ C.nf (DMX.getRow 1) aDMX,
                                 C.bench "column" $ C.nf (DMX.getCol 1) aDMX,
                                 C.bench "identity" $ C.nf identDMX n
                                 ]
                  ]


{-# LANGUAGE TypeFamilies #-}

import Test.HUnit
import Test.Framework
import Test.Framework.Providers.HUnit

import Numeric.LinearAlgebra.HMatrix
import Numeric.Minimization.QuadProgPP

import Numeric.AD

import Data.SQP

qpTest :: Assertion
qpTest =
  -- From example here:
  -- http://www.mathworks.com/help/optim/ug/quadprog.html#zmw57dd0e69051
  let a = matrix 2 [1, (-1), (-1), 2]
      b = vector [(-2), (-6)]
      e = matrix 2 [ (-1), (-1)
                   , 1, (-2)
                   , (-2), (-1)
                   , 1, 0
                   , 0, 1 ]
      f = vector [2, 2, 3, 0, 0]
      result = solveQuadProg (a, b) Nothing $ Just (e, f)
  in  result @?= Right (vector [0.6666666666666669,
                                1.3333333333333335],
                        -8.22222222222222)

-- sqpTest :: Assertion
-- sqpTest =
--   let a = matrix 2 [1, (-1), (-1), 2]
--       b = vector [(-2), (-6)]
--       e = matrix 2 [ (1), (1)
--                    , (-1), (2)
--                    , (2), (1)
--                    , (-1), 0
--                    , 0, (-1) ]
--       f = vector [(-2), (-2), (-3), 0, 0]
--       affineIneqs x = (e, f)
--       ineqs x = (e #> x) + f
--       cost x = (x `dot` ((0.5 * (a #> x)) + b)) + 0.0
--       problem = Problem
--                 { _cost = cost
--                 , _approxCost = const (a, b, 0)
--                 , _trueIneqs = ineqs
--                 , _approxAffineIneqs = affineIneqs
--                 , _numVariables = 2
--                 , _numIneqs = 5
--                 }
--       (xResult, fResult) = optimize problem (vector [0, 0])
--       (xTrue, fTrue) = (vector [0.6666666666666669,
--                              1.3333333333333335],
--                         -8.22222222222222)
--       xError = norm_2 $ xTrue - xResult
--       fError = abs $ fTrue - fResult
--   in  (xError < 1e-8, fError < 1e-8) @?= (True, True)

getQuadraticApproximation :: ([Double] -> Double)
                          -> ([Double] -> [Double])
                          -> ([Double] -> [[Double]])
                          -> Vector Double
                          -> (Matrix Double, Vector Double, Double)
getQuadraticApproximation fL fL' fL'' x =
  let xL = toList x
      f' = fromList $ fL' xL
      f'' = makeSymmetric $ fromLists $ fL'' xL
      c  = fL xL - (x `dot` f') + 0.5 * (x `dot` (f'' #> x))
  in  ( f'', f' - (f'' #> x), c)

getLinearApproximation :: ([Double] -> [Double])
                       -> ([Double] -> [[Double]])
                       -> Vector Double
                       -> (Matrix Double, Vector Double)
getLinearApproximation fL fL' x =
  let xL = toList x
      f' = fromLists $ fL' xL
  in  (f', vector (fL xL) - (f' #> x))

-- rosenbrockTest :: Assertion
-- rosenbrockTest =
--   -- From: http://www.mathworks.com/help/optim/ug/constrained-nonlinear-optimization-algorithms.html#f26633
--   let f [x, y] = (1 - x)^(2 :: Int) + 100*(y - x^(2 :: Int))^(2 :: Int)
--       cost x = let xL = toList x in f xL
--       convexCost = getQuadraticApproximation f
--       g [x, y] = x^(2 :: Int) + y^(2 :: Int) - 1.5
--       ineqs x = let xL = toList x in vector [g xL]
--       convexIneqs = getLinearApproximation f
--       problem = Problem
--                 { _cost = cost
--                 , _approxCost = convexCost
--                 , _trueIneqs = ineqs
--                 , _approxAffineIneqs = convexIneqs
--                 , _numVariables = 2
--                 , _numIneqs = 1
--                 }
--       (xResult, fResult) = optimize problem (vector [(-1.9), 2.0])
--       (xTrue, fTrue) = (vector [0.9072, 0.8228],
--                         0.00861632761)
--       xError = norm_2 $ xTrue - xResult
--       fError = abs $ fTrue - fResult
--   in  (xError < 1e-4, fError < 1e-4) @?= (True, True)

makeSymmetric :: Matrix Double -> Matrix Double
makeSymmetric m =
  let upperTriIxs = concat
        [ [(r, c) | c <- [(r + 1)..(cols m - 1)]] | r <- [0..(rows m - 1)] ]
      symElems (r, c) = [ ((r, c), m ! r ! c)
                        , ((c, r), m ! r ! c) ]
  in  (diag . takeDiag) m + assoc (size m) 0.0 (concatMap symElems upperTriIxs)

-- From: http://www.mathworks.com/help/optim/ug/nonlinear-equality-and-inequality-constraints.html
nonlinearEqTest :: Assertion
nonlinearEqTest =
  let f [x, y] = exp(x) * (4*x*x + 2*y*y + 4*x*y + 2*y + 1)
      g [x, y] = [negate (x * y) - 10]
      h [x, y] = [(x * x) + y - 1]

      cost = f . toList
      approxCost = getQuadraticApproximation f (grad f) (hessian f)
      ineq = vector . g . toList
      approxIneq = getLinearApproximation g (jacobian g)
      eq = vector . h . toList
      approxEq = getLinearApproximation h (jacobian h)
      problem = Problem
                { _cost = cost
                , _approxCost = approxCost
                , _trueIneqs = ineq
                , _approxAffineIneqs = approxIneq
                , _eqs = eq
                , _approxEqs = approxEq
                , _numVariables = 2
                , _numIneqs = 1
                , _numEqs = 1
                }
      (xResult, fResult) = optimize problem $ vector [(-1.0), 1.0]
      (xTrue, fTrue) = (vector [(-0.7529), 0.4332], 1.5093)
  in  (xResult, fResult) @?= (xTrue, fTrue)

main = defaultMain tests
  where tests = [ -- testCase "qp-test" qpTest
                -- , testCase "sqp-test" sqpTest
                -- , testCase "rosenbrock-test" rosenbrockTest
                  testCase "nonlineareq-test" nonlinearEqTest
                ]


module RotationPath where

import Control.Monad.Writer (Writer)
import qualified Control.Monad.Writer as Writer
import Data.DList (DList)
import qualified Data.DList as DList
import Numeric.LinearAlgebra ((<>))

import Cube (Cube)
import qualified Cube as Cube
import Rotation (MatrixRotation)
import qualified Rotation as Rotation

type RotationPath a = Writer (DList MatrixRotation) a
type CubeMutation = Cube -> RotationPath Cube

rotate :: MatrixRotation -> CubeMutation
rotate rotation cube = do
  let cube' = Rotation.toCubeRotation (Rotation.toVectorRotation rotation) cube
  Writer.tell $ DList.singleton rotation
  return cube'

between :: MatrixRotation -> MatrixRotation -> CubeMutation -> CubeMutation
between prefix suffix λ cube = do
  let (cube', rotationLog) = Writer.runWriter . λ $ Rotation.rotate prefix cube
  Writer.tell $ DList.map (\r -> suffix <> r <> prefix) rotationLog
  return $ Rotation.rotate suffix cube'

rotations :: [MatrixRotation] -> CubeMutation
rotations [] cube = return cube
rotations (r:rs) cube = rotate r cube >>= rotations rs


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# OPTIONS_GHC -fno-warn-orphans  #-}

module Nudge (
    nudgeProp
  , nudgeProp2
  ) where

import           Control.Monad
import           Data.Bifunctor
import           Data.Finite
import           Data.Kind
import           Data.Maybe
import           Data.Proxy
import           GHC.TypeLits
import           Hedgehog
import           Lens.Micro
import           Lens.Micro.Platform                   ()
import           Numeric.Backprop
import qualified Data.Ix                               as Ix
import qualified Data.Vector.Sized                     as SV
import qualified Hedgehog.Gen                          as Gen
import qualified Hedgehog.Range                        as Range
import qualified Numeric.LinearAlgebra                 as HU
import qualified Numeric.LinearAlgebra.Static          as H
import qualified Numeric.LinearAlgebra.Static.Backprop as B

nudge :: Double
nudge = 1e-6

eps :: Double
eps = 1e-11

class (Backprop c, Show c, Show (TIx c)) => Testing c where
    type TIx c :: Type
    allIx  :: c -> [TIx c]
    ixLens :: TIx c -> Lens' c Double
    scalarize :: Reifies s W => BVar s c -> BVar s Double
    genTest :: Gen c

sized
    :: forall s t d. H.Sized t s d
    => Lens' s (d t)
sized f = fmap (fromJust . H.create) . f . H.extract

ixContainer
    :: forall t d. HU.Container d t
    => HU.IndexOf d
    -> Lens' (d t) t
ixContainer i = lens (`HU.atIndex` i)
                     (\xs x -> HU.accum xs (\_ _ -> x) [(i, x)])

instance Testing Double where
    type TIx Double = ()
    allIx _ = [()]
    ixLens _ = id
    scalarize = abs
    genTest = Gen.filter ((> eps) . (**2)) $
         Gen.double (Range.linearFracFrom 0 (-5) 5)

instance KnownNat n => Testing (H.R n) where
    type TIx (H.R n) = Int
    allIx v = [0 .. H.size v - 1]
    ixLens i = sized . ixContainer i
    scalarize = B.norm_2V
    genTest = H.vector <$> replicateM n genTest
      where
        n = fromInteger $ natVal (Proxy @n)

instance (KnownNat n, KnownNat m) => Testing (H.L n m) where
    type TIx (H.L n m) = (Int, Int)
    allIx m = Ix.range ((0,0), bimap pred pred (H.size m))
    ixLens i = sized . ixContainer i
    scalarize = sqrt . B.sumElements . (**2)
    genTest = H.matrix <$> replicateM nm genTest
      where
        nm = fromInteger $ natVal (Proxy @n) * natVal (Proxy @m)

instance Testing (HU.Vector Double) where
    type TIx (HU.Vector Double) = Int
    allIx v = [0 .. HU.size v - 1]
    ixLens = ixContainer
    scalarize = liftOp1 . op1 $ \xs -> (HU.sumElements xs, (`HU.konst` HU.size xs))
    genTest = HU.fromList <$> replicateM 3 genTest

instance Testing (HU.Matrix Double) where
    type TIx (HU.Matrix Double) = (Int, Int)
    allIx m = Ix.range ((0,0), bimap pred pred (HU.size m))
    ixLens = ixContainer
    scalarize = liftOp1 . op1 $ \xs -> (HU.sumElements xs, (`HU.konst` HU.size xs))
    genTest = HU.fromLists <$> (replicateM 3 . replicateM 2) genTest

instance (KnownNat n, Testing a, Num a) => Testing (SV.Vector n a) where
    type TIx (SV.Vector n a) = (Finite n, TIx a)
    allIx = fst . SV.imapM (\i x -> ((fromIntegral i,) <$> allIx x , x))
    ixLens (i,j) = SV.ix i . ixLens j
    scalarize = scalarize . liftOp1 o . (^ (2 :: Int))
      where
        o :: Op '[SV.Vector n a] a
        o = op1 $ \xs -> (SV.sum xs, SV.replicate)
    genTest = SV.replicateM genTest

instance (Testing a, Testing b) => Testing (a, b) where
    type TIx (a, b) = Either (TIx a) (TIx b)
    allIx (x, y) = (Left  <$> allIx x)
                ++ (Right <$> allIx y)
    ixLens (Left  i) = _1 . ixLens i
    ixLens (Right j) = _2 . ixLens j
    scalarize t = B.norm_2V (B.vec2 (scalarize (t ^^. _1))
                                    (scalarize (t ^^. _2))
                            )
    genTest = (,) <$> genTest <*> genTest

instance (Testing a, Testing b, Testing c) => Testing (a, b, c) where
    type TIx (a, b, c) = Either (TIx a) (Either (TIx b) (TIx c))
    allIx (x, y, z) = (Left          <$> allIx x)
                   ++ (Right . Left  <$> allIx y)
                   ++ (Right . Right <$> allIx z)
    ixLens (Left         i ) = _1 . ixLens i
    ixLens (Right (Left  j)) = _2 . ixLens j
    ixLens (Right (Right k)) = _3 . ixLens k
    scalarize t = B.norm_2V (B.vec3 (scalarize (t ^^. _1))
                                    (scalarize (t ^^. _2))
                                    (scalarize (t ^^. _3))
                            )
    genTest = (,,) <$> genTest <*> genTest <*> genTest

validGrad
    :: Monad m
    => Lens' c Double
    -> c
    -> c
    -> (c -> Double)
    -> PropertyT m (Double, Double)
validGrad l x0 g f = forAll $ Gen.double (Range.constantFrom 0 (-nudge) nudge) <&> \d ->
    let x   = x0 & l %~ (+d)
        old = f x0 + (g ^. l) * d
        new = f x
    in  (old, new)

nudgeProp
    :: forall c d. (Testing c, Testing d)
    => (forall s. Reifies s W => BVar s c -> BVar s d)
    -> Property
nudgeProp f = property $ do
    (inp, i) <- forAll $ do
      inp <- genTest
      i   <- Gen.element (allIx inp)
      return (inp, i)
    let (r,gr) = backprop (scalarize . f) inp
    when (r**2 < eps) discard
    (old, new) <- validGrad (ixLens i) inp gr (evalBP (scalarize . f))
    footnoteShow (r, gr, old, new, (old - new)**2, ((old - new)/old)**2)
    assert $ ((old - new)/old)**2 < eps

nudgeProp2
    :: forall c d e. (Testing c, Testing d, Testing e)
    => (forall s. Reifies s W => BVar s c -> BVar s d -> BVar s e)
    -> Property
nudgeProp2 f = property $ do
    (inpC, inpD, i) <- forAll $ do
      inpC <- genTest
      inpD <- genTest
      i    <- Gen.element (allIx (inpC, inpD))
      return (inpC, inpD, i)
    let (r, gr) = backprop2 (\x -> scalarize . f x) inpC inpD
    when (r**2 < eps) discard
    (old, new) <- validGrad (ixLens i) (inpC, inpD) gr
          (evalBP (\t -> scalarize $ f (t ^^. _1) (t ^^. _2)))
    footnoteShow (r, gr, old, new, (old - new)**2, ((old - new)/old)**2)
    assert $ ((old - new)/old)**2 < eps

instance (HU.Container HU.Vector a, Num a) => Backprop (HU.Matrix a) where
    -- TODO: make more efficient?
    zero = HU.cmap (const 0)
    add  = HU.add
    one  = HU.cmap (const 1)

instance (KnownNat n, Num a) => Backprop (SV.Vector n a) where
    zero = (0 <$)
    add  = (+)
    one  = (1 <$)


{-# LANGUAGE BangPatterns #-}
module Main where

import Criterion.Main
import Data.Massiv.Array as A
--import Data.Massiv.Array.SIMD
--import Data.Massiv.Array.Manifest.Vector as A
import Data.Massiv.Array.Unsafe as A
import Data.Massiv.Bench as A
import Prelude as P hiding ((<>))
import Data.Semigroup

--import Statistics.Matrix as S
--import Statistics.Matrix.Fast as SF

multArrsAlt :: Array U Ix2 Double -> Array U Ix2 Double -> Array U Ix2 Double
multArrsAlt arr1 arr2
  | n1 /= m2 =
    error $
    "(|*|): Inner array dimensions must agree, but received: " ++
    show (size arr1) ++ " and " ++ show (size arr2)
  | otherwise = compute $
    makeArrayR D (getComp arr1 <> getComp arr2) (Sz (m1 :. n2)) $ \(i :. j) ->
      A.foldlS (+) 0 (A.zipWith (*) (unsafeOuterSlice arr1 i) (unsafeOuterSlice arr2' j))
  where
    Sz2 m1 n1 = size arr1
    Sz2 m2 n2 = size arr2
    arr2' = computeAs U $ A.transpose arr2
{-# INLINE multArrsAlt #-}

main :: IO ()
main = do
  let !sz@(Sz2 _m _n) = Sz2 600 600
      !arr = arrRLightIx2 U Seq sz
      -- !arr2 = computeAs U arr
      -- !mat = S.Matrix m n $ A.toVector arr
      -- !mat' = S.transpose mat
      -- !arrV = arrRLightIx2 V Seq sz
      -- !arrV2 = computeAs V arrV
  defaultMain
    [ env (return (computeAs U (A.transpose arr))) $ \ arr' ->
        bgroup
          "Mult"
          [ bgroup
              "Seq"
              [ bench "(|*|)" $ whnfIO (setComp Seq arr |*| arr')
              -- , bench "multiplyTranspose" $
              --   whnf (computeAs U . multiplyTransposed (setComp Seq arr)) arr2
              -- , bench "multiplyTransposeSIMD" $
              --   whnf (computeAs P . multiplyTransposedSIMD arrV) arrV2
              , bench "multArrsAlt" $ whnf (multArrsAlt (setComp Seq arr)) arr'
              --, bench "multiply (dense-linear-algebra)" $ whnf (SF.multiply mat) mat'
              ]
          , bgroup
              "Par"
              [ bench "(|*|)" $ whnfIO (setComp Par arr |*| arr')
              , bench "fused (|*|)" $ whnfIO (setComp Par arr |*| A.transpose arr)
              -- , bench "multiplyTranspose" $
              --   whnf (computeAs U . multiplyTransposed (setComp Par arr)) arr2
              -- , bench "multiplyTransposeSIMD" $
              --   whnf (computeAs P . multiplyTransposedSIMD (setComp Par arrV)) arrV2
              , bench "multArrsAlt" $ whnf (multArrsAlt (setComp Par arr)) arr'
              ]
          ]
    ]


module NNSpec where

import Numeric.LinearAlgebra as LA
import Test.Hspec
import NN

spec :: Spec
spec =
  describe "Toy Neural Network" $

--    it "Train Point" $ do
--        let wT1 = (3 >< 2) $ repeat 1.0
--            b1 = (3 >< 1) $ repeat 0.0
--            vT1 = (1 >< 3) $ repeat 1.0
--            learningRate = 0.001
--            res = trainPoint wT1 b1 vT1 learningRate 2 1 10
--        res `shouldBe` (wT1, b1, vT1)

    it "Test" $ do
        let wT = (3 >< 2) $ repeat 1.0
            b = (3 >< 1) $ repeat 0.0
            vT = (1 >< 3) $ repeat 1.0
            learningRate = 0.001
            epochs = 10
            xyPairs = [((2.0, 1.0), 10.0), ((3.0, 3.0), 21.0), ((4.0, 5.0), 32.0), ((6.0, 6.0), 42.0)]
            (wT', b', vT') = trainNeuralNetwork xyPairs wT b vT learningRate epochs
            x = [(1.0,1.0), (2.0,2.0), (3.0,3.0), (5.0,5.0), (10.0,10.0)]
            expected = fmap (snd . fmap (* 7)) x
            res = testNeuralNetwork x wT' b' vT' learningRate
            res' = map ((< 0.1) . abs) $ zipWith (-) res expected
        res' `shouldBe` [True, True, True, True, True]


{-# LANGUAGE AllowAmbiguousTypes, DataKinds, RankNTypes, TypeFamilies,
             TypeOperators #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
module TestMnistCNN (testTrees, shortTestForCITrees) where

import Prelude

import           Control.Arrow (first)
import           Control.Monad (foldM)
import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.Proxy (Proxy (Proxy))
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, SomeNat (..), someNatVal, type (<=))
import           Numeric.LinearAlgebra (Matrix, Vector)
import qualified Numeric.LinearAlgebra as HM
import           System.IO (hPutStrLn, stderr)
import           System.Random
import           Test.Tasty
import           Test.Tasty.HUnit hiding (assert)
import           Test.Tasty.QuickCheck hiding (label, scale, shuffle)
import           Text.Printf

import HordeAd
import HordeAd.Core.DualClass (HasRanks (dKonst2), IsPrimal (dZero))
import HordeAd.Tool.MnistCnnShaped
import HordeAd.Tool.MnistTools

import TestCommon

testTrees :: [TestTree]
testTrees = [ mnistCNNTestsShort
            , mnistCNNTestsLong
            ]

shortTestForCITrees :: [TestTree]
shortTestForCITrees = [ mnistCNNTestsShort
                      ]

-- * The simplest possible convolutional net, based on
-- https://www.ritchieng.com/machine-learning/deep-learning/tensorflow/convnets/#Problem-1
-- but with different initialization and no batches and the picture size
-- evolves differently (@conv2@ used instead of @convSame2@). Theirs is not
-- real convolution but, most likely, correlation, and their padding
-- only preserves size, while ours in @conv2@ increases it,
-- not to put less weigth onto information from the outer rows and columns.

patch_size, depth0, num_hidden0, final_image_size :: Int
patch_size = 5
depth0 = 16
num_hidden0 = 64
final_image_size = 10  -- if size was not increased: 7, see below

lenMnistCNN :: Int -> Int -> Int -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
lenMnistCNN final_image_sz depth num_hidden =
  ( depth + depth
  , [num_hidden, sizeMnistLabel]
  , replicate (depth + depth * depth) (patch_size, patch_size)
    ++ [(num_hidden, final_image_sz * final_image_sz * depth)]
    ++ [(sizeMnistLabel, num_hidden)]
  , []
  )

-- This is simple convolution with depth 1.
convDataMnistCNN :: DualMonad d r m
                 => DualNumberVariables d r -> Matrix r -> Int
                 -> m (DualNumber d (Matrix r))
convDataMnistCNN variables x offset = do
  let ker = var2 variables offset
      bias = var0 variables offset
  yConv@(D u _) <- conv2 ker (D x (dKonst2 dZero (HM.size x)))  -- == (scalar x)
  yRelu <- reluAct $ yConv + konst2 bias (HM.size u)
  maxPool2 2 2 yRelu

-- This simulates convolution of nontrivial depth, without using tensors.
convMiddleMnistCNN :: DualMonad d r m
                   => Int -> DualNumberVariables d r
                   -> [DualNumber d (Matrix r)] -> Int
                   -> m (DualNumber d (Matrix r))
convMiddleMnistCNN depth variables ms1 k = do
  let conv (m, n) = do
        let ker = var2 variables ((1 + k) * depth + n)
        conv2 ker m
  ms2 <- mapM conv $ zip ms1 [0 ..]
  yConv@(D u _) <- returnLet $ sum ms2
  let bias = var0 variables (depth + k)
  yRelu <- reluAct $ yConv + konst2 bias (HM.size u)
  maxPool2 2 2 yRelu

convMnistCNN :: DualMonad d r m
             => Int -> Matrix r  -- 28x28
             -> DualNumberVariables d r
             -> m (DualNumber d (Vector r))
convMnistCNN depth x variables = do
  ms1 <- mapM (convDataMnistCNN variables x) [0 .. depth - 1]
  ms3 <- mapM (convMiddleMnistCNN depth variables ms1) [0 .. depth - 1]
  let flattenAppend m = append1 (flatten1 m)
  v <- returnLet $ foldr flattenAppend (seq1 V.empty) ms3
  let weigthsDense = var2 variables (depth + depth * depth)
      biasesDense = var1 variables 0
      denseLayer = weigthsDense #>! v + biasesDense
  denseRelu <- reluAct denseLayer
  let weigthsReadout = var2 variables (depth + depth * depth + 1)
      biasesReadout = var1 variables 1
  returnLet $ weigthsReadout #>! denseRelu + biasesReadout

convMnistLossCNN :: DualMonad d r m
                 => Int -> MnistData2 r
                 -> DualNumberVariables d r
                 -> m (DualNumber d r)
convMnistLossCNN depth (x, target) variables = do
  result <- convMnistCNN depth x variables
  lossSoftMaxCrossEntropyV target result

convMnistTestCNN
  :: forall r. IsScalar 'DModeGradient r
  => Proxy r -> Int -> [MnistData2 r] -> Domains r -> r
convMnistTestCNN _ depth inputs parameters =
  let matchesLabels :: MnistData2 r -> Bool
      matchesLabels (glyph, label) =
        let nn variables = do
              m <- convMnistCNN depth glyph variables
              softMaxActV m
            value = primalValue nn parameters
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)
{-# SPECIALIZE convMnistTestCNN :: Proxy Double -> Int -> [MnistData2 Double] -> Domains Double -> Double #-}

-- Here, unlike in
-- https://www.ritchieng.com/machine-learning/deep-learning/tensorflow/convnets/#Problem-1
-- we don't batch yet.
convMnistTestCaseCNN
  :: String
  -> Int
  -> Int
  -> (Int
      -> MnistData2 Double
      -> DualNumberVariables 'DModeGradient Double
      -> DualMonadGradient Double (DualNumber 'DModeGradient Double))
  -> (Proxy Double -> Int -> [MnistData2 Double]
      -> Domains Double -> Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> TestTree
convMnistTestCaseCNN prefix epochs maxBatches trainWithLoss testLoss
                     final_image_sz widthHidden widthHidden2 gamma expected =
  let ( (nParams0, nParams1, nParams2, nParamsX)
       , totalParams, range, parameters0 ) =
        initializerFixed 44 0.05
        (lenMnistCNN final_image_sz widthHidden widthHidden2)
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show widthHidden, show widthHidden2
                        , show nParams0, show nParams1, show nParams2
                        , show nParamsX
                        , show totalParams, show gamma, show range]
  in testCase name $ do
       hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
              prefix epochs maxBatches
       trainData <- loadMnistData2 trainGlyphsPath trainLabelsPath
       testData <- loadMnistData2 testGlyphsPath testLabelsPath
       -- Mimic how backprop tests and display it, even though tests
       -- should not print, in principle.
       let runBatch :: Domains Double
                    -> (Int, [MnistData2 Double])
                    -> IO (Domains Double)
           runBatch (!params0, !params1, !params2, !paramsX) (k, chunk) = do
             let f = trainWithLoss widthHidden
                 res = fst $ sgd gamma f chunk
                                 (params0, params1, params2, paramsX)
                 !trainScore = testLoss (Proxy @Double)
                                       widthHidden chunk res
                 !testScore = testLoss (Proxy @Double)
                                      widthHidden testData res
                 !lenChunk = length chunk
             hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
             hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
             hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
             return res
       let runEpoch :: Int
                    -> Domains Double
                    -> IO (Domains Double)
           runEpoch n params2 | n > epochs = return params2
           runEpoch n params2 = do
             hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
                 chunks = take maxBatches
                          $ zip [1 ..] $ chunksOf 5000 trainDataShuffled
             !res <- foldM runBatch params2 chunks
             runEpoch (succ n) res
       res <- runEpoch 1 parameters0
       let testErrorFinal = 1 - testLoss (Proxy @Double)
                                         widthHidden testData res
       testErrorFinal @?= expected

-- * Another flavour of the simplest possible convolutional net, based on
-- https://www.ritchieng.com/machine-learning/deep-learning/tensorflow/convnets/#Problem-1
-- but with different initialization and no batches.
-- Also, if @conv2@ was used instead of @convSame2@,
-- the picture size would evolve differently. Theirs is not
-- real convolution but, most likely, correlation, and their padding
-- only preserves size, while ours in @conv2@ increases it,
-- not to put less weigth onto information from the outer rows and columns.

final_image_sizeS :: Int
final_image_sizeS = 7

-- This is simple convolution with depth 1.
convDataMnistCNNS :: DualMonad d r m
                  => DualNumberVariables d r -> Matrix r -> Int
                  -> m (DualNumber d (Matrix r))
convDataMnistCNNS variables x offset = do
  let ker = var2 variables offset
      bias = var0 variables offset
  yConv@(D u _) <- convSame2 ker (constant x)
  yRelu <- reluAct $ yConv + konst2 bias (HM.size u)
  maxPool2 2 2 yRelu

-- This simulates convolution of nontrivial depth, without using tensors.
convMiddleMnistCNNS :: DualMonad d r m
                    => Int -> DualNumberVariables d r
                    -> [DualNumber d (Matrix r)] -> Int
                    -> m (DualNumber d (Matrix r))
convMiddleMnistCNNS depth variables ms1 k = do
  let conv (m, n) = do
        let ker = var2 variables ((1 + k) * depth + n)
        convSame2 ker m
  ms2 <- mapM conv $ zip ms1 [0 ..]
  yConv@(D u _) <- returnLet $ sum ms2
  let bias = var0 variables (depth + k)
  yRelu <- reluAct $ yConv + konst2 bias (HM.size u)
  maxPool2 2 2 yRelu

convMnistCNNS :: DualMonad d r m
              => Int -> Matrix r  -- 28x28
              -> DualNumberVariables d r
              -> m (DualNumber d (Vector r))
convMnistCNNS depth x variables = do
  ms1 <- mapM (convDataMnistCNNS variables x) [0 .. depth - 1]
  ms3 <- mapM (convMiddleMnistCNNS depth variables ms1) [0 .. depth - 1]
  let flattenAppend m = append1 (flatten1 m)
  v <- returnLet $ foldr flattenAppend (seq1 V.empty) ms3
  let weigthsDense = var2 variables (depth + depth * depth)
      biasesDense = var1 variables 0
      denseLayer = weigthsDense #>! v + biasesDense
  denseRelu <- reluAct denseLayer
  let weigthsReadout = var2 variables (depth + depth * depth + 1)
      biasesReadout = var1 variables 1
  returnLet $ weigthsReadout #>! denseRelu + biasesReadout

convMnistLossCNNS :: DualMonad d r m
                  => Int -> MnistData2 r
                  -> DualNumberVariables d r
                  -> m (DualNumber d r)
convMnistLossCNNS depth (x, target) variables = do
  result <- convMnistCNNS depth x variables
  lossSoftMaxCrossEntropyV target result

convMnistTestCNNS
  :: forall r. IsScalar 'DModeGradient r
  => Proxy r -> Int -> [MnistData2 r] -> Domains r -> r
convMnistTestCNNS _ depth inputs parameters =
  let matchesLabels :: MnistData2 r -> Bool
      matchesLabels (glyph, label) =
        let nn variables = do
              m <- convMnistCNNS depth glyph variables
              softMaxActV m
            value = primalValue nn parameters
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)
{-# SPECIALIZE convMnistTestCNNS :: Proxy Double -> Int -> [MnistData2 Double] -> Domains Double -> Double #-}

-- * A variant of @convMnistCNN@ with @conv2'@.

-- This is simple convolution with depth 1.
convDataMnistCNNP :: DualMonad d r m
                  => DualNumberVariables d r -> Matrix r -> Int
                  -> m (DualNumber d (Matrix r))
convDataMnistCNNP variables x offset = do
  let ker = var2 variables offset
      bias = var0 variables offset
  yConv@(D u _) <-
    returnLet $ conv2' ker (D x (dKonst2 dZero (HM.size x)))  -- == (scalar x)
  yRelu <- reluAct $ yConv + konst2 bias (HM.size u)
  maxPool2 2 2 yRelu

-- This simulates convolution of nontrivial depth, without using tensors.
convMiddleMnistCNNP :: DualMonad d r m
                    => Int -> DualNumberVariables d r
                    -> [DualNumber d (Matrix r)] -> Int
                    -> m (DualNumber d (Matrix r))
convMiddleMnistCNNP depth variables ms1 k = do
  let conv (m, n) = do
        let ker = var2 variables ((1 + k) * depth + n)
        returnLet $ conv2' ker m
  ms2 <- mapM conv $ zip ms1 [0 ..]
  yConv@(D u _) <- returnLet $ sum ms2
  let bias = var0 variables (depth + k)
  yRelu <- reluAct $ yConv + konst2 bias (HM.size u)
  maxPool2 2 2 yRelu

convMnistCNNP :: DualMonad d r m
              => Int -> Matrix r  -- 28x28
              -> DualNumberVariables d r
              -> m (DualNumber d (Vector r))
convMnistCNNP depth x variables = do
  ms1 <- mapM (convDataMnistCNNP variables x) [0 .. depth - 1]
  ms3 <- mapM (convMiddleMnistCNNP depth variables ms1) [0 .. depth - 1]
  let flattenAppend m = append1 (flatten1 m)
  v <- returnLet $ foldr flattenAppend (seq1 V.empty) ms3
  let weigthsDense = var2 variables (depth + depth * depth)
      biasesDense = var1 variables 0
      denseLayer = weigthsDense #>! v + biasesDense
  denseRelu <- reluAct denseLayer
  let weigthsReadout = var2 variables (depth + depth * depth + 1)
      biasesReadout = var1 variables 1
  returnLet $ weigthsReadout #>! denseRelu + biasesReadout

convMnistLossCNNP :: DualMonad d r m
                  => Int -> MnistData2 r
                  -> DualNumberVariables d r
                  -> m (DualNumber d r)
convMnistLossCNNP depth (x, target) variables = do
  result <- convMnistCNNP depth x variables
  lossSoftMaxCrossEntropyV target result

convMnistTestCNNP
  :: forall r. IsScalar 'DModeGradient r
  => Proxy r -> Int -> [MnistData2 r] -> Domains r -> r
convMnistTestCNNP _ depth inputs parameters =
  let matchesLabels :: MnistData2 r -> Bool
      matchesLabels (glyph, label) =
        let nn variables = do
              m <- convMnistCNNP depth glyph variables
              softMaxActV m
            value = primalValue nn parameters
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)
{-# SPECIALIZE convMnistTestCNNP :: Proxy Double -> Int -> [MnistData2 Double] -> Domains Double -> Double #-}

-- * A variant of @convMnistCNN@ with shaped tensors, including mini-batches

convMnistTestCaseCNNT
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , r ~ Double, d ~ 'DModeGradient, m ~ DualMonadGradient Double )
  => String
  -> Int
  -> Int
  -> (forall kheight_minus_1' kwidth_minus_1' num_hidden' out_channels'
             in_height' in_width' batch_size'.
      ( KnownNat kheight_minus_1', KnownNat kwidth_minus_1'
      , KnownNat num_hidden', KnownNat out_channels'
      , KnownNat in_height', KnownNat in_width', KnownNat batch_size'
      , 1 <= kheight_minus_1'
      , 1 <= kwidth_minus_1'
      , DualMonad d r m )
      => Proxy kheight_minus_1'
      -> Proxy kwidth_minus_1'
      -> Proxy num_hidden'
      -> Proxy out_channels'
      -> ( OS.Array '[batch_size', in_height', in_width'] r
         , OS.Array '[batch_size', SizeMnistLabel] r )
      -> DualNumberVariables d r
      -> m (DualNumber d r))
  -> (forall kheight_minus_1' kwidth_minus_1' num_hidden' out_channels'
             in_height' in_width'.
      ( KnownNat kheight_minus_1', KnownNat kwidth_minus_1'
      , KnownNat num_hidden', KnownNat out_channels'
      , KnownNat in_height', KnownNat in_width'
      , 1 <= kheight_minus_1'
      , 1 <= kwidth_minus_1'
      , IsScalar d r )
      => Proxy r
      -> Proxy kheight_minus_1'
      -> Proxy kwidth_minus_1'
      -> Proxy num_hidden'
      -> Proxy out_channels'
      -> [( OS.Array '[in_height', in_width'] r
          , OS.Array '[SizeMnistLabel] r )]
      -> Domains r
      -> r)
  -> (forall kheight_minus_1' kwidth_minus_1' num_hidden' out_channels'
             in_height' in_width'.
      ( KnownNat kheight_minus_1', KnownNat kwidth_minus_1'
      , KnownNat num_hidden', KnownNat out_channels'
      , KnownNat in_height', KnownNat in_width' )
      => Proxy kheight_minus_1'
      -> Proxy kwidth_minus_1'
      -> Proxy num_hidden'
      -> Proxy out_channels'
      -> Proxy in_height'
      -> Proxy in_width'
      -> (Int, [Int], [(Int, Int)], [OT.ShapeL]))
  -> Double
  -> Double
  -> TestTree
convMnistTestCaseCNNT prefix epochs maxBatches trainWithLoss ftest flen
                      gamma expected =
  let proxy_kheight_minus_1 = Proxy @kheight_minus_1
      proxy_kwidth_minus_1 = Proxy @kwidth_minus_1
      proxy_num_hidden = Proxy @num_hidden
      proxy_out_channels  = Proxy @out_channels
      proxy_in_height = Proxy @in_height
      proxy_in_width = Proxy @in_width
      batch_size = valueOf @batch_size
      ((_, _, _, nParamsX), totalParams, range, parametersInit) =
        initializerFixed 44 0.05
          (flen proxy_kheight_minus_1 proxy_kwidth_minus_1
                proxy_num_hidden proxy_out_channels
                proxy_in_height proxy_in_width)
      name = prefix ++ ": "
             ++ unwords [ show epochs, show maxBatches
                        , show (valueOf @num_hidden :: Int), show batch_size
                        , show nParamsX, show totalParams
                        , show gamma, show range ]
      packBatchS :: [( OS.Array '[in_height, in_width] r
                    , OS.Array '[SizeMnistLabel] r )]
                -> ( OS.Array '[batch_size, in_height, in_width] r
                   , OS.Array '[batch_size, SizeMnistLabel] r )
      packBatchS l =
        let (inputs, targets) = unzip l
        in (OS.ravel $ OSB.fromList inputs, OS.ravel $ OSB.fromList targets)
      shapeBatchS :: MnistData r
                  -> ( OS.Array '[in_height, in_width] r
                     , OS.Array '[SizeMnistLabel] r )
      shapeBatchS (input, target) = (OS.fromVector input, OS.fromVector target)
  in testCase name $ do
    hPutStrLn stderr $ printf "\n%s: Epochs to run/max batches per epoch: %d/%d"
           prefix epochs maxBatches
    trainData <- map shapeBatchS
                 <$> loadMnistData trainGlyphsPath trainLabelsPath
    testData <- take 100  -- TODO: reduced for now, because too slow
                . map shapeBatchS
                <$> loadMnistData testGlyphsPath testLabelsPath
     -- There is some visual feedback, because some of these take long.
    let runBatch :: Domains r
                 -> (Int, [( OS.Array '[in_height, in_width] r
                           , OS.Array '[SizeMnistLabel] r )])
                 -> IO (Domains r)
        runBatch parameters@(!_, !_, !_, !_) (k, chunk) = do
          let f = trainWithLoss proxy_kheight_minus_1 proxy_kwidth_minus_1
                                proxy_num_hidden proxy_out_channels
              chunkS = map packBatchS
                       $ filter (\ch -> length ch >= batch_size)
                       $ chunksOf batch_size chunk
              res = fst $ sgd gamma f chunkS parameters
              !trainScore = ftest (Proxy @r)
                                 proxy_kheight_minus_1 proxy_kwidth_minus_1
                                 proxy_num_hidden proxy_out_channels
                                 chunk res
              !testScore = ftest (Proxy @r)
                                proxy_kheight_minus_1 proxy_kwidth_minus_1
                                proxy_num_hidden proxy_out_channels
                                testData res
              !lenChunk = length chunk
          hPutStrLn stderr $ printf "\n%s: (Batch %d with %d points)" prefix k lenChunk
          hPutStrLn stderr $ printf "%s: Training error:   %.2f%%" prefix ((1 - trainScore) * 100)
          hPutStrLn stderr $ printf "%s: Validation error: %.2f%%" prefix ((1 - testScore ) * 100)
          return res
    let runEpoch :: Int -> Domains r -> IO (Domains r)
        runEpoch n params2 | n > epochs = return params2
        runEpoch n params2 = do
          hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n
          let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData
              chunks = take maxBatches
                       $ zip [1 ..]
                       $ chunksOf (2 * batch_size) trainDataShuffled
                           -- TODO: (10 * batch_size) takes forever
          !res <- foldM runBatch params2 chunks
          runEpoch (succ n) res
    res <- runEpoch 1 parametersInit
    let testErrorFinal = 1 - ftest (Proxy @r)
                                   proxy_kheight_minus_1 proxy_kwidth_minus_1
                                   proxy_num_hidden proxy_out_channels
                                   testData res
    testErrorFinal @?= expected

mnistCNNTestsLong :: TestTree
mnistCNNTestsLong = testGroup "MNIST CNN long tests"
  [ {-convMnistTestCaseCNN "artificial 5 4 3 2 1" 5 4
                         convMnistLossCNN convMnistTestCNN final_image_size
                         3 2 1 0.8991
  , convMnistTestCaseCNN "S artificial 5 4 3 2 1" 5 4
                         convMnistLossCNNS convMnistTestCNNS final_image_sizeS
                         3 2 1 0.8991
  , -}convMnistTestCaseCNN "P artificial 5 4 3 2 1" 5 4
                         convMnistLossCNNP convMnistTestCNNP final_image_size
                         3 2 1 0.8991
  , convMnistTestCaseCNNT @4 @4 @2 @3 @SizeMnistHeight @SizeMnistWidth @1
                          "T artificial 5 4 3 2 1" 5 4
                          convMnistLossFusedS convMnistTestS convMnistLenS
                          0.02 0.98
  , convMnistTestCaseCNN "1 epoch 1 batch" 1 1
                         convMnistLossCNN convMnistTestCNN
                         final_image_size depth0 num_hidden0
                         0.02 5.989999999999995e-2
{-
  , convMnistTestCaseCNN "2 epochs but only 1 batch" 2 1
                         convMnistLossCNN convMnistTestCNN
                         final_image_size depth0 num_hidden0
                         0.02 8.879999999999999e-2  -- dummy results everywhere
  , convMnistTestCaseCNN "1 epoch all batches" 1 99
                         convMnistLossCNN convMnistTestCNN
                         final_image_size depth0 num_hidden0
                         0.02 5.1100000000000034e-2
-}
  , convMnistTestCaseCNN "S1 epoch 1 batch" 1 1
                         convMnistLossCNNS convMnistTestCNNS
                         final_image_sizeS depth0 num_hidden0
                         0.02 4.800000000000004e-2
{-
  , convMnistTestCaseCNN "S2 epochs but only 1 batch" 2 1
                         convMnistLossCNNS convMnistTestCNNS
                         final_image_sizeS depth0 num_hidden0
                         0.02 8.879999999999999e-2
  , convMnistTestCaseCNN "S1 epoch all batches" 1 99
                         convMnistLossCNNS convMnistTestCNNS
                         final_image_sizeS depth0 num_hidden0
                         0.02 5.1100000000000034e-2
-}
  , convMnistTestCaseCNN "P1 epoch 1 batch" 1 1
                         convMnistLossCNNP convMnistTestCNNP
                         final_image_size depth0 num_hidden0
                         0.02 5.989999999999995e-2
{-
  , convMnistTestCaseCNN "P2 epochs but only 1 batch" 2 1
                         convMnistLossCNNP convMnistTestCNNP
                         final_image_size depth0 num_hidden0
                         0.02 4.94e-2
  , convMnistTestCaseCNN "P1 epoch all batches" 1 99
                         convMnistLossCNNP convMnistTestCNNP
                         final_image_size depth0 num_hidden0
                         0.02 2.7000000000000024e-2
-}
  , convMnistTestCaseCNNT @4 @4 @64 @16 @SizeMnistHeight @SizeMnistWidth @16
                          "T1 epoch 1 batch" 1 1
                          convMnistLossFusedS convMnistTestS convMnistLenS
                          0.02 0.8200000000000001
  , testProperty "Compare gradients and two forward derivatives for a single 2d convolution implemented from primitive operations and as a hardwired primitive" $
      forAll (choose (1, 30)) $ \seed ->
      forAll (choose (1, 50)) $ \seedDs ->
      forAll (choose (1, 100)) $ \widthHidden ->
      forAll (choose (1, 150)) $ \widthHidden2 ->
      forAll (choose (0, seed + widthHidden - 2)) $ \ix1 ->
      forAll (choose (0, seedDs + widthHidden2 - 2)) $ \ix2 ->
      forAll (choose (0.01, 10)) $ \range ->
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let paramShape =
              (0, [], [(seed, widthHidden2), (widthHidden, seedDs)], [])
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f, fP :: forall d r m. (DualMonad d r m)
                  => DualNumberVariables d r -> m (DualNumber d r)
            f variables = do
              let ker = var2 variables 0
                  x = var2 variables 1
              c <- conv2 ker x
              cx <- returnLet $ from2X c
              cx1 <- returnLet $ indexX cx ix1
              cx2 <- returnLet $ indexX cx1 ix2
              returnLet $ fromX0 cx2
            fP variables = do
              let ker = var2 variables 0
                  x = var2 variables 1
              c <- returnLet $ conv2' ker x
              cx <- returnLet $ from2X c
              cx1 <- returnLet $ indexX cx ix1
              cx2 <- returnLet $ indexX cx1 ix2
              returnLet $ fromX0 cx2
        in
            qcPropDom f  parameters ds parametersPerturbation 1 .&&.
            qcPropDom fP parameters ds parametersPerturbation 1 .&&.
            cmpTwoSimple f fP parameters ds

  , testProperty "Compare gradients and two forward derivatives for convMnistTestCNN and convMnistTestCNNP" $
      \seed ->
      forAll (choose (0, sizeMnistLabel - 1)) $ \seedDs ->
      forAll (choose (1, 20)) $ \depth ->
      forAll (choose (1, 30)) $ \num_hidden ->
      forAll (choose (0.01, 0.5)) $ \range ->
      forAll (choose (0.01, 10)) $ \rangeDs ->
        let createRandomVector n seedV = HM.randomVector seedV HM.Uniform n
            glyph = HM.reshape 28 $ createRandomVector (28 * 28) seed
            label = HM.konst 0 sizeMnistLabel V.// [(seedDs, 1)]
            mnistData :: MnistData2 Double
            mnistData = (glyph, label)
            paramShape = lenMnistCNN final_image_size depth num_hidden
            (_, _, _, parameters) = initializerFixed seed range paramShape
            (_, _, _, ds) = initializerFixed seedDs rangeDs paramShape
            (_, _, _, parametersPerturbation) =
              initializerFixed (seed + seedDs) 1e-7 paramShape
            f, fP, fT :: forall d r m. (DualMonad d r m, r ~ Double)
                  => DualNumberVariables d r -> m (DualNumber d r)
            f = convMnistLossCNN depth mnistData
            fP = convMnistLossCNNP depth mnistData
            fT = case ( someNatVal $ toInteger num_hidden
                      , someNatVal $ toInteger depth ) of
              ( Just (SomeNat proxy_num_hidden)
               ,Just (SomeNat proxy_out_channel) ) ->
                convMnistLossFusedS (Proxy @4) (Proxy @4)
                                    proxy_num_hidden proxy_out_channel
                                    (packBatch @1 [shapeBatch
                                                  $ first HM.flatten mnistData])
              _ -> error "fT panic"
            paramsToT (p0, p1, p2, _) =
              let qX = V.fromList
                    [ OT.fromVector [depth, 1, 5, 5]
                      $ V.concat $ map HM.flatten
                      $ take depth $ V.toList p2
                    , OT.fromVector [depth] $ V.take depth p0
                    , OT.fromVector [depth, depth, 5, 5]
                      $ V.concat $ map HM.flatten
                      $ take (depth * depth) (drop depth $ V.toList p2)
                    , OT.fromVector [depth] $ V.drop depth p0
                    , let m = p2 V.! (depth + depth * depth)
                      in OT.fromVector [num_hidden, HM.cols m]
                         $ HM.flatten m
                    , OT.fromVector [num_hidden] $ p1 V.! 0
                    , OT.fromVector [sizeMnistLabel, num_hidden]
                      $ HM.flatten
                      $ p2 V.! (depth + depth * depth + 1)
                    , OT.fromVector [sizeMnistLabel] $ p1 V.! 1
                    ]
              in (V.empty, V.empty, V.empty, qX)
            parametersT = paramsToT parameters
            dsT = paramsToT ds
        in
            qcPropDom f  parameters  ds  parametersPerturbation 1 .&&.
            qcPropDom fP parameters  ds  parametersPerturbation 1 .&&.
            qcPropDom fT parametersT dsT parametersPerturbation 1 .&&.
            cmpTwoSimple f fP parameters ds .&&.
            cmpTwo f fT parameters parametersT ds dsT
  ]

mnistCNNTestsShort :: TestTree
mnistCNNTestsShort = testGroup "MNIST CNN short tests"
  [ convMnistTestCaseCNN "artificial 1 1 1 1 1" 1 1
                         convMnistLossCNN convMnistTestCNN final_image_size
                         1 1 1 0.9026
  , convMnistTestCaseCNN "S artificial 1 1 1 1 1" 1 1
                         convMnistLossCNNS convMnistTestCNNS final_image_sizeS
                         1 1 1 0.9026
  , convMnistTestCaseCNN "P artificial 1 1 1 1 1" 1 1
                         convMnistLossCNNP convMnistTestCNNP final_image_size
                         1 1 1 0.9026
  , convMnistTestCaseCNNT @4 @4 @1 @1 @SizeMnistHeight @SizeMnistWidth @1
                          "T artificial 1 1 1 1 1" 1 1
                          convMnistLossFusedS convMnistTestS convMnistLenS
                          1 0.85
{-
  , convMnistTestCaseCNN "artificial 1 2 3 4 5" 1 2
                         convMnistLossCNN convMnistTestCNN final_image_size
                         3 4 5 0.902
  , convMnistTestCaseCNN "S artificial 1 2 3 4 5" 1 2
                         convMnistLossCNNS convMnistTestCNNS final_image_sizeS
                         3 4 5 0.902
  , convMnistTestCaseCNN "P artificial 1 2 3 4 5" 1 2
                         convMnistLossCNNP convMnistTestCNNP final_image_size
                         3 4 5 0.8972
-}
  , convMnistTestCaseCNNT @4 @4 @4 @3 @SizeMnistHeight @SizeMnistWidth @5
                          "T artificial 1 2 3 4 5" 1 2
                          convMnistLossFusedS convMnistTestS convMnistLenS
                          6 0.92
  ]


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-|

Module      : Grenade.Layers.Dropout
Description : Defines the Dropout Layer

Dropout is a regularization method used to prevent overfitting
of the network to the training dataset.
It works by randomly "dropping out", i.e. ignoring, outputs of
a percentage of neurons.
This is strategy is effective in preventing overfitting as it
forces a NN not to rely on a single feature for computing the
output of a neural network
-}
module Grenade.Layers.Dropout (
    Dropout (..)
  , randomDropout
  ) where

import           Control.DeepSeq
import           Control.Monad.Primitive        (PrimBase, PrimState)
import           Data.Proxy
import           Data.Serialize
import           GHC.Generics                   hiding (R)
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static   hiding (Seed)
import           System.Random.MWC

import           Grenade.Core

-- Dropout layer help to reduce overfitting.
-- Idea here is that the vector is a shape of 1s and 0s, which we multiply the input by.
-- After backpropogation, we return a new matrix/vector, with different bits dropped out.
-- The provided argument is the proportion to drop in each training iteration (like 1% or
-- 5% would be reasonable).
data Dropout (pct :: Nat) =
  Dropout
    { dropoutActive :: Bool     -- ^ Add possibility to deactivate dropout
    , dropoutSeed   :: !Int     -- ^ Seed
    }
  deriving (Generic)

instance NFData (Dropout pct) where rnf (Dropout a s) = rnf a `seq` rnf s
instance Show (Dropout pct) where show (Dropout _ _) = "Dropout"
instance Serialize (Dropout pct) where
  put (Dropout act seed) = put act >> put seed
  get = Dropout <$> get <*> get

instance UpdateLayer (Dropout pct) where
  type Gradient (Dropout pct) = ()
  runUpdate _ (Dropout act seed) _ = Dropout act (seed+1)
  runSettingsUpdate set (Dropout _ seed) = Dropout (setDropoutActive set) seed
  reduceGradient _ = ()

instance RandomLayer (Dropout pct) where
  createRandomWith _ = randomDropout

randomDropout :: (PrimBase m) => Gen (PrimState m) -> m (Dropout pct)
randomDropout gen = Dropout True <$> uniform gen

instance (KnownNat pct, KnownNat i) => Layer (Dropout pct) ('D1 i) ('D1 i) where
  type Tape (Dropout pct) ('D1 i) ('D1 i) = R i
  runForwards (Dropout act seed) (S1D x)
    | not act = (v, S1D $ dvmap (rate *) x) -- multily with rate to normalise throughput
    | otherwise = (v, S1D $ v * x)
    where
      rate = (/100) $ fromIntegral $ max 0 $ min 100 $ natVal (Proxy :: Proxy pct)
      v = dvmap mask $ randomVector seed Uniform
      mask r
        | not act || r < rate = 1
        | otherwise = 0
  runBackwards (Dropout _ _) v (S1D x) = ((), S1D $ x * v)


module Sound.Filter
(
    lzero2cp
    , lpole2cp
    , rlzero2cp
    , rlpole2cp
    , Cx
    , Z
)
where

import qualified Data.Complex as Cp

import Sound.InfList
import Sound.Time

lpole2cp :: Z -> L Double -> L Double
lpole2cp p x0 =
    y0
    where
        y0 = lzip3 (\ sy1 sy2 sx0 -> a * sy1 - b * sy2 + sx0) y1 y2 x0
        y1 = ldelay y0
        y2 = ldelay y1
        r = Cp.realPart p
        i = Cp.imagPart p
        a = 2 * r
        b = r * r + i * i

rlpole2cp :: Z -> Rated (L Double) -> Rated (L Double)
rlpole2cp p = rmap (lpole2cp p)

lzero2cp :: Z -> L Double -> L Double
lzero2cp p x0 =
    lzip3 (\ sx0 sx1 sx2 -> sx0 - a * sx1 + b * sx2) x0 x1 x2
    where
        x1 = ldelay x0
        x2 = ldelay x1
        r = Cp.realPart p
        i = Cp.imagPart p
        a = 2 * r
        b = r * r + i * i

rlzero2cp :: Z -> Rated (L Double) -> Rated (L Double)
rlzero2cp p = rmap (lzero2cp p)

-- | complex number
type Cx = Cp.Complex Double

-- | z-plane coordinate
type Z = Cx


module LayersSpec
  ( spec
  ) where

import           Control.Monad
import qualified Data.Set                as Set
import           Data.Tuple.Extra
import           Debug.Trace
import           Foundation.Monad
import           Layers
import           Learn
import           Mnist
import           Numeric.LinearAlgebra   as A
import           Synapse
import           Test.Hspec
import           Test.Hspec.QuickCheck   (prop)
import           Test.QuickCheck
import           Test.QuickCheck.Gen
import           Test.QuickCheck.Monadic

runTest :: IO ()
runTest = hspec spec

spec :: Spec
spec = do
  describe "predict" propsPredict
  describe "evaluate" propsEvaluate
  describe "Affine" propsAffine
  describe "Sigmoid" propsSigmoid
  describe "ReLU" propsReLU
  describe "Joined" propsJoined
  describe "Matrix" propsMatrix

propsPredict =
  prop "index" $
  forAll (genVectorN 100) $ \v ->
    let r = predict ReLUForward v
        a = maxIndex v
     in r `shouldBe` a

propsEvaluate = do
  prop "100%" $
    forAll (vectorOf 100 $ genVectorN 100) $ \vs ->
      let ns = map maxIndex vs
          r = evaluate ReLUForward $ zip ns vs
       in r `shouldBe` 1
  prop "0%" $
    forAll (vectorOf 100 $ genVectorN 100) $ \vs ->
      let ns = map ((+ 1) . maxIndex) vs
          r = evaluate ReLUForward $ zip ns vs
       in r `shouldBe` 0
  prop "50%" $
    forAll (vectorOf 100 $ genVectorN 100) $ \vs ->
      let (vA, vB) = splitAt 50 vs
          nA = map maxIndex vA
          nB = map ((+ 1) . maxIndex) vB
          ns = nA ++ nB
          r = evaluate ReLUForward $ zip ns vs
       in r `shouldBe` 0.5

propsAffine = do
  prop "forward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          c <- choose (10, 100)
          f <- genAffine a b
          s <- genMatrix c a
          return (f, s)) $ \(af, x) ->
      let AffineForward weights bias = af
          (AffineBackward rW rB rX, r) = forward af x
          e = affinem weights bias x
       in (rW, rB, rX, r) `shouldBe` (weights, bias, x, e)
  prop "backward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          c <- choose (10, 100)
          d <- genMatrix c b
          f <- genAffine a b
          s <- genMatrix c a
          return (f, s, d)) $ \(af, x, d) ->
      let AffineForward w b = af
          rate = 1.0
          (AffineForward w' b', d') = backward rate (AffineBackward w b x) d
          (dx, dw, db) = affinemBackward w x d
          rW = w - dw
          rB = b - db
       in (w', b', d') `shouldBe` (rW, rB, dx)
  prop "backward with rate" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          c <- choose (10, 100)
          r <- arbitrary `suchThat` (\a -> 0 < a && a < 1)
          d <- genMatrix c b
          f <- genAffine a b
          s <- genMatrix c a
          return (r, f, s, d)) $ \(rate, af, x, d) ->
      let AffineForward w b = af
          (AffineForward w' b', d') = backward rate (AffineBackward w b x) d
          (dx, dw, db) = affinemBackward w x d
          rW = w - A.scale rate dw
          rB = b - A.scale rate db
       in (w', b', d') `shouldBe` (rW, rB, dx)

propsSigmoid = do
  prop "forward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          genMatrix a b) $ \x ->
      let (SigmoidBackward y, r) = forward SigmoidForward x
          e = sigmoidm x
       in (y, r) `shouldBe` (r, e)
  prop "backward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          d <- genMatrix a b
          y <- genMatrix a b
          return (y, d)) $ \(y, d) ->
      let (SigmoidForward, d') = backward 0 (SigmoidBackward y) d
          e = sigmoidBackward y d
       in d' `shouldBe` e

propsReLU = do
  prop "forward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          genMatrix a b) $ \x ->
      let (ReLUBackward y, r) = forward ReLUForward x
          e = relum x
       in (y, r) `shouldBe` (x, e)
  prop "backward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          x <- genMatrix a b
          d <- genMatrix a b
          return (x, d)) $ \(x, d) ->
      let (ReLUForward, d') = backward 0 (ReLUBackward x) d
          e = relumBackward x d
       in d' `shouldBe` e

propsJoined = do
  prop "forward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          c <- choose (10, 100)
          f <- genAffine a b
          x <- genMatrix c a
          return (f, x)) $ \(fL1, x) ->
      let fL2 = SigmoidForward
          fJ = fL1 ~> fL2
          (JoinedBackwardLayer bL1 bL2, y) = forward fJ x
          (bL1', yL1) = forward fL1 x
          (bL2', yL2) = forward fL2 yL1
       in (bL1, bL2, y) `shouldBe` (bL1', bL2', yL2)
  prop "backward" $
    forAll
      (do a <- choose (10, 100)
          b <- choose (10, 100)
          c <- choose (10, 100)
          d <- genMatrix c b
          y <- genMatrix c b
          x <- genMatrix c a
          fL1 <- genAffine a b
          let AffineForward weights bias = fL1
          let bL1 = AffineBackward weights bias x
          return (bL1, SigmoidBackward y, d)) $ \(bL1, bL2, d) ->
      let bJ = bL1 <~ bL2
          (JoinedForwardLayer fL1 SigmoidForward, bD) = backward 1.0 bJ d
          (SigmoidForward, dL2) = backward 1.0 bL2 d
          (fL1', dL1) = backward 1.0 bL1 dL2
       in (fL1, bD) `shouldBe` (fL1', dL1)

propsMatrix = do
  prop "scale" $
    forAll genElements $ \(a, b, x, y) ->
      let m = b `A.scale` toMatrix x y a
          m' = toMatrix x y (a * b)
       in m `shouldBe` m'
  prop "multiply" $
    forAll genElements $ \(a, b, x, y) ->
      let m = toMatrix x y b * toMatrix x y a
          m' = toMatrix x y (a * b)
       in m `shouldBe` m'
  prop "subtract" $
    forAll genElements $ \(a, b, x, y) ->
      let m = toMatrix x y a - toMatrix x y b
          m' = toMatrix x y (a - b)
       in m `shouldBe` m'
  prop "add" $
    forAll genElements $ \(a, b, x, y) ->
      let m = toMatrix x y a + toMatrix x y b
          m' = toMatrix x y (a + b)
       in m `shouldBe` m'
  where
    toMatrix x y a = fromLists $ replicate x $ replicate y a
    genElements = do
      a <- arbitrary
      b <- arbitrary
      x <- choose (10, 100)
      y <- choose (10, 100)
      return (a :: R, b, x, y)

genVectorN :: Int -> Gen (Vector R)
genVectorN n = fmap fromList $ vectorOf n $ choose (1, 100)

genAffine :: Int -> Int -> Gen (ForwardLayer R)
genAffine nIn nOut = do
  weights <- fmap fromLists $ vectorOf nIn $ vectorOf nOut arbitrary
  bias <- fromList <$> vectorOf nOut arbitrary
  return $ AffineForward weights bias

genMatrix :: Int -> Int -> Gen (Matrix R)
genMatrix nRows nCols =
  fmap fromLists $ vectorOf nRows $ vectorOf nCols arbitrary