{-
module Solve
  ( Ord
  , DoubleMap
  , innerKeys
  , toMatrix
  , toMatrix'
  , trans
  , solve
  , smartSolve ) where
-}
module Solve where

import           Data.List             (elemIndex, nub, sort, union, (\\))
import           Data.List.Extra       (groupSortOn)
import           Data.Map.Strict       (Map)
import qualified Data.Map.Strict       as M
import           Data.Maybe            (catMaybes, mapMaybe)
import           Data.Tuple            (swap)
import           Numeric.LinearAlgebra

type DoubleMap k l a = Map k (Map l a)

pullMaybe :: Monad m => (m a, m b) -> m (a, b)
pullMaybe = (>>= (fmap swap . sequence . swap)) . sequence

innerKeys :: Eq l => DoubleMap k l a -> [l]
innerKeys = nub . concatMap M.keys . M.elems

keyToInt :: Ord k => k -> DoubleMap k l a -> Maybe Int
keyToInt n = elemIndex n . sort . M.keys

innerKeyToInt :: Ord l => l -> DoubleMap k l a -> Maybe Int
innerKeyToInt n = elemIndex n . sort . innerKeys

toAssocList :: DoubleMap k l a -> [((k, l), a)]
toAssocList = fmap f . concatMap sequence . M.toList . fmap M.toList
  where f (n, (i, a)) = ((n, i), a)

toAssocInt :: (Ord k, Ord l) =>  DoubleMap k l a -> [((Int, Int), a)]
toAssocInt b = mapMaybe f . toAssocList $ b
  where
    f ((n, i), a) = pullMaybe (pullMaybe (keyToInt n b, innerKeyToInt i b), Just a)

fromAssocList :: (Ord k, Ord l) => [((k, l), a)] -> DoubleMap k l a
fromAssocList = fmap M.fromList . M.fromList . fmap shift . groupSortOn fst . fmap f
  where
    shift xs = (fst . head $ xs, fmap snd xs)
    f ((k, l), a) = (k, (l, a))

trans :: (Ord k, Ord l) => DoubleMap k l a -> DoubleMap l k a
trans = fromAssocList . fmap f . toAssocList
  where f ((l, k), a) = ((k, l), a)

-- Outer keys are rows
toMatrix :: (Container Vector a, Num a, Ord k, Ord l)
  => DoubleMap k l a -> Matrix a
toMatrix b = assoc (M.size b, length $ innerKeys b) 0 (toAssocInt b)

-- Column matrix from a Map
toMatrix' :: (Container Vector a, Num a, Ord k) => Map k a -> Matrix a
toMatrix' = toMatrix . fmap (M.fromList . pure . ((),))

-- Are the rows independent?
isInd :: (Field a, Ord k, Ord l) => DoubleMap k l a -> Bool
isInd m = rank (toMatrix m) == M.size m

-- Choices of n elements from a list, fast algorithm from stack overflow.
-- https://stackoverflow.com/questions/21265454/subsequences-of-length-n-from-list-performance/21288092#21288092
choices :: Int -> [a] -> [[a]]
choices n xs =
  let l = length xs
  in if n>l then [] else subsequencesBySize xs !! (l-n)
  where
    subsequencesBySize [] = [[[]]]
    subsequencesBySize (y:ys) =
      let next = subsequencesBySize ys
      in zipWith (++) ([]:next) (map (map (y:)) next ++ [[]])

-- SubMaps with n keys, excluding some.
allSizeNSub :: Ord k => Int -> Map k v -> [Map k v]
allSizeNSub n m =
  [ M.filterWithKey (\k _ -> k `elem` keyList) m
  | keyList <- choices n . M.keys $ m ]

-- The max rank submatricies with independent rows, excluding some.
allMaxInd :: (Field a, Ord k, Ord l)
  => DoubleMap k l a -> [DoubleMap k l a]
allMaxInd m = filter isInd . allSizeNSub (rank $ toMatrix m) $ m

-- The max rank submatricies with independent rows,
-- including or excluding certain rows.
chooseInd :: (Field a, Ord k, Ord l)
  => [k] -> [k] -> DoubleMap k l a -> [DoubleMap k l a]
chooseInd fixedKeys free m = filter hasKeys . allMaxInd $ m'
  where
    hasKeys b = null $ fixedKeys \\ M.keys b
    m' = M.filterWithKey (\k _ -> k `notElem` free) m

-- Maximal sets of independent columns.
chooseCols :: (Field a, Ord k, Ord l) => DoubleMap k l a -> [DoubleMap k l a]
chooseCols = fmap trans . allMaxInd . trans

-- All submatricies with independent rows and columns, respecting row preferences.
chooseBoth :: (Field a, Ord k, Ord l)
  => [k] -> [k] -> DoubleMap k l a -> [DoubleMap k l a]
chooseBoth fixedKeys free m = concatMap chooseCols (chooseInd fixedKeys free m)

-- All legal constraints where extra terms to zero.
allExtraZero :: (Field a, Ord k, Num v, Ord l)
  => Map k v -> [k] -> DoubleMap k l a -> [Map k v]
allExtraZero fixed free = fmap (M.union fixed . fmap (const 0))
  . chooseInd (M.keys fixed) free

-- Solve Ax = B, where B is a column matrix.
innerSolver :: (Field a, Ord l, Ord k)
  => DoubleMap k l a -> Map k a -> Maybe (Map l a)
innerSolver b c
  | M.keys b == M.keys c = fromCol <$> linearSolve (toMatrix b) (toMatrix' c)
  | otherwise = Nothing
  where
    fromCol = M.fromList . zip (sort . innerKeys $ b) . concat . toLists

-- Takes a list of fixed values, a list of free values,
-- and a DoubleMap, and attempts to solve.
solve :: (Eq a, Field a, Ord k, Ord l)
  => Map k a -> [k] -> DoubleMap k l a -> [Map l a]
solve fixed free m = nub . catMaybes $
  [ innerSolver x y
  | x <- chooseBoth (M.keys fixed) free m
  , y <- allExtraZero fixed free m ]

-- Rows with a unique nonzero entry.
looseEnds :: (Eq a, Ord k, Num a) => DoubleMap k l a -> [k]
looseEnds m =
  [ k
  | k <- M.keys m
  , let rowK = M.findWithDefault M.empty k m
  , (== 1) . M.size . M.filter (/= 0) $ rowK ]

-- Tries setting loose ends to be free,
-- excluding those already declared to be fixed.
smartSolve :: (Field a, Ord k, Ord l, Ord a)
  => Map k a -> [k] -> DoubleMap k l a -> [Map l a]
smartSolve fixed free m = case filter (all (>= 0)) . solve fixed free' $ m of
  [] -> solve fixed free m
  xs -> xs
  where
    free' = union free $ looseEnds m \\ M.keys fixed


{-# LANGUAGE FlexibleContexts #-}
{-# language BangPatterns        #-}
{-# language LambdaCase          #-}
{-# options_ghc -Wno-unused-imports #-}
module Data.Vector.FFT (
  fft, ifft
  -- * Useful results
  , crossCorrelation
  ) where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad(..))

import Control.Monad.ST (runST)
import Data.Bits (shiftR,shiftL,(.&.),(.|.))
import Data.Bool (Bool,otherwise)
import Data.Complex (Complex(..),conjugate)
import Data.Foldable (forM_)

import Data.Vector.Unboxed as V (Vector, Unbox, map, zipWith, length, unsafeFreeze, (!))
import qualified Data.Vector.Unboxed.Mutable as VM (MVector, read, write, new, length)
import qualified Data.Vector.Generic as VG (Vector(..), copy)

import Prelude hiding (read)

{-# RULES
"fft/ifft" forall x. fft (ifft x) = x
"ifft/fft" forall x. ifft (fft x) = x
  #-}

-- | (Circular) cross-correlation of two vectors
--
-- Defined via the FFT and IFFT for computational efficiency
--
-- NB the source vectors should have matching length for meaningful results
crossCorrelation :: Vector (Complex Double)
                 -> Vector (Complex Double)
                 -> Vector (Complex Double)
crossCorrelation v1 v2 = ifft $ (cmap conjugate v1hat) `prod` v2hat
  where
    prod = V.zipWith (*)
    v1hat = fft v1
    v2hat = fft v2

-- | Radix-2 decimation-in-time fast Fourier Transform.
--
--   The given array (and therefore the output as well) is zero-padded to the next power of two if necessary.
fft :: Vector (Complex Double) -> Vector (Complex Double)
fft arr = runST $ do
  marr <- copyPadded arr
  mfft marr
  V.unsafeFreeze marr
{-# inlinable [1] fft #-}

-- | Inverse fast Fourier transform.
--
--   The given array (and therefore the output as well) is zero-padded to the next power of two if necessary.
ifft :: Vector (Complex Double) -> Vector (Complex Double)
ifft arr = do
  let lenComplex = intToComplexDouble (V.length arr)
  cmap ((/ lenComplex) . conjugate) . fft . cmap conjugate $ arr
{-# inlinable [1] ifft #-}

-- | Copy the source vector into a zero-padded mutable one
copyPadded :: (PrimMonad m, Num a, Unbox a) =>
              Vector a -> m (VM.MVector (PrimState m) a)
copyPadded arr = do
  let
    len = V.length arr
    l2 = nextPow2 len
  marr <- VM.new l2
  forM_ [0 .. l2 - 1] $ \i -> do
    let x | i < len = arr V.! i
          | otherwise = 0
    VM.write marr i x
  pure marr
{-# inline copyPadded #-}

-- | Radix-2 decimation-in-time fast Fourier Transform.
--   The given array must have a length that is a power of two,
--   though this property is not checked.
mfft :: (PrimMonad m) => VM.MVector (PrimState m) (Complex Double) -> m ()
mfft mut = do {
    let len = VM.length mut
  ; let bitReverse !i !j = do {
          ; if i == len - 1
              then stage 0 1
              else do {
                  when (i < j) $ swap mut i j
                ; let inner k l = if k <= l
                        then inner (k `shiftR` 1) (l - k)
                        else bitReverse (i + 1) (l + k)
                ; inner (len `shiftR` 1) j
              }
        }
        stage l l1 = if l == (log2 len)
          then pure ()
          else do {
              let !l2 = l1 `shiftL` 1
                  !e = (negate twoPi) / (intToDouble l2)
                  flight j !a = if j == l1
                    then stage (l + 1) l2
                    else do {
                        let butterfly i = if i >= len
                              then flight (j + 1) (a + e)
                              else do {
                                  let i1 = i + l1
                                ; xi1 :+ yi1 <- VM.read mut i1
                                ; let !co = cos a
                                      !si = sin a
                                      d = (co * xi1 - si * yi1) :+ (si * xi1 + co * yi1)
                                ; ci <- VM.read mut i
                                ; VM.write mut i1 (ci - d)
                                ; VM.write mut i (ci + d)
                                ; butterfly (i + l2)
                              }
                      ; butterfly j
                    }
            ; flight 0 0
         }
  ; bitReverse 0 0
}

-- wildcard cases should never happen. if they do, really bad things will happen.
b,s :: Int -> Int
b = \case { 0 -> 0x02; 1 -> 0x0c; 2 -> 0xf0; 3 -> 0xff00; 4 -> wordToInt 0xffff0000; 5 -> wordToInt 0xffffffff00000000; _ -> 0; }
s = \case { 0 -> 1; 1 -> 2; 2 -> 4; 3 -> 8; 4 -> 16; 5 -> 32; _ -> 0; }
{-# inline b #-}
{-# inline s #-}

-- | Next power of 2
nextPow2 :: Int -> Int
nextPow2 n
  | mod n 2 == 0 = n
  | otherwise = (2 :: Int) ^ (log2 n + 1)

log2 :: Int -> Int
log2 v0 = if v0 <= 0
  then error $ "Data.Vector.FFT: nonpositive input, got " ++ show v0
  else go 5 0 v0
  where
    go !i !r !v
      | i == -1 = r
      | v .&. b i /= 0 =
          let si = s i
          in go (i - 1) (r .|. si) (v `shiftR` si)
      | otherwise = go (i - 1) r v

{-# inline swap #-}
swap :: (PrimMonad m, Unbox a) =>
        VM.MVector (PrimState m) a -> Int -> Int -> m ()
swap mut i j = do
  atI <- VM.read mut i
  atJ <- VM.read mut j
  VM.write mut i atJ
  VM.write mut j atI

twoPi :: Double
{-# inline twoPi #-}
twoPi = 6.283185307179586

intToDouble :: Int -> Double
{-# inline intToDouble #-}
intToDouble = fromIntegral

wordToInt :: Word -> Int
{-# inline wordToInt #-}
wordToInt = fromIntegral

intToComplexDouble :: Int -> Complex Double
{-# inline intToComplexDouble #-}
intToComplexDouble = fromIntegral

{-# inline cmap #-}
cmap :: (Floating a, Unbox a) => (Complex a -> Complex a) -> V.Vector (Complex a) -> V.Vector (Complex a)
cmap = V.map

--

-- {-# inline copyWhole #-}
-- copyWhole :: (PrimMonad m, VG.Vector Vector a, Unbox a) => V.Vector a -> m (VM.MVector (PrimState m) a)
-- copyWhole arr = do
--   let len = V.length arr
--   marr <- VM.new len
--   VG.copy marr arr
--   pure marr

-- {-# inline arrOK #-}
-- arrOK :: Unbox a => Vector a -> Bool
-- arrOK arr =
--   let n = V.length arr
--   in (1 `shiftL` log2 n) == n


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
module AI.Singularity.Data.Network where
import GHC.TypeLits
import qualified Numeric.LinearAlgebra.Data as LD
import qualified Numeric.LinearAlgebra.Devel as LDev
import qualified Numeric.LinearAlgebra as LA
import qualified Control.Monad.State as S
import Control.Lens
import Data.Proxy
import Data.List (foldl')
import Control.Monad.ST
import System.Random
import Control.Monad
import Control.Arrow
import qualified Data.Vector.Storable as VS

import AI.Singularity.Utils
import AI.Singularity.Data.Matrix
import AI.Singularity.Data.Vector

data GeneralNetwork vec mat inp out where
  FFLayer :: forall vec mat m n. (KnownNat m, KnownNat n) =>
    mat (m+1) n ->
    (forall a. Floating a => a -> a) ->
    GeneralNetwork vec mat (vec m) (vec n)
  FFSeq :: forall vec mat a b c.
    GeneralNetwork  vec mat a b ->
    GeneralNetwork  vec mat b c ->
    GeneralNetwork  vec mat a c
  FromFunc :: (Num a) =>
    (forall b. Num b => b -> b) ->
    GeneralNetwork vec mat a a
  FromBinFunc :: forall a vec mat. Fractional a =>
    (forall b. Fractional b => b -> b -> b) ->
    GeneralNetwork vec mat (a,a) a
  FFDiv :: forall vec mat a' a b b'.
    GeneralNetwork vec mat a b ->
    GeneralNetwork vec mat a' b' ->
    GeneralNetwork vec mat (a, a') (b, b')
  Split :: forall vec mat a b c. Num a =>
    GeneralNetwork vec mat a b ->
    GeneralNetwork vec mat a c ->
    GeneralNetwork vec mat a (b,c)
  SplitNet :: forall vec mat n k l c. (KnownNat n, KnownNat k, KnownNat l, (k + l) ~ n) =>
    GeneralNetwork vec mat (vec n) c ->
    GeneralNetwork vec mat (vec k, vec l) c
  Recurse :: forall vec mat a b. (Num b) =>
    GeneralNetwork vec mat (a,b) b ->
    GeneralNetwork vec mat [a] b
  First :: forall vec mat a b c. (Num c) =>
    GeneralNetwork vec mat a b ->
    GeneralNetwork vec mat (a,c) b
  Second :: forall vec mat a b c. (Num c) =>
    GeneralNetwork vec mat a b ->
    GeneralNetwork vec mat (c,a) b

type Network inp out = GeneralNetwork Vector Matrix inp out

type TrainedNetwork inp out = GeneralNetwork Vector (GeneralInTrainingMat LD.Matrix) inp out

type STNetwork s = GeneralNetwork Vector (STInTrainMat s)

changeMatType :: Applicative f => (forall m n. (KnownNat m, KnownNat n) => mat m n -> f (mat' m n)) ->
                                  GeneralNetwork vec mat inp out -> f (GeneralNetwork vec mat' inp out)
{-# INLINE changeMatType #-}
changeMatType conv (FFLayer m f)    = (\x -> FFLayer x f) <$> conv m
changeMatType conv (FFSeq n1 n2)    = FFSeq <$> changeMatType conv n1 <*> changeMatType conv n2
changeMatType conv (FFDiv n1 n2)    = FFDiv <$> changeMatType conv n1 <*> changeMatType conv n2
changeMatType conv (FromFunc f)     = pure $ FromFunc f
changeMatType conv (FromBinFunc f)  = pure $ FromBinFunc f
changeMatType conv (Split n1 n2)    = Split <$> changeMatType conv n1 <*> changeMatType conv n2
changeMatType conv (SplitNet n1)    = SplitNet <$> changeMatType conv n1
changeMatType conv (Recurse n)      = Recurse <$> changeMatType conv n
changeMatType conv (First n)        = First <$> changeMatType conv n
changeMatType conv (Second n)       = Second <$> changeMatType conv n

matToTrain :: (KnownNat m, KnownNat n) => Matrix m n -> Identity (InTrainMat m n)
matToTrain = pure . (\x -> GeneralInTrainingMat x 0 0)

saveMat :: Matrix m n -> S.State [LD.Matrix Float] (Matrix m n)
saveMat (Mat m) = do
  l <- S.get
  S.put (m:l)
  return (Mat m)

saveWeights = map LD.toLists . reverse . flip S.execState [] . changeMatType saveMat

loadMat :: Matrix m n -> S.State [LD.Matrix Float] (Matrix m n)
loadMat _ = do
  (a:l) <- S.get
  S.put l
  return (Mat a)

loadWeights w = flip S.evalState (map LD.fromLists w) . changeMatType loadMat

toTrain :: Network inp out -> TrainedNetwork inp out
toTrain = runIdentity . changeMatType matToTrain

fromTrain :: TrainedNetwork inp out -> Network inp out
fromTrain = runIdentity . changeMatType (Identity . view trMatr)

toSTTrain :: Network inp out -> ST s (STNetwork s inp out)
toSTTrain = changeMatType toSTInTrainM

instance Show (Network inp out) where
  show (FFLayer m _)     = "weights:\n" ++ show m
  show (FFSeq net1 net2) = show net1 ++ "\n" ++ show net2
  show (FromFunc _ )     = "FromFunc"
  show (FFDiv a b)       = show a ++ "\n\tdiv \n" ++ show b
  show (SplitNet n)      = "\tSplitNet \n" ++ show n
  show (Recurse n)       = "\tRec\n" ++ show n

createLayer :: forall m n. (KnownNat m, KnownNat n) => Matrix (m+1) n -> (forall a. Floating a => a -> a)  -> Network (Vector m) (Vector n)
createLayer weights activator = FFLayer weights activator

randomLayer :: forall m n k. (KnownNat m, KnownNat n, KnownNat k, k ~ (m + 1)) => (forall a. Floating a => a -> a) -> IO (Network (Vector m) (Vector n))
randomLayer activator = (\w -> createLayer w activator) <$> randomMat

identLayer :: Num a => Network a a
identLayer = FromFunc id

generalConductSignal :: forall mat a b. (forall m n. mat m n -> Matrix m n) -> a -> GeneralNetwork Vector mat a b -> b
generalConductSignal _ v (FromFunc f)        = f v
generalConductSignal _ v (FromBinFunc f)     = uncurry f v
generalConductSignal c v (FFSeq net1 net2)   = generalConductSignal c (generalConductSignal c v net1) net2
generalConductSignal c v (FFLayer weights f) = cmap f (c weights #> consV 1 v)
generalConductSignal c v (FFDiv n1 n2)       = (generalConductSignal c (fst v) n1, generalConductSignal c (snd v) n2)
generalConductSignal c v (Split n1 n2)       = (generalConductSignal c v n1, generalConductSignal c v n2)
generalConductSignal c v (SplitNet n)        = generalConductSignal c (uncurry appendV v) n
generalConductSignal c v (First n)           = generalConductSignal c (fst v) n
generalConductSignal c v (Second n)          = generalConductSignal c (snd v) n
generalConductSignal c v (Recurse n)         = foldl' (\b a -> generalConductSignal c (a,b) n) 0 v

conductSignalInTrain = generalConductSignal (view trMatr)

conductSignal :: forall a b. a -> Network a b -> b
conductSignal = generalConductSignal id

runNetwork :: Network a b -> a -> b
runNetwork = flip conductSignal

compose :: Network m n -> Network k m -> Network k n
compose net2 net1 = FFSeq net1 net2

sigmoid :: Floating a => a -> a
sigmoid = (1/) . (1+). exp . negate

softmax :: forall n. KnownNat n => Vector n -> Vector n
softmax = exp >>> id &&& (foldV (+) 0 >>> fromValV) >>> uncurry (/)

atFirst :: Num a => Network b c -> Network (b,a) (c,a)
atFirst = flip FFDiv identLayer
atSecond :: Num a => Network b c -> Network (a,b) (a,c)
atSecond = FFDiv identLayer

takeFst :: (Num c, Num a) => Network (a,c) a
takeFst = First identLayer
takeSnd :: (Num c, Num a) => Network (c,a) a
takeSnd = Second identLayer

prelstm :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) => IO (Network [Vector m] (Vector n, Vector k))
prelstm = do
  let takeInp  = First identLayer
      takeOut  = Second . Second $ identLayer
      -- takeMem  = Second . First $ identLayer
      splitter = Split takeInp takeOut
  x <- SplitNet <$> randomLayer sigmoid
  y <- SplitNet <$> randomLayer sigmoid
  z <- SplitNet <$> randomLayer tanh
  k <- SplitNet <$> randomLayer sigmoid
  l <- randomLayer tanh
  let
    gateCell   = Split (FFSeq splitter x) (Second . First $ identLayer) `FFSeq` FromBinFunc (*)
    secondCell = Split gateCell (splitter `FFSeq` Split y z `FFSeq` FromBinFunc (*)) `FFSeq` FromBinFunc (+)
    thirdCell  = Split secondCell (splitter `FFSeq` k) `FFSeq` Split takeFst (atFirst l `FFSeq` FromBinFunc (*))
      -- get2 pre = Split pre (Split)
  return . Recurse $ thirdCell
-- MnistReader

gru :: forall m n. (KnownNat m, KnownNat n) => IO (Network [Vector m] (Vector n))
gru = do
  z <- SplitNet <$> randomLayer sigmoid
  r <- SplitNet <$> randomLayer sigmoid
  h <- SplitNet <$> randomLayer tanh
  let
    rxt = Split r takeSnd `FFSeq` FromBinFunc (*)
    ht' = Split rxt takeFst `FFSeq` h
    ht  = Split (z `FFSeq` Split (FromFunc (1-)) identLayer) (Split takeSnd ht') `FFSeq` FromBinFunc (*) `FFSeq` FromBinFunc (+)
  return $ Recurse ht
    --beg = Split (Split x takeFst `FFSeq` FromBinFunc (*)) takeFst `FFSeq` z

iden :: forall a. (Num a) => a -> a
iden x = x


{-# LANGUAGE BangPatterns, TypeFamilies, TypeOperators, FlexibleInstances, FlexibleContexts, GADTs #-}
{-# LANGUAGE MultiParamTypeClasses, UndecidableInstances #-}
module Data.NeuralNetwork.Backend.HMatrix.Layers where

import Numeric.LinearAlgebra hiding (R, C)
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Storable.Mutable as SVM
import System.Random.MWC
import System.Random.MWC.Distributions
import Control.Monad.ST
import Control.Monad (liftM2, forM_, when)
import GHC.Float
import Data.STRef
import Data.Functor.Identity
import Control.DeepSeq
import Data.NeuralNetwork

type R = Float

-- We parameterise the activation layer T, where the parameter indicates how
-- elements are contained:
--   SinglC :. Vector, SinglC :. Matrix, MultiC :. Vector, MultiC :.  Matrix
-- SinglC means the input has only one channel, while
-- MultiC means the input has more than one.
--
-- type function composition
data (f :: * -> *) :. (g :: * -> *) :: * -> *
-- type function: Identity
data SinglC :: * -> *
data MultiC :: * -> *

-- Tags for each form of layer
data F
data C
data A
data M
data T (c :: * -> *)
data S a b

data RunLayer :: * -> * where
  -- Densely connected layer
  -- input:   vector of size m
  -- output:  vector of size n
  -- weights: matrix of size m x n
  -- biases:  vector of size n
  Full :: !(Matrix R) -> !(Vector R) -> RunLayer F
  -- convolutional layer
  -- input:  channels of 2D floats, of the same size (a x b), # of input channels:  m
  -- output: channels of 2D floats, of the same size (c x d), # of output channels: n
  --         where c = a + 2*padding + 1 - s
  --               d = b + 2*padding + 1 - t
  -- feature:  matrix of (s x t), # of features: m x n
  -- padding:  number of 0s padded at each side of channel
  -- biases:   bias for each output, # of biases: n
  Conv  :: !(V.Vector (V.Vector (Matrix R))) -> !(V.Vector R) -> Int -> RunLayer C
  -- Reshape from channels of matrix to a single vector
  -- input:  m channels of 2D matrices
  --         assuming that all matrices are of the same size a x b
  -- output: 1D vector of the concatenation of all input channels
  --         its size: m x a x b
  As1D  :: RunLayer A
  -- max pooling layer
  -- input:  channels of 2D floats, of the same size (a x b), # of input channels:  m
  --         assuming that a and b are both multiple of stride
  -- output: channels of 2D floats, of the same size (c x d), # of output channels: m
  --         where c = a / stride
  --               d = b / stride
  MaxP :: Int -> RunLayer M
  -- Activator
  -- the input can be either a 1D vector, 2D matrix, or channels of either.
  Activation :: (R->R, R->R) -> RunLayer (T c)
  -- stacking two components a and b
  -- the output of a should matches the input of b
  Stack :: !(RunLayer a) -> !(RunLayer b) -> RunLayer (S a b)

instance Component (RunLayer F) where
    type Run (RunLayer F) = Identity
    type Inp (RunLayer F) = Vector R
    type Out (RunLayer F) = Vector R
    -- trace is (input, weighted-sum)
    newtype Trace (RunLayer F) = DTrace (Vector R, Vector R)
    forwardT (Full w b) !inp =
        let !bv = (inp <# w) `add` b
        in return $ DTrace (inp,bv)
    output (DTrace (_,!a)) = a
    backward l (DTrace (!iv,!bv)) !odelta rate =
        let Full w b = l
            !d = scale (negate rate) odelta
            !m = iv `outer` d
            -- back-propagated error at input
            !idelta = w #> odelta
            -- update to weights
            ---- for what reason, could this expression: w `add` (iv `outer` d)
            ---- entails a huge space leak? especially, neither 'seq' nor
            ---- 'deepseq' helps a bit. The only workaround is to expand the
            ---- add function, and call SV.force on the result vector, which
            ---- explcitly copy and drop reference to orignal computed result.
            !w'= w `add` m
            -- !w'= let (r,c) = size w
            --          dat1 = flatten (tr' w)
            --          dat2 = flatten (tr' m)
            --      in matrixFromVector ColumnMajor r c $ SV.force $ dat1 `add` dat2
            !b'= b `add` d
        in return $ (Full w' b', idelta)

instance Component (RunLayer C) where
    type Run (RunLayer C) = Identity
    type Inp (RunLayer C) = V.Vector (Matrix R)
    type Out (RunLayer C) = V.Vector (Matrix R)
    -- trace is (input, convoluted output)
    newtype Trace (RunLayer C) = CTrace (Inp (RunLayer C), V.Vector (Matrix R))
    forwardT (Conv fs bs p) !inp =
        let !ov = V.zipWith feature
                            (tr fs) -- feature matrix indexed majorly by each output
                            bs      -- biases by each output
        in return $ CTrace (inp,ov)
      where
        !osize = let (x,y) = size (V.head inp)
                     (u,v) = size (V.head $ V.head fs)
                 in (x+2*p-u+1, y+2*p-v+1)
        -- transpose the features matrix
        tr :: V.Vector (V.Vector a) -> V.Vector (V.Vector a)
        tr uv = let n   = V.length (V.head uv)
                    !vu = V.map (\i -> V.map (V.! i) uv) $ V.enumFromN 0 n
                in vu
        feature :: V.Vector (Matrix R) -> R -> Matrix R
        feature f b = V.foldl1' add (V.zipWith (layerCorr2 p) f inp) `add` konst b osize
    output (CTrace (_,a)) = a
    backward l (CTrace (!iv,!av)) !odelta rate =
      let Conv fs bs p = l
          -- update to the feature matrix
          m :: V.Vector (V.Vector (Matrix R))
          !m = V.zipWith (\flts chn ->
                            -- chn:  a single input channel
                            -- flts: all features used for chn
                            V.zipWith (\f d ->
                              let upd = scale (negate rate) (layerCorr2 p chn d)
                              in f `add` upd
                            ) flts odelta
                          ) fs iv
          -- update to the biases
          b :: V.Vector R
          !b = V.zipWith (\b d -> b + (negate rate) * sumElements d) bs odelta
          -- back-propagated error at input
          idelta :: V.Vector (Matrix R)
          !idelta = V.map (\f -> V.foldl1' add $ V.zipWith (layerConv2 p) f odelta) fs
      in --trace ("CL:" ++ show odelta)
         return $ (Conv m b p, idelta)

instance Component (RunLayer A) where
  type Run (RunLayer A) = Identity
  type Inp (RunLayer A) = V.Vector (Matrix R)
  type Out (RunLayer A) = Vector R
  -- trace keeps information of (m, axb, b, output)
  newtype Trace (RunLayer A) = ReshapeTrace (Int, Int, Int, Vector R)
  forwardT _ !inp =
    let !b = V.length inp
        (!r,!c) = size (V.head inp)
        !o = V.foldr' (\x y -> flatten x SV.++ y) SV.empty inp
    in return $ ReshapeTrace (b, r*c, c, o)
  output (ReshapeTrace (_,_,_,a)) = a
  backward a (ReshapeTrace (b,n,c,_)) !odelta _ =
    let !idelta = V.fromList $ map (reshape c) $ takesV (replicate b n) odelta
    in return $ (a, idelta)

instance Component (RunLayer M) where
  type Run (RunLayer M) = Identity
  type Inp (RunLayer M) = V.Vector (Matrix R)
  type Out (RunLayer M) = V.Vector (Matrix R)
  -- trace is (dimension of pools, index of max in each pool, pooled matrix)
  -- for each channel.
  newtype Trace (RunLayer M) = PTrace (V.Vector (IndexOf Matrix, Vector Int, Matrix R))
  -- forward is to divide the input matrix in stride x stride sub matrices,
  -- and then find the max element in each sub matrices.
  forwardT (MaxP stride) !inp = return $ PTrace $ V.map mk inp
    where
      mk inp = let (!i,!v) = pool stride inp in (size v, i, v)
  output (PTrace a) = V.map (\(_,_,!o) ->o) a
  -- use the saved index-of-max in each pool to propagate the error.
  backward l@(MaxP stride) (PTrace t) odelta _ =
      let !idelta = V.zipWith gen t odelta in return $ (l, idelta)
    where
      gen (!si,!iv,_) od = unpool stride iv od

instance (Component (RunLayer a),
          Component (RunLayer b),
          Run (RunLayer a) ~ Identity,
          Run (RunLayer b) ~ Identity,
          Out (RunLayer a) ~ Inp (RunLayer b)
         ) => Component (RunLayer (S a b)) where
    type Run (RunLayer (S a b)) = Identity
    type Inp (RunLayer (S a b)) = Inp (RunLayer a)
    type Out (RunLayer (S a b)) = Out (RunLayer b)
    newtype Trace (RunLayer (S a b)) = TTrace (Trace (RunLayer b), Trace (RunLayer a))
    forwardT (Stack a b) !i = do
        !tra <- forwardT a i
        !trb <- forwardT b (output tra)
        return $ TTrace (trb, tra)
    output (TTrace !a) = output (fst a)
    backward (Stack a b) (TTrace (!trb,!tra)) !odelta rate = do
        (b', !odelta') <- backward b trb odelta  rate
        (a', !idelta ) <- backward a tra odelta' rate
        return (Stack a' b', idelta)

instance (Container c R, Pairwise c R) => Component (RunLayer (T (MultiC :. c))) where
    type Run (RunLayer (T (MultiC :. c))) = Identity
    type Inp (RunLayer (T (MultiC :. c))) = V.Vector (c R)
    type Out (RunLayer (T (MultiC :. c))) = V.Vector (c R)
    newtype Trace (RunLayer (T (MultiC :. c))) = TTraceM (V.Vector (Trace (RunLayer (T (SinglC :. c)))))
    forwardT (Activation ac) !inp =
      TTraceM <$> V.mapM (forwardT (Activation ac)) inp
    output (TTraceM a) = V.map output a
    backward a@(Activation ac) (TTraceM ts) !odelta r = do
      idelta <- V.zipWithM (\t d -> snd <$> backward (Activation ac) t d r) ts odelta
      return (a, idelta)

instance (Container c R, Pairwise c R) => Component (RunLayer (T (SinglC :. c))) where
    type Run (RunLayer (T (SinglC :. c))) = Identity
    type Inp (RunLayer (T (SinglC :. c))) = c R
    type Out (RunLayer (T (SinglC :. c))) = c R
    newtype Trace (RunLayer (T (SinglC :. c))) = TTraceS (c R, c R)
    forwardT (Activation (af,_)) !inp = return $ TTraceS (inp, cmap af inp)
    output (TTraceS (_,!a)) = a
    backward a@(Activation (_,ag)) (TTraceS (!iv,_)) !odelta _ = return $ (a, odelta `hadamard` cmap ag iv)

newFLayer :: Int                -- number of input values
          -> Int                -- number of neurons (output values)
          -> IO (RunLayer F)    -- new layer
newFLayer m n =
    withSystemRandom . asGenIO $ \gen -> do
        -- we build the weights in column major because in the back-propagation
        -- algo, the computed update to weights is in column major. So it is
        -- good for performance to keep the matrix always in column major.
        w <- buildMatrix (normal 0 0.01 gen) ColumnMajor (m,n)
        b <- return $ konst 1 n
        return $ Full w b

newCLayer :: Int                -- number of input channels
          -> Int                -- number of output channels
          -> Int                -- size of each feature
          -> Int                -- size of padding
          -> IO (RunLayer C)    -- new layer
newCLayer inpsize outsize sfilter npadding =
  withSystemRandom . asGenIO $ \gen -> do
      fs <- V.replicateM inpsize $ V.replicateM outsize $
              buildMatrix (truncNormal 0 0.1 gen) RowMajor (sfilter, sfilter)
      bs <- return $ V.replicate outsize 0.1
      return $ Conv fs bs npadding
  where
    truncNormal m s g = do
      x <- standard g
      if x >= 2.0 || x <= -2.0
        then truncNormal m s g
        else return $! m + s * x

buildMatrix g order (nr, nc) = do
  vals <- SV.replicateM (nr*nc) (double2Float <$> g)
  return $ matrixFromVector order nr nc vals

layerCorr2 :: Int -> Matrix R -> Matrix R -> Matrix R
layerCorr2 p k m = corr2 k padded
  where
    padded = zeroPadded p m
    (w,_)  = size k

layerConv2 :: Int -> Matrix R -> Matrix R -> Matrix R
layerConv2 p k m = corr2 (fliprl . flipud $ k) padded
  where
    padded = zeroPadded p m

-- max pool, picking out the maximum element
-- in each stride x stride sub-matrices.
-- assuming that the original matrix row and column size are
-- both multiple of stride
pool :: Int -> Matrix Float -> (Vector Int, Matrix Float)
pool 1 mat = let (r,c) = size mat in (SV.replicate (r*c) 0, mat)
-- pool 2 mat | orderOf mat == RowMajor = c_max_pool2_f mat
pool stride mat = runST $ do
  ori <- unsafeThawMatrix mat
  mxv <- newUndefinedMatrix RowMajor r' c'
  mxi <- newUndefinedVector (r'*c')
  forM_ [0..r'-1] $ \i -> do
    forM_ [0..c'-1] $ \j -> do
      (n,v) <- unsafeMaxIndEle ori (i*stride) (j*stride) stride stride
      unsafeWriteVector mxi (i*c'+j) n
      unsafeWriteMatrix mxv i j v
  a <- unsafeFreezeVector mxi
  b <- unsafeFreezeMatrix mxv
  return (a,b)
  where
    (r,c) = size mat
    r'    = r `div` stride
    c'    = c `div` stride
    unsafeMaxIndEle mm x y r c = do
      mp <- newSTRef 0
      mv <- newSTRef (-10000.0)
      forM_ [0..r-1] $ \ i -> do
        forM_ [0..c-1] $ \ j -> do
          v1 <- unsafeReadMatrix mm (x+i) (y+j)
          v0 <- readSTRef mv
          when (v1 > v0) $ do
            writeSTRef mv v1
            writeSTRef mp (i*2+j)
      p <- readSTRef mp
      v <- readSTRef mv
      return (p, v)

-- the reverse of max pool.
-- assuming idx and mat are of the same size
unpool :: Int -> Vector Int -> Matrix Float -> Matrix Float
unpool stride idx mat = runSTMatrix $ do
  mat' <- newMatrix' 0 r' c'
  forM_ [0..r-1] $ \i -> do
    forM_ [0..c-1] $ \j -> do
      let pos     = idx SV.! (i*c+j)
      let (oi,oj) = pos `divMod` 2
      let val     = mat `atIndex` (i,j)
      unsafeWriteMatrix mat' (i*stride+oi) (j*stride+oj) val
  return mat'
  where
    (r,c) = size mat
    (r',c') = (r*stride, c*stride)

-- a slightly faster way to pading the matrix
-- camparing to fromBlocks provided by hmatrix.
zeroPadded :: Int -> Matrix Float -> Matrix Float
zeroPadded p mat = runSTMatrix $ do
  mat' <- newMatrix' 0 r' c'
  setMatrix mat' p p mat
  return mat'
  where
    (r,c) = size mat
    (r',c') = (r+2*p,c+2*p)

-- a slightly faster version of newMatrix, which based
-- directly on lower level Vector.Storable creation.
newMatrix' :: SVM.Storable t => t -> Int -> Int -> ST s (STMatrix s t)
newMatrix' v r c = do
  vec <- SVM.replicate (r*c) v
  vec <- SV.unsafeFreeze vec
  unsafeThawMatrix $ reshape c vec

-- The hmatrix does has a SIMD enabled pairwise product named 'mul' in the
-- 'Container' class, but unfortunately it is not exported at all.
-- We define a slow hadamard product here.
class Container c e => Pairwise c e where
  hadamard :: c e -> c e -> c e
instance Pairwise Vector R where
  hadamard = zipVectorWith (*)
instance Pairwise Matrix R where
  hadamard a b = reshape (cols a) $ hadamard (flatten a) (flatten b)


{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2014
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-----------------------------------------------------------------------------

module Numeric.AD.Newton
  (
  -- * Newton's Method (Forward AD)
    findZero
  , inverse
  , fixedPoint
  , extremum
  -- * Gradient Ascent/Descent (Reverse AD)
  , gradientDescent
  , gradientAscent
  , conjugateGradientDescent
  , conjugateGradientAscent
  ) where

import Data.Foldable (all, sum)
import Data.Reflection (Reifies)
import Data.Traversable
import Numeric.AD.Internal.Combinators
import Numeric.AD.Internal.Forward (Forward)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Or
import Numeric.AD.Internal.Reverse (Reverse, Tape)
import Numeric.AD.Internal.Type (AD(..))
import Numeric.AD.Mode
import Numeric.AD.Mode.Reverse as Reverse (gradWith')
import Numeric.AD.Rank1.Kahn as Kahn (Kahn, grad)
import qualified Numeric.AD.Rank1.Newton as Rank1
import Prelude hiding (all, mapM, sum)

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,2.5,2.05,2.000609756097561,2.0000000929222947,2.000000000000002,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a]
findZero f = Rank1.findZero (runAD.f.AD)
{-# INLINE findZero #-}

-- | The 'inverse' function inverts a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- Example:
--
-- >>> last $ take 10 $ inverse sqrt 1 (sqrt 10)
-- 10.0
inverse :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a -> [a]
inverse f = Rank1.inverse (runAD.f.AD)
{-# INLINE inverse  #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Newton's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a]
fixedPoint f = Rank1.fixedPoint (runAD.f.AD)
{-# INLINE fixedPoint #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Newton's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream
-- becomes constant ("it converges"), no further elements are returned.
--
-- >>> last $ take 10 $ extremum cos 1
-- 0.0
extremum :: (Fractional a, Eq a) => (forall s. AD s (On (Forward (Forward a))) -> AD s (On (Forward (Forward a)))) -> a -> [a]
extremum f = Rank1.extremum (runAD.f.AD)
{-# INLINE extremum #-}

-- | The 'gradientDescent' function performs a multivariate
-- optimization, based on the naive-gradient-descent in the file
-- @stalingrad\/examples\/flow-tests\/pre-saddle-1a.vlad@ from the
-- VLAD compiler Stalingrad sources.  Its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- It uses reverse mode automatic differentiation to compute the gradient.
gradientDescent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> [f a]
gradientDescent f x0 = go x0 fx0 xgx0 0.1 (0 :: Int)
  where
    (fx0, xgx0) = Reverse.gradWith' (,) f x0
    go x fx xgx !eta !i
      | eta == 0     = [] -- step size is 0
      | fx1 > fx     = go x fx xgx (eta/2) 0 -- we stepped too far
      | zeroGrad xgx = [] -- gradient is 0
      | otherwise    = x1 : if i == 10
                            then go x1 fx1 xgx1 (eta*2) 0
                            else go x1 fx1 xgx1 eta (i+1)
      where
        zeroGrad = all (\(_,g) -> g == 0)
        x1 = fmap (\(xi,gxi) -> xi - eta * gxi) xgx
        (fx1, xgx1) = Reverse.gradWith' (,) f x1
{-# INLINE gradientDescent #-}

-- | Perform a gradient descent using reverse mode automatic differentiation to compute the gradient.
gradientAscent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> [f a]
gradientAscent f = gradientDescent (negate . f)
{-# INLINE gradientAscent #-}

-- | Perform a conjugate gradient descent using reverse mode automatic differentiation to compute the gradient, and using forward-on-forward mode for computing extrema.
--
-- >>> let sq x = x * x
-- >>> let rosenbrock [x,y] = sq (1 - x) + 100 * sq (y - sq x)
-- >>> rosenbrock [0,0]
-- 1
-- >>> rosenbrock (conjugateGradientDescent rosenbrock [0, 0] !! 5) < 0.1
-- True
conjugateGradientDescent
  :: (Traversable f, Ord a, Fractional a)
  => (forall s. Chosen s => f (Or s (On (Forward (Forward a))) (Kahn a)) -> Or s (On (Forward (Forward a))) (Kahn a))
  -> f a -> [f a]
conjugateGradientDescent f = conjugateGradientAscent (negate . f)
{-# INLINE conjugateGradientDescent #-}

lfu :: Functor f => (f (Or F a b) -> Or F a b) -> f a -> a
lfu f = runL . f . fmap L

rfu :: Functor f => (f (Or T a b) -> Or T a b) -> f b -> b
rfu f = runR . f . fmap R

-- | Perform a conjugate gradient ascent using reverse mode automatic differentiation to compute the gradient.
conjugateGradientAscent
  :: (Traversable f, Ord a, Fractional a)
  => (forall s. Chosen s => f (Or s (On (Forward (Forward a))) (Kahn a)) -> Or s (On (Forward (Forward a))) (Kahn a))
  -> f a -> [f a]
conjugateGradientAscent f x0 = takeWhile (all (\a -> a == a)) (go x0 d0 d0 delta0)
  where
    dot x y = sum $ zipWithT (*) x y
    d0 = Kahn.grad (rfu f) x0
    delta0 = dot d0 d0
    go xi _ri di deltai = xi : go xi1 ri1 di1 deltai1
      where
        ai = last $ take 20 $ Rank1.extremum (\a -> lfu f $ zipWithT (\x d -> auto x + a * auto d) xi di) 0
        xi1 = zipWithT (\x d -> x + ai*d) xi di
        ri1 = Kahn.grad (rfu f) xi1
        deltai1 = dot ri1 ri1
        bi1 = deltai1 / deltai
        di1 = zipWithT (\r d -> r + bi1 * d) ri1 di
{-# INLINE conjugateGradientAscent #-}


module XAlg.Languages.Arithmetic.Expr (
-- * Core data types
  Expr (..),
  Expr_ (..),

-- * Parsing
  lit,
  term,
  expr,

-- * Evaluation
  eval

) where

import XAlg.Languages.Arithmetic.Number
import XAlg.Foundation.Recursion
import XAlg.Foundation.Core
import XAlg.Foundation.Parsing

import Data.Set
import Data.Data
import qualified Data.Map as M

import Text.Megaparsec
import Text.Megaparsec.String
import Text.Megaparsec.Expr
import Text.Megaparsec.Lexer (float)
import qualified Data.Complex as C

-- | Arithmetical expression.
data Expr = Lit Number
          | Var String
          | Plus Expr Expr
          | Minus Expr Expr
          | Mult Expr Expr
          | Div Expr Expr
          | Power Expr Expr
          | Neg Expr
          deriving (Show, Eq, Ord, Data)

infixl 6 `Plus`
infixl 6 `Minus`
infixl 8 `Mult`
infixl 8 `Div`
infixl 9 `Power`

instance Num Expr where
  (*) = Mult
  (+) = Plus
  negate = Neg
  abs = undefined
  signum = undefined
  fromInteger = Lit . Int

instance RecursiveData Expr where
  nodeEquality (Lit x)     (Lit y)     = x == x
  nodeEquality (Var x)     (Var y)     = x == x
  nodeEquality (Plus _ _)  (Plus _ _)  = True
  nodeEquality (Minus _ _) (Minus _ _) = True
  nodeEquality (Mult _ _)  (Mult _ _)  = True
  nodeEquality (Div _ _)   (Div _ _)   = True
  nodeEquality (Power _ _) (Power _ _) = True
  nodeEquality (Neg _)     (Neg _)     = True
  nodeEquality _           _           = False

  children (Lit x)       = []
  children (Var x)       = []
  children (Plus x1 x2)  = [x1,x2]
  children (Minus x1 x2) = [x1,x2]
  children (Mult x1 x2)  = [x1,x2]
  children (Div x1 x2)   = [x1,x2]
  children (Power x1 x2) = [x1,x2]
  children (Neg x)       = [x]

  mapChildren f (Lit x)       = Lit x
  mapChildren f (Var x)       = Var x
  mapChildren f (Plus x1 x2)  = recursively f x1 `Plus` recursively f x2
  mapChildren f (Minus x1 x2) = recursively f x1 `Minus` recursively f x2
  mapChildren f (Mult x1 x2)  = recursively f x1 `Mult` recursively f x2
  mapChildren f (Div x1 x2)   = recursively f x1 `Div` recursively f x2
  mapChildren f (Power x1 x2) = recursively f x1 `Power` recursively f x2
  mapChildren f (Neg x)       = Neg (recursively f x)

instance SymbolicLang Expr Number where
  variable = Var
  isVariable (Var v) = Just v
  isVariable _       = Nothing
  literal = Lit
  isLiteral (Lit x) = Just x
  isLiteral _       = Nothing
  evaluate = eval

-----------------------------------------------------------------------------
--    Functorial data type
-----------------------------------------------------------------------------

data Expr_ a = Lit_ Number
             | Var_ String
             | Plus_ a a
             | Minus_ a a
             | Mult_ a a
             | Div_ a a
             | Power_ a a
             | Neg_ a
             deriving (Functor, Foldable, Traversable, Data)

instance Fixable Expr_ Expr where
  foldFix (Lit_ x) = Lit x
  foldFix (Var_ x) = Var x
  foldFix (Plus_ x1 x2) = Plus x1 x2
  foldFix (Minus_ x1 x2) = Minus x1 x2
  foldFix (Mult_ x1 x2) = Mult x1 x2
  foldFix (Div_ x1 x2) = Div x1 x2
  foldFix (Power_ x1 x2) = Power x1 x2
  foldFix (Neg_ x) = Neg x

  unfoldFix (Lit x) = Lit_ x
  unfoldFix (Var x) = Var_ x
  unfoldFix (Plus x1 x2) = Plus_ x1 x2
  unfoldFix (Minus x1 x2) = Minus_ x1 x2
  unfoldFix (Mult x1 x2) = Mult_ x1 x2
  unfoldFix (Div x1 x2) = Div_ x1 x2
  unfoldFix (Power x1 x2) = Power_ x1 x2
  unfoldFix (Neg x) = Neg_ x

instance ParaSymb Expr_ where
  isVariable_ (Var_ v) = Just v
  isVariable_ _        = Nothing
  variable_ = Var_

deriving instance Show a => Show (Expr_ a)
deriving instance Show (Fix Expr_)

deriving instance Eq a => Eq (Expr_ a)
deriving instance Eq (Fix Expr_)

-----------------------------------------------------------------------------
--    Parsing
-----------------------------------------------------------------------------

imaginaryFloat :: Parser (C.Complex Double)
imaginaryFloat = fmap (0 C.:+)
  $ char 'i' *> float <|> float <* char 'i'

imaginaryInt :: Parser (C.Complex Double)
imaginaryInt = fmap ((0 C.:+) . fromInteger)
  $ char 'i' *> integ <|> integ <* char 'i'

imaginary :: Parser (C.Complex Double)
imaginary = imaginaryInt <|> imaginaryFloat <?> "imaginary number"

-- | Parse a numeric literal.
lit :: Parser Number
lit = try (Complex <$> imaginary)
  <|> try (Real <$> float)
  <|> Int <$> integ

-- | Parser an arithmetic expression term.
term :: Parser Expr
term = parens expr <|> var <|> Lit <$> lit <?> "term"

-- | Parse an arithmetic expression.
expr :: Parser Expr
expr = makeExprParser term table <?> "expression"

-- Must be a better way to get spaces:
table :: [[Operator Parser Expr]]
table = [ prefix "-" Neg
        , binary "^" Power
        , binary "*" Mult ++ binary "/" Div
        , binary "+" Plus ++ binary "-" Minus
        ]

instance Parsable Expr where
  exprParser = expr

-----------------------------------------------------------------------------
--    Pretty printing
-----------------------------------------------------------------------------

-- Only Minus is special, second argument has higher precedence.
showNodes :: Algebra Expr_ (Int, String)
showNodes (Lit_ x)       = (5, show x)
showNodes (Var_ x)       = (5, x)
showNodes (Plus_ x1 x2)  = (1, applyParens 1 x1 ++ " + " ++ applyParens 1 x2)
showNodes (Mult_ x1 x2)  = (3, applyParens 3 x1 ++ "*" ++ applyParens 3 x2)
showNodes (Neg_ x)       = (5, "-" ++ applyParens 5 x)
showNodes (Minus_ x1 x2) = (1, applyParens 1 x1 ++ " - " ++ applyParens 2 x2)
showNodes (Div_ x1 x2)   = (3, applyParens 3 x1 ++ "/" ++ applyParens 3 x2)
showNodes (Power_ x1 x2) = (4, applyParens 4 x1 ++ "^" ++ applyParens 4 x2)

-- target prec, (node prec, shown term) -> apply parens if necessary
applyParens :: Int -> (Int,String) -> String
applyParens t (n,str) = if n < t then parenthesize str else str

parenthesize :: String -> String
parenthesize s = "(" ++ s ++ ")"

instance Renderable Expr where
  render = snd . catamap showNodes

instance Transcribe Expr

instance XAlgLang Expr Expr_ Number

-------

countNode :: Algebra Expr_ Int
countNode (Lit_ x)       = 1
countNode (Var_ x)       = 1
countNode (Plus_ x1 x2)  = x1 + x2
countNode (Mult_ x1 x2)  = x1 + x2
countNode (Minus_ x1 x2) = x1 + x2
countNode (Div_ x1 x2)   = x1 + x2
countNode (Power_ x1 x2) = x1 + x2
countNode (Neg_ x)       = succ x

countNodes :: Expr -> Int
countNodes = catamap countNode

-------

data Mode = IntMode
          | RealMode
          | ComplexMode deriving (Eq,Ord)

instance Show Mode where
  show IntMode     = "ℤ"
  show RealMode    = "ℝ"
  show ComplexMode = "ℂ"

litMode :: Number -> Mode
litMode (Real _)   = RealMode
litMode (Int   _)   = IntMode
litMode (Complex _) = ComplexMode

modeNodes :: Algebra Expr_ Mode
modeNodes (Lit_ x)       = litMode x
modeNodes (Var_ x)       = IntMode
modeNodes (Plus_ x1 x2)  = x1 `max` x2
modeNodes (Minus_ x1 x2) = x1 `max` x2
modeNodes (Mult_ x1 x2)  = x1 `max` x2
modeNodes (Div_ x1 x2)   = x1 `max` x2
modeNodes (Power_ x1 x2) = x1 `max` x2
modeNodes (Neg_ x)       = x

---------

evalNode :: Algebra Expr_ (EvalResult Expr Number)
evalNode (Lit_ x)       = ValueResult x
evalNode (Var_ x)       = FunctionResult (singleton x) (Var x)
evalNode (Plus_ x1 x2)  = combineWith Plus (+) x1 x2
evalNode (Minus_ x1 x2) = combineWith Minus (-) x1 x2
evalNode (Mult_ x1 x2)  = combineWith Mult (*) x1 x2
evalNode (Div_ x1 x2)   = combineWith Div (/) x1 x2
evalNode (Power_ x1 x2) = combineWith Power (**) x1 x2
evalNode (Neg_ x)       = applyWith Neg negate x

evalNodeCtx :: Env Expr -> Algebra Expr_ (EvalResult Expr Number)
evalNodeCtx m (Var_ s) = case M.lookup s m of
  Just x  -> catamap (evalNodeCtx m) x
  Nothing -> FunctionResult (singleton s) (Var s)
evalNodeCtx _ x = evalNode x

eval :: Expr -> EvalResult Expr Number
eval = catamap evalNode

-- helpers

type BinaryOp a = a -> a -> a
type UnaryOp a = a -> a

combineWith :: SymbolicLang x s => BinaryOp x -> BinaryOp s -> BinaryOp (EvalResult x s)
combineWith f g (FunctionResult vs x) (FunctionResult ws y) = FunctionResult (union vs ws) (f x y)
combineWith f g (FunctionResult vs x) (ValueResult t)       = FunctionResult vs $ f x $ literal t
combineWith f g (ValueResult s)       (FunctionResult ws y) = FunctionResult ws $ f (literal s) y
combineWith f g (ValueResult s)       (ValueResult t)       = ValueResult $ g s t

applyWith :: SymbolicLang x s => UnaryOp x -> UnaryOp s -> UnaryOp (EvalResult x s)
applyWith f g (FunctionResult s x) = FunctionResult s (f x)
applyWith f g (ValueResult r)      = ValueResult (g r)


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Trustworthy #-}

{- |
Module      :  Physics.Learn.Schrodinger1D
Copyright   :  (c) Scott N. Walck 2015-2018
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains functions to
solve the (time dependent) Schrodinger equation
in one spatial dimension for a given potential function.
-}

module Physics.Learn.Schrodinger1D
    (
    -- * Potentials
      freeV
    , harmonicV
    , squareWell
    , doubleWell
    , stepV
    , wall
    -- * Initial wavefunctions
    --    , harm
    , coherent
    , gaussian
    , movingGaussian
    -- * Utilities
    , stateVectorFromWavefunction
    , hamiltonianMatrix
    , expectX
    , picture
    , xRange
    , listForm
    )
    where

import Data.Complex
    ( Complex(..)
    , magnitude
    )
import Graphics.Gloss
    ( Picture(..)
    , yellow
    , black
    , Display(..)
    , display
    )
-- import Math.Polynomial.Hermite
--     ( evalPhysHermite
--     )
import Numeric.LinearAlgebra
    ( R
    , C
    , Vector
    , Matrix
    , (|>)
    , (<.>)
    , fromLists
    , toList
    , size
    )
import Physics.Learn.QuantumMat
    ( probVector
    , timeEv
    )

--i :: Complex Double
--i = 0 :+ 1

----------------
-- Potentials --
----------------

-- | Free potential.
--   The potential energy is zero everywhere.
freeV
    :: Double  -- ^ position
    -> Double  -- ^ potential energy
freeV _x = 0

-- | Harmonic potential.
--   This is the potential energy of a linear spring.
harmonicV
    :: Double  -- ^ spring constant
    -> Double  -- ^ position
    -> Double  -- ^ potential energy
harmonicV k x = k * x**2 / 2

-- | A double well potential.
--   Potential energy is a quartic function of position
--   that gives two wells, each approximately harmonic
--   at the bottom of the well.
doubleWell
    :: Double  -- ^ width (for both wells and well separation)
    -> Double  -- ^ energy height of barrier between wells
    -> Double  -- ^ position
    -> Double  -- ^ potential energy
doubleWell a v0 x = v0 * ((x**2 - a**2)/a**2)**2

-- | Finite square well potential.
--   Potential is zero inside the well,
--   and constant outside the well.
--   Well is centered at the origin.
squareWell
    :: Double  -- ^ well width
    -> Double  -- ^ energy height of well
    -> Double  -- ^ position
    -> Double  -- ^ potential energy
squareWell l v0 x
    | abs x < l/2  = 0
    | otherwise    = v0

-- | A step barrier potential.
--   Potential is zero to left of origin.
stepV
    :: Double  -- ^ energy height of barrier (to the right of origin)
    -> Double  -- ^ position
    -> Double  -- ^ potential energy
stepV v0 x
    | x < 0      = 0
    | otherwise  = v0

-- | A potential barrier with thickness and height.
wall
    :: Double  -- ^ thickness of wall
    -> Double  -- ^ energy height of barrier
    -> Double  -- ^ position of center of barrier
    -> Double  -- ^ position
    -> Double  -- ^ potential energy
wall w v0 x0 x
    | abs (x-x0) < w/2  = v0
    | otherwise         = 0

---------------------------
-- Initial wavefunctions --
---------------------------

-- -- | Harmonic oscillator stationary state
-- harm :: Int          -- ^ nonnegative integer n identifying stationary state
--      -> Double       -- ^ x / sqrt(hbar/(m * omega)), i.e. position
--                      --   in units of sqrt(hbar/(m * omega))
--      -> C            -- ^ complex amplitude
-- harm n u
--     = exp (-u**2/2) * evalPhysHermite n u / sqrt (2^n * fact n * sqrt pi) :+ 0

coherent
    :: R       -- ^ length scale = sqrt(hbar / m omega)
    -> C       -- ^ parameter z
    -> R -> C  -- ^ wavefunction
coherent l z x
    = ((1/(pi*l**2))**0.25 * exp(-x**2/(2*l**2)) :+ 0)
      * exp(-z**2/2 + (sqrt(2/l**2) * x :+ 0) * z)

gaussian
    :: R       -- ^ width parameter
    -> R       -- ^ center of wave packet
    -> R -> C  -- ^ wavefunction
gaussian a x0 x = exp(-(x-x0)**2/(2*a**2)) / sqrt(a * sqrt pi) :+ 0

movingGaussian
    :: R       -- ^ width parameter
    -> R       -- ^ center of wave packet
    -> R       -- ^ l0 = hbar / p0
    -> R -> C  -- ^ wavefunction
movingGaussian a x0 l0 x = exp((0 :+ x/l0) - ((x-x0)**2/(2*a**2) :+ 0)) / (sqrt(a * sqrt pi) :+ 0)

---------------
-- Utilities --
---------------

fact :: Int -> Double
fact 0 = 1
fact n = fromIntegral n * fact (n-1)

linspace :: Double -> Double -> Int -> [Double]
linspace left right num
    = let dx = (right - left) / fromIntegral (num - 1)
      in [ left + dx * fromIntegral n | n <- [0..num-1]]

-- | Transform a wavefunction into a state vector.
stateVectorFromWavefunction :: R         -- ^ lowest x
                            -> R         -- ^ highest x
                            -> Int       -- ^ dimension of state vector
                            -> (R -> C)  -- ^ wavefunction
                            -> Vector C  -- ^ state vector
stateVectorFromWavefunction left right num psi
    = (num |>) [psi x | x <- linspace left right num]

hamiltonianMatrix :: R         -- ^ lowest x
                  -> R         -- ^ highest x
                  -> Int       -- ^ dimension of state vector
                  -> R         -- ^ hbar
                  -> R         -- ^ mass
                  -> (R -> R)  -- ^ potential energy function
                  -> Matrix C  -- ^ Hamiltonian Matrix
hamiltonianMatrix xmin xmax num hbar m pe
    = let coeff = -hbar**2/(2*m)
          dx = (xmax - xmin) / fromIntegral (num - 1)
          diagKEterm = -2 * coeff / dx**2
          offdiagKEterm = coeff / dx**2
          xs = linspace xmin xmax num
      in fromLists [[case abs(i-j) of
                       0  -> (diagKEterm + pe x) :+ 0
                       1  -> offdiagKEterm :+ 0
                       _  -> 0
                          | j <- [1..num] ] | (i,x) <- zip [1..num] xs]

expectX :: Vector C  -- ^ state vector
        -> Vector R  -- ^ vector of x values
        -> R         -- ^ <X>, expectation value of X
expectX psi xs = probVector psi <.> xs

glossScaleX :: Int -> (Double,Double) -> Double -> Float
glossScaleX screenWidth (xmin,xmax) x
    = let w = fromIntegral screenWidth :: Double
      in realToFrac $ (x - xmin) / (xmax - xmin) * w - w / 2

glossScaleY :: Int -> (Double,Double) -> Double -> Float
glossScaleY screenHeight (ymin,ymax) y
    = let h = fromIntegral screenHeight :: Double
      in realToFrac $ (y - ymin) / (ymax - ymin) * h - h / 2

glossScalePoint :: (Int,Int)        -- ^ (screenWidth,screenHeight)
                -> (Double,Double)  -- ^ (xmin,xmax)
                -> (Double,Double)  -- ^ (ymin,ymax)
                -> (Double,Double)  -- ^ (x,y)
                -> (Float,Float)
glossScalePoint (screenWidth,screenHeight) xMinMax yMinMax (x,y)
    = (glossScaleX screenWidth  xMinMax x
      ,glossScaleY screenHeight yMinMax y)

-- | Produce a gloss 'Picture' of state vector
--   for 1D wavefunction.
picture :: (Double, Double)    -- ^ y range
        -> [Double]            -- ^ xs
        -> Vector C            -- ^ state vector
        -> Picture
picture (ymin,ymax) xs psi
    = Color
      yellow
      (Line
       [glossScalePoint
        (screenWidth,screenHeight)
        (head xs, last xs)
        (ymin,ymax)
        p | p <- zip xs (map magSq $ toList psi)])
           where
             magSq = \z -> magnitude z ** 2
             screenWidth  = 1000
             screenHeight =  750

-- options for representing wave functions
-- 1.  A function R -> C
-- 2.  ([R],Vector C), where lengths match
-- 3.  [(R,C)]
-- 4.  (R,R,Vector C)  -- xmin, xmax, state vector (assumes even spacing)

-- 2,4 are best for evolution

listForm :: (R,R,Vector C) -> ([R],Vector C)
listForm (xmin,xmax,v)
    = let dt = (xmax - xmin) / fromIntegral (size v - 1)
      in ([xmin, xmin + dt .. xmax],v)

{-
-- | Given an initial state vector and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   interval for one timestep.
simulate1D :: [Double] -> Vector C -> (Float -> (Float,[Double],Vector C) -> (Float,[Double],Vector C)) -> IO ()
simulate1D xs initial statePropFunc
    = simulate display black 10 (0,initial) displayFunc (const statePropFunc)
      where
        display = InWindow "Animation" (screenWidth,screenHeight) (10,10)
        displayFunc (_t,v) = Color yellow (Line [(

      white (\tFloat -> Pictures [Color blue (Line (points (realToFrac tFloat)))
                                 ,axes (screenWidth,screenHeight) (xmin,xmax) (ymin,ymax)])

-- | Produce a state propagation function from a time-dependent Hamiltonian.
--   The float is dt.
statePropGloss :: (Double -> Matrix C) -> Float -> (Float,Vector C) -> (Float,Vector C)
statePropGloss ham dt (tOld,v)
    = (tNew, timeEv (realToFrac dt) (ham tMid) v)
      where
        tNew = tOld + dt
        tMid = realToFrac $ (tNew + tOld) / 2

-- | Given an initial state vector and a time-dependent Hamiltonian,
--   produce a visualization of a 1D wavefunction.
evolutionBlochSphere :: Vector C -> (Double -> Matrix C) -> IO ()
evolutionBlochSphere psi0 ham
    = simulateBlochSphere 0.01 psi0 (stateProp ham)

-}

{-
def triDiagMatrixMult(square_arr,arr):
    num = len(arr)
    result = array([0 for n in range(num)],dtype=complex128)
    result[0] = square_arr[0][0] * arr[0] + square_arr[0][1] * arr[1]
    for n in range(1,num-1):
        result[n] = square_arr[n][n-1] * arr[n-1] + square_arr[n][n] * arr[n] \
            + square_arr[n][n+1] * arr[n+1]
    result[num-1] = square_arr[num-1][num-2] * arr[num-2] \
        + square_arr[num-1][num-1] * arr[num-1]
    return result
-}

------------------
-- Main program --
------------------

-- n is number of points
-- n-1 is number of intervals
xRange :: R -> R -> Int -> [R]
xRange xmin xmax n
    = let dt = (xmax - xmin) / fromIntegral (n - 1)
      in [xmin, xmin + dt .. xmax]

{-
if __name__ == '__main__':
    m = 1
    omega = 10
    xmin = -2.0
    xmax =  2.0
    num = 256
    num = 128
    dt = 0.0002
    dt = 0.01
    xs = linspace(xmin,xmax,num)
    dx = xs[1] - xs[0]

    super = lambda x: (harm0(m,omega)(x) + harm1(m,omega)(x))/sqrt(2)
    shiftedHarm = lambda x: harm0(m,omega)(x-1)
    coh = coherent(m,omega,1)

    print sum(conj(psi)*psi)*dx

    harmV = harmonicV(m * omega**2)

    V = doubleWell(1,0.1*hbar*omega)
    V = squareWell(1.0,hbar*omega)
    V = harmonicV(m*omega**2)
    V = stepV(10*hbar*omega)
    V = wall(0.1,14.0*hbar*omega,0)
    V = freeV

    H = matrixH(m,xmin,xmax,num,V)
    I = matrixI(num)

    (vals,vecs) = eigh(H)

    E0 = vals[0]
    E1 = vals[1]
    psi0 = normalize(transpose(vecs)[0],dx)
    psi1 = normalize(transpose(vecs)[1],dx)

    psi = func2psi(gaussian(0.3,1),xmin,xmax,num)
    psi = func2psi(coh,xmin,xmax,num)
    psi = func2psi(movingGaussian(0.3,10,-1),xmin,xmax,num)

    psi = psi0
    psi = psi1
    psi = (psi0 + psi1)/sqrt(2)

    E = sum(conj(psi)*triDiagMatrixMult(H,psi)).real*dx

    Escale = hbar*omega

    print E
    print Escale

    leftM  = I + 0.5 * i * H / hbar * dt
    rightM = I - 0.5 * i * H / hbar * dt

    box = display(title='Schrodinger Equation',width=1000,height=1000)

    c = curve(pos = psi2rho(psi,xs))
    c.color = color.blue
    c.radius = 0.02

    ball = sphere(radius=0.05,color=color.red,pos=(expectX(psi,xs),0,0))

    pot_curve = [(x,V(x)/Escale,0) for x in xs if V(x)/Escale < xmax]
    pot = curve(color=color.green,pos=pot_curve,radius=0.01)

    Eline = curve(color=(1,1,0),pos=[(x,E/Escale) for x in xs])
    axis = curve(color=color.white,pos=[(x,0) for x in xs])

    while 1:
        psi = solve(leftM,triDiagMatrixMult(rightM,psi))
        c.pos = psi2rho(psi,xs)
        ball.x = expectX(psi,xs)

To Do:
add combinators for potentials
to shift horizontally and vertically,
and to add potentials

-}

-- Are we committed to SI units for hbar?  No.
-- harmonic oscillator functions depend only on sqrt(hbar/m omega)
-- which is a length parameter
-- for moving gaussian, could give hbar/p0 instead of p0
-- (is that debrogie wavelength?  I think it's h/p0)


{-# LANGUAGE TemplateHaskell #-}
module Lib
    ( toPicture, export, mandelbrot, burningShip
    , FractalParams(..), Size(..), Position(..), Zoom, Image
    , size, position, zoom, innerColor, outerColor
    , updatePos
    ) where
import Codec.BMP(writeBMP, packRGBA32ToBMP24)
import Codec.Picture(readBitmap, writeDynamicPng)
import Control.Lens(makeLenses, (^.), (+~))
import qualified Data.ByteString as BS
import Data.Complex(Complex(..),realPart,imagPart)
import Data.Word(Word8)
import qualified Graphics.Gloss as G
import System.Directory(removeFile)

type ColorPart = Word8
type Color = (ColorPart, ColorPart, ColorPart)
type Image = ([Color], Size)
type Zoom a = a

data Size = Size {width :: Int, height :: Int}

data Position a = Pos {_x :: a, _y :: a}
makeLenses ''Position

data FractalParams a =
  F { _size :: Size
    , _position :: Position a
    , _zoom :: Zoom a
    , _innerColor :: G.Color
    , _outerColor :: G.Color
    }
makeLenses ''FractalParams

updatePos :: (Real a, Fractional (Zoom b)) =>
             a -> a -> FractalParams b -> FractalParams b
updatePos x' y' fp = newCoord x x' . newCoord y (-y') $ fp
  where newCoord c c' = position.c +~ realToFrac c'/_zoom fp

toList :: Color -> [ColorPart]
toList (r, g, b) = [r, g, b, 255]

toPicture :: Image -> G.Picture
toPicture (pxs, s) = G.bitmapOfByteString (width s) (height s) (pixelsToByteString pxs) True
  where pixelsToByteString = BS.pack . concatMap (reverse . toList)

{-# SPECIALIZE burningShip :: FractalParams Float -> Image #-}
{-# SPECIALIZE burningShip :: FractalParams Double -> Image #-}
burningShip :: (Enum (Zoom a), RealFloat (Zoom a)) => FractalParams a -> Image
burningShip = generateFractal $ iterations abs (255::Int) 3

{-# SPECIALIZE mandelbrot :: FractalParams Float -> Image #-}
{-# SPECIALIZE mandelbrot :: FractalParams Double -> Image #-}
mandelbrot :: (Enum (Zoom a), RealFloat (Zoom a)) => FractalParams a -> Image
mandelbrot = generateFractal $ iterations id (255::Int) 3

generateFractal :: (Enum (Zoom a), RealFloat (Zoom a), Real b) =>
                   (Complex (Zoom a) -> b) -> FractalParams a -> Image
generateFractal iterationFun fp = (map (getPixel' . iterationFun) coords, _size fp)
  where getPixel' = getPixel (toFloats innerColor) (toFloats outerColor)
        toFloats = G.rgbaOfColor . (fp ^.)
        coords = getCoordinates (_size fp) (toComplex $ _position fp) (_zoom fp)
        toComplex p = _x p :+ _y p

getCoordinates :: (Enum a, RealFloat a) => Size -> Complex a -> a -> [Complex a]
getCoordinates s p z = [ (r/z :+ i/z) + p
                       | i <- reverse . range $ height s
                       , r <- range $ width s]
  where range n = let half = (fromIntegral n-1)/2 in [-half..half]

{-# SPECIALIZE iterations :: (Float -> Float) -> Int -> Float -> Complex Float -> Int #-}
{-# SPECIALIZE iterations :: (Double -> Double) -> Int -> Double -> Complex Double -> Int #-}
iterations :: (Eq a, Num a, RealFloat b) => (b -> b) -> a -> b -> Complex b -> a
iterations f iterLimit zlimit c = go c 0
  where magIsGTzLimit (r:+i) = r * r + i * i > zlimit * zlimit
        go prev count | prev `seq` count == iterLimit || magIsGTzLimit prev = count
                      | otherwise = go (prev'*prev' + c) (count + 1)
          where prev' = f (realPart prev) :+ f (imagPart prev)

getPixel :: (Integral a, Real b) => (Float,Float,Float,x) -> (Float,Float,Float,y) -> b -> (a,a,a)
getPixel (r1, g1, b1, _) (r2, g2, b2, _) n = (avg r1 r2, avg g1 g2, avg b1 b2)
  where avg start end = truncate $ 255 * (end + frac * (start - end))
        frac = logBase 1.25 (1+realToFrac n::Float)/25

export :: FilePath -> Image -> IO ()
export path (pxs, s) = do
  writeBMP path . packRGBA32ToBMP24 (width s) (height s) .
    BS.pack . concatMap toList $ pxs
  (Right im) <- readBitmap path
  _ <- writeDynamicPng (path ++ ".png") im
  removeFile path


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE NoStarIsType #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}

-- | Arrays with a fixed shape, with a HMatrix representation.
module NumHask.Array.HMatrix
  ( -- $usage
    Array (..),

    -- * Representation
    --
    -- With no functor instance, we instead supply the representable API
    index,
    tabulate,

    -- * Conversion
    shape,
    toDynamic,
    toFixed,
    fromFixed,

    -- * Operators
    reshape,
    transpose,
    diag,
    ident,
    singleton,
    selects,
    selectsExcept,
    folds,
    concatenate,
    insert,
    append,
    reorder,
    expand,
    slice,
    squeeze,

    -- * Scalar
    --
    -- Scalar specialisations
    Scalar,
    fromScalar,
    toScalar,

    -- * Vector
    --
    -- Vector specialisations.
    Vector,

    -- * Matrix
    --
    -- Matrix specialisations.
    Matrix,
    col,
    row,
    safeCol,
    safeRow,
    mmult,
  )
where

import qualified Data.Vector as V
import GHC.Exts (IsList (..))
import GHC.TypeLits
import qualified NumHask.Array.Dynamic as D
import qualified NumHask.Array.Fixed as F
import NumHask.Array.Shape
import NumHask.Prelude as P
import qualified Numeric.LinearAlgebra as H
import qualified Numeric.LinearAlgebra.Devel as H
import qualified Prelude
import Data.Proxy (Proxy(..))

-- $setup
-- >>> :set -XDataKinds
-- >>> :set -XOverloadedLists
-- >>> :set -XTypeFamilies
-- >>> :set -XFlexibleContexts
-- >>> import GHC.TypeLits
-- >>> import GHC.Int
-- >>> let s = [1] :: Array ('[] :: [Nat]) Int64 -- scalar
-- >>> let v = [1,2,3] :: Array '[3] Int64       -- vector
-- >>> let m = [0..11] :: Array '[3,4] Int64     -- matrix
-- >>> let a = [1..24] :: Array '[2,3,4] Int64

-- | a multidimensional array with a type-level shape
--
-- >>> let a = [1..24] :: Array '[2,3,4] Int64
-- >>> a
-- [[[1, 2, 3, 4],
--   [5, 6, 7, 8],
--   [9, 10, 11, 12]],
--  [[13, 14, 15, 16],
--   [17, 18, 19, 20],
--   [21, 22, 23, 24]]]
--
-- >>> [1,2,3] :: Array '[2,2] Int64
-- [[*** Exception: NumHaskException {errorMessage = "shape mismatch"}
newtype Array s a = Array {unArray :: H.Matrix a}
  deriving (Show, Generic)

instance
  ( Additive a,
    HasShape s,
    H.Container H.Vector a,
    Num a
  ) =>
  Additive (Array s a)
  where
  (+) (Array x1) (Array x2) = Array $ H.add x1 x2

  zero = Array $ H.konst zero (n, m)
    where
      s = shapeVal (toShape @s)
      [n, m] = s

instance
  ( Multiplicative a,
    HasShape s,
    H.Container H.Vector a,
    Num (H.Vector a),
    Num a
  ) =>
  Multiplicative (Array s a)
  where
  (*) (Array x1) (Array x2) = Array $ H.liftMatrix2 (Prelude.*) x1 x2

  one = Array $ H.konst one (n, m)
    where
      s = shapeVal (toShape @s)
      [n, m] = s

-- (<.>) (Array a) (Array b) = H.sumElements $ H.liftMatrix2 (Prelude.*) a b

instance
  ( Subtractive a,
    HasShape s,
    H.Container H.Vector a,
    Num a
  ) =>
  Subtractive (Array s a)
  where
  negate (Array x1) = Array $ H.cmap negate x1

instance
  (HasShape s, Multiplicative a, H.Container H.Vector a, Num a) =>
  MultiplicativeAction (Array s a) a
  where
  (.*) s (Array r) = Array $ H.cmap (s *) r
  {-# INLINE (.*) #-}

  (*.) (Array r) s = Array $ H.cmap (*s) r
  {-# INLINE (*.) #-}

instance
  (HasShape s, Additive a, H.Container H.Vector a, Num a) =>
  AdditiveAction (Array s a) a
  where
  (.+) s (Array r) = Array $ H.cmap (s +) r
  {-# INLINE (.+) #-}

  (+.) (Array r) s = Array $ H.cmap (+s) r
  {-# INLINE (+.) #-}

instance
  (HasShape s, Subtractive a, H.Container H.Vector a, Num a) =>
  SubtractiveAction (Array s a) a
  where
  (.-) s (Array r) = Array $ H.cmap (s -) r
  {-# INLINE (.-) #-}

  (-.) (Array r) s = Array $ H.cmap (\x -> x - s) r
  {-# INLINE (-.) #-}

instance
  (HasShape s, Divisive a, H.Container H.Vector a, Num a) =>
  DivisiveAction (Array s a) a
  where
  (./) s (Array r) = Array $ H.cmap (s /) r
  {-# INLINE (./) #-}

  (/.) (Array r) s = Array $ H.cmap (/ s) r
  {-# INLINE (/.) #-}

-- | from flat list
instance
  ( HasShape s,
    H.Element a
  ) =>
  IsList (Array s a)
  where
  type Item (Array s a) = a

  fromList l =
    bool
      (throw (NumHaskException "shape mismatch"))
      (Array $ H.reshape n $ H.fromList l)
      ((length l == 1 && null s) || (length l == size s))
    where
      s = shapeVal (toShape @s)
      n = Prelude.last s

  toList (Array v) = H.toList $ H.flatten v

-- | Get shape of an Array as a value.
--
-- >>> shape a
-- [2,3,4]
shape :: forall a s. (HasShape s) => Array s a -> [Int]
shape _ = shapeVal $ toShape @s
{-# INLINE shape #-}

-- | Convert to a dynamic array.
toDynamic :: (HasShape s, H.Element a) => Array s a -> D.Array a
toDynamic a@(Array h) = D.fromFlatList (shape a) (mconcat $ H.toLists h)

-- | Convert to a fixed array.
toFixed :: (HasShape s, H.Element a) => Array s a -> F.Array s a
toFixed (Array h) = fromList (mconcat $ H.toLists h)

-- | Convert from a fixed array.
fromFixed :: (HasShape s, H.Element a) => F.Array s a -> Array s a
fromFixed a = fromList (P.toList a)

-- | with no fmap, we supply the representable API
index ::
  forall s a.
  ( HasShape s,
    H.Element a,
    H.Container H.Vector a
  ) =>
  Array s a ->
  [Int] ->
  a
index (Array v) i = H.flatten v `H.atIndex` flatten s i
  where
    s = shapeVal (toShape @s)

-- | tabulate an array with a generating function
--
-- >>> tabulate [2,3,4] ((1+) . flatten [2,3,4]) == a
-- True
tabulate ::
  forall s a.
  ( HasShape s,
    H.Element a
  ) =>
  ([Int] -> a) ->
  Array s a
tabulate f =
  fromList (V.toList $ V.generate (size s) (f . shapen s))
  where
    s = shapeVal (toShape @s)

-- | Reshape an array (with the same number of elements).
--
-- >>> reshape a :: Array '[4,3,2] Int64
-- [[[1, 2],
--   [3, 4],
--   [5, 6]],
--  [[7, 8],
--   [9, 10],
--   [11, 12]],
--  [[13, 14],
--   [15, 16],
--   [17, 18]],
--  [[19, 20],
--   [21, 22],
--   [23, 24]]]
reshape ::
  forall a s s'.
  ( Size s ~ Size s',
    HasShape s,
    HasShape s',
    H.Container H.Vector a
  ) =>
  Array s a ->
  Array s' a
reshape a = tabulate (index a . shapen s . flatten s')
  where
    s = shapeVal (toShape @s)
    s' = shapeVal (toShape @s')

-- | Reverse indices eg transposes the element A/ijk/ to A/kji/.
--
-- >>> index (transpose a) [1,0,0] == index a [0,0,1]
-- True
transpose :: forall a s. (H.Element a, H.Container H.Vector a, HasShape s, HasShape (Reverse s)) => Array s a -> Array (Reverse s) a
transpose a = tabulate (index a . reverse)

-- | The identity array.
--
-- >>> ident :: Array '[3,2] Int64
-- [[1, 0],
--  [0, 1],
--  [0, 0]]
ident :: forall a s. (H.Element a, H.Container H.Vector a, HasShape s, Additive a, Multiplicative a) => Array s a
ident = tabulate (bool zero one . isDiag)
  where
    isDiag [] = True
    isDiag [_] = True
    isDiag [x, y] = x == y
    isDiag (x : y : xs) = x == y && isDiag (y : xs)

-- | Extract the diagonal of an array.
--
-- >>> diag (ident :: Array '[3,2] Int64)
-- [1, 1]
diag ::
  forall a s.
  ( HasShape s,
    HasShape '[Minimum s],
    H.Element a,
    H.Container H.Vector a
  ) =>
  Array s a ->
  Array '[Minimum s] a
diag a = tabulate go
  where
    go [] = throw (NumHaskException "Rank Underflow")
    go (s' : _) = index a (replicate (length ds) s')
    ds = shapeVal (toShape @s)

-- | Create an array composed of a single value.
--
-- >>> singleton one :: Array '[3,2] Int64
-- [[1, 1],
--  [1, 1],
--  [1, 1]]
singleton :: (H.Element a, H.Container H.Vector a, HasShape s) => a -> Array s a
singleton a = tabulate (const a)

-- | Select an array along dimensions.
--
-- >>> let s = selects (Proxy :: Proxy '[0,1]) [1,1] a
-- >>> :t s
-- s :: Array '[4] Int64
--
-- >>> s
-- [17, 18, 19, 20]
selects ::
  forall ds s s' a.
  ( HasShape s,
    HasShape ds,
    HasShape s',
    s' ~ DropIndexes s ds,
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy ds ->
  [Int] ->
  Array s a ->
  Array s' a
selects _ i a = tabulate go
  where
    go s = index a (addIndexes s ds i)
    ds = shapeVal (toShape @ds)

-- | Select an index /except/ along specified dimensions.
--
-- >>> let s = selectsExcept (Proxy :: Proxy '[2]) [1,1] a
-- >>> :t s
-- s :: Array '[4] Int64
--
-- >>> s
-- [17, 18, 19, 20]
selectsExcept ::
  forall ds s s' a.
  ( HasShape s,
    HasShape ds,
    HasShape s',
    s' ~ TakeIndexes s ds,
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy ds ->
  [Int] ->
  Array s a ->
  Array s' a
selectsExcept _ i a = tabulate go
  where
    go s = index a (addIndexes i ds s)
    ds = shapeVal (toShape @ds)

-- | Fold along specified dimensions.
--
-- >>> folds sum (Proxy :: Proxy '[1]) a
-- [68, 100, 132]
folds ::
  forall ds st si so a b.
  ( HasShape st,
    HasShape ds,
    HasShape si,
    HasShape so,
    si ~ DropIndexes st ds,
    so ~ TakeIndexes st ds,
    H.Element a,
    H.Container H.Vector a,
    H.Element b,
    H.Container H.Vector b
  ) =>
  (Array si a -> b) ->
  Proxy ds ->
  Array st a ->
  Array so b
folds f d a = tabulate go
  where
    go s = f (selects d s a)

-- | Concatenate along a dimension.
--
-- >>> :t concatenate (Proxy :: Proxy 1) a a
-- concatenate (Proxy :: Proxy 1) a a :: Array '[2, 6, 4] Int64
concatenate ::
  forall a s0 s1 d s.
  ( CheckConcatenate d s0 s1 s,
    Concatenate d s0 s1 ~ s,
    HasShape s0,
    HasShape s1,
    HasShape s,
    KnownNat d,
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy d ->
  Array s0 a ->
  Array s1 a ->
  Array s a
concatenate _ s0 s1 = tabulate go
  where
    go s =
      bool
        (index s0 s)
        ( index
            s1
            ( addIndex
                (dropIndex s d)
                d
                ((s !! d) - (ds0 !! d))
            )
        )
        ((s !! d) >= (ds0 !! d))
    ds0 = shapeVal (toShape @s0)
    d = fromIntegral $ natVal @d Proxy

-- | Insert along a dimension at a position.
--
-- >>> insert (Proxy :: Proxy 2) (Proxy :: Proxy 0) a ([100..105])
-- [[[100, 1, 2, 3, 4],
--   [101, 5, 6, 7, 8],
--   [102, 9, 10, 11, 12]],
--  [[103, 13, 14, 15, 16],
--   [104, 17, 18, 19, 20],
--   [105, 21, 22, 23, 24]]]
insert ::
  forall a s s' d i.
  ( DropIndex s d ~ s',
    CheckInsert d i s,
    KnownNat i,
    KnownNat d,
    HasShape s,
    HasShape s',
    HasShape (Insert d s),
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy d ->
  Proxy i ->
  Array s a ->
  Array s' a ->
  Array (Insert d s) a
insert _ _ a b = tabulate go
  where
    go s
      | s !! d == i = index b (dropIndex s d)
      | s !! d < i = index a s
      | otherwise = index a (decAt d s)
    d = fromIntegral $ natVal @d Proxy
    i = fromIntegral $ natVal @i Proxy

-- | Insert along a dimension at the end.
--
-- >>>  :t append (Proxy :: Proxy 0) a
-- append (Proxy :: Proxy 0) a
--   :: Array '[3, 4] Int64 -> Array '[3, 3, 4] Int64
append ::
  forall a d s s'.
  ( DropIndex s d ~ s',
    CheckInsert d (Dimension s d - 1) s,
    KnownNat (Dimension s d - 1),
    KnownNat d,
    HasShape s,
    HasShape s',
    HasShape (Insert d s),
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy d ->
  Array s a ->
  Array s' a ->
  Array (Insert d s) a
append d = insert d (Proxy :: Proxy (Dimension s d - 1))

-- | Change the order of dimensions.
--
-- >>> let r = reorder (Proxy :: Proxy '[2,0,1]) a
-- >>> :t r
-- r :: Array '[4, 2, 3] Int64
reorder ::
  forall a ds s.
  ( HasShape ds,
    HasShape s,
    HasShape (Reorder s ds),
    CheckReorder ds s,
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy ds ->
  Array s a ->
  Array (Reorder s ds) a
reorder _ a = tabulate go
  where
    go s = index a (addIndexes [] ds s)
    ds = shapeVal (toShape @ds)

-- | Product two arrays using the supplied binary function.
--
-- For context, if the function is multiply, and the arrays are tensors,
-- then this can be interpreted as a tensor product.
--
-- https://en.wikipedia.org/wiki/Tensor_product
--
-- The concept of a tensor product is a dense crossroad, and a complete treatment is elsewhere.  To quote:
--
-- ... the tensor product can be extended to other categories of mathematical objects in addition to vector spaces, such as to matrices, tensors, algebras, topological vector spaces, and modules. In each such case the tensor product is characterized by a similar universal property: it is the freest bilinear operation. The general concept of a "tensor product" is captured by monoidal categories; that is, the class of all things that have a tensor product is a monoidal category.
--
-- >>> expand (*) v v
-- [[1, 2, 3],
--  [2, 4, 6],
--  [3, 6, 9]]
expand ::
  forall s s' a b c.
  ( HasShape s,
    HasShape s',
    HasShape ((++) s s'),
    H.Element a,
    H.Container H.Vector a,
    H.Element b,
    H.Container H.Vector b,
    H.Element c
  ) =>
  (a -> b -> c) ->
  Array s a ->
  Array s' b ->
  Array ((++) s s') c
expand f a b = tabulate (\i -> f (index a (take r i)) (index b (drop r i)))
  where
    r = rank (shape a)

-- | Select elements along positions in every dimension.
--
-- >>> let s = slice (Proxy :: Proxy '[[0,1],[0,2],[1,2]]) a
-- >>> :t s
-- s :: Array '[2, 2, 2] Int64
--
-- >>> s
-- [[[2, 3],
--   [10, 11]],
--  [[14, 15],
--   [22, 23]]]
--
-- >>> let s = squeeze $ slice (Proxy :: Proxy '[ '[0], '[0], '[0]]) a
-- >>> :t s
-- s :: Array '[] Int64
--
-- >>> s
-- 1
slice ::
  forall (pss :: [[Nat]]) s s' a.
  ( HasShape s,
    HasShape s',
    KnownNatss pss,
    KnownNat (Rank pss),
    s' ~ Ranks pss,
    H.Element a,
    H.Container H.Vector a
  ) =>
  Proxy pss ->
  Array s a ->
  Array s' a
slice pss a = tabulate go
  where
    go s = index a (zipWith (!!) pss' s)
    pss' = natValss pss

-- | Remove single dimensions.
--
-- >>> let a = [1..24] :: Array '[2,1,3,4,1] Int64
-- >>> a
-- [[[[[1],
--     [2],
--     [3],
--     [4]],
--    [[5],
--     [6],
--     [7],
--     [8]],
--    [[9],
--     [10],
--     [11],
--     [12]]]],
--  [[[[13],
--     [14],
--     [15],
--     [16]],
--    [[17],
--     [18],
--     [19],
--     [20]],
--    [[21],
--     [22],
--     [23],
--     [24]]]]]
-- >>> squeeze a
-- [[[1, 2, 3, 4],
--   [5, 6, 7, 8],
--   [9, 10, 11, 12]],
--  [[13, 14, 15, 16],
--   [17, 18, 19, 20],
--   [21, 22, 23, 24]]]
--
-- >>> squeeze ([1] :: Array '[1,1] Double)
-- 1.0
squeeze ::
  forall s t a.
  (t ~ Squeeze s) =>
  Array s a ->
  Array t a
squeeze (Array x) = Array x

-- $scalar
-- Scalar specialisations

-- | <https://en.wikipedia.org/wiki/Scalarr_(mathematics) Wiki Scalar>
--
-- An Array '[] a despite being a Scalar is never-the-less a one-element vector under the hood. Unification of representation is unexplored.
type Scalar a = Array ('[] :: [Nat]) a

-- | Unwrapping scalars is probably a performance bottleneck.
--
-- >>> let s = [3] :: Array ('[] :: [Nat]) Int64
-- >>> fromScalar s
-- 3
fromScalar :: (H.Element a, H.Container H.Vector a, HasShape ('[] :: [Nat])) => Array ('[] :: [Nat]) a -> a
fromScalar a = index a ([] :: [Int])

-- | Convert a number to a scalar.
--
-- >>> :t toScalar 2
-- toScalar 2 :: Num a => Array '[] a
toScalar :: (H.Element a, H.Container H.Vector a, HasShape ('[] :: [Nat])) => a -> Array ('[] :: [Nat]) a
toScalar a = fromList [a]

-- | <https://en.wikipedia.org/wiki/Vector_(mathematics_and_physics) Wiki Vector>
type Vector s a = Array '[s] a

-- | <https://en.wikipedia.org/wiki/Matrix_(mathematics) Wiki Matrix>
type Matrix m n a = Array '[m, n] a

instance
  ( Multiplicative a,
    P.Distributive a,
    Subtractive a,
    H.Numeric a,
    KnownNat m,
    HasShape '[m, m],
    H.Element a,
    H.Container H.Vector a
  ) =>
  Multiplicative (Matrix m m a)
  where
  (*) = mmult

  one = ident

-- | Extract specialised to a matrix.
--
-- >>> row 1 m
-- [4, 5, 6, 7]
row :: forall m n a. (H.Element a, H.Container H.Vector a, KnownNat m, KnownNat n, HasShape '[m, n]) => Int -> Matrix m n a -> Vector n a
row i (Array a) = fromList $ H.toList $ H.subVector (i * n) n (H.flatten a)
  where
    n = fromIntegral $ natVal @n Proxy

-- | Row extraction checked at type level.
--
-- >>> safeRow (Proxy :: Proxy 1) m
-- [4, 5, 6, 7]
--
-- >>> safeRow (Proxy :: Proxy 3) m
-- ...
-- ... index outside range
-- ...
safeRow :: forall m n a j. (H.Element a, H.Container H.Vector a, 'True ~ CheckIndex j m, KnownNat j, KnownNat m, KnownNat n, HasShape '[m, n]) => Proxy j -> Matrix m n a -> Vector n a
safeRow _j (Array a) = fromList $ H.toList $ H.subVector (j * n) n (H.flatten a)
  where
    n = fromIntegral $ natVal @n Proxy
    j = fromIntegral $ natVal @j Proxy

-- | Extract specialised to a matrix.
--
-- >>> col 1 m
-- [1, 5, 9]
col :: forall m n a. (H.Element a, H.Container H.Vector a, KnownNat m, KnownNat n, HasShape '[m, n]) => Int -> Matrix m n a -> Vector n a
col i (Array a) = Array $ H.takeColumns i a

-- | Column extraction checked at type level.
--
-- >>> safeCol (Proxy :: Proxy 1) m
-- [1, 5, 9]
--
-- >>> safeCol (Proxy :: Proxy 4) m
-- ...
-- ... index outside range
-- ...
safeCol :: forall m n a j. (H.Element a, H.Container H.Vector a, 'True ~ CheckIndex j n, KnownNat j, KnownNat m, KnownNat n, HasShape '[m, n]) => Proxy j -> Matrix m n a -> Vector n a
safeCol _j (Array a) = Array $ H.takeColumns j a
  where
    j = fromIntegral $ natVal @j Proxy

-- | Matrix multiplication.
--
-- This is dot sum (*) specialised to matrices
--
-- >>> let a = [1, 2, 3, 4] :: Array '[2, 2] Int64
-- >>> let b = [5, 6, 7, 8] :: Array '[2, 2] Int64
-- >>> a
-- [[1, 2],
--  [3, 4]]
--
-- >>> b
-- [[5, 6],
--  [7, 8]]
--
-- >>> mmult a b
-- [[19, 22],
--  [43, 50]]
mmult ::
  forall m n k a.
  ( KnownNat k,
    KnownNat m,
    KnownNat n,
    HasShape [m, n],
    Ring a,
    H.Numeric a
  ) =>
  Array [m, k] a ->
  Array [k, n] a ->
  Array [m, n] a
mmult (Array x) (Array y) = Array $ x H.<> y


{-# LANGUAGE DataKinds, TypeFamilies #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
-- | Matrix-based (meaning that dual numbers for gradient computation
-- consider matrices, not scalars, as the primitive differentiable type)
-- implementation of fully connected neutral network for classification
-- of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistFcnnMatrix where

import Prelude

import           Control.Exception (assert)
import qualified Data.Array.DynamicS as OT
import qualified Data.Vector.Generic as V
import           GHC.Exts (inline)
import           Numeric.LinearAlgebra (Vector)

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, var1, var2)
import HordeAd.Tool.MnistData

fcnnMnistLen2 :: Int -> Int -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
fcnnMnistLen2 widthHidden widthHidden2 =
  ( 0
  , [widthHidden, widthHidden2, sizeMnistLabel]
  , [ (widthHidden, sizeMnistGlyph)
    , (widthHidden2, widthHidden)
    , (sizeMnistLabel, widthHidden2) ]
  , []
  )

-- | Fully connected neural network for the MNIST digit classification task.
-- There are two hidden layers and both use the same activation function.
-- The output layer uses a different activation function.
-- The width of the layers is determined by the dimensions of the matrices
-- and vectors given as dual number parameters (variables).
-- The dimensions, in turn, can be computed by the @len*@ functions
-- on the basis of the requested widths, see above.
fcnnMnist2 :: forall d r m. DualMonad d r m
           => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> Vector r
           -> DualNumberVariables d r
           -> m (DualNumber d (Vector r))
fcnnMnist2 factivationHidden factivationOutput input variables = do
  let !_A = assert (sizeMnistGlyph == V.length input) ()
      weightsL0 = var2 variables 0
      biasesV0 = var1 variables 0
      weightsL1 = var2 variables 1
      biasesV1 = var1 variables 1
      weightsL2 = var2 variables 2
      biasesV2 = var1 variables 2
  let hiddenLayer1 = weightsL0 #>!! input + biasesV0
  nonlinearLayer1 <- factivationHidden hiddenLayer1
  let hiddenLayer2 = weightsL1 #>! nonlinearLayer1 + biasesV1
  nonlinearLayer2 <- factivationHidden hiddenLayer2
  let outputLayer = weightsL2 #>! nonlinearLayer2 + biasesV2
  factivationOutput outputLayer

-- | The neural network applied to concrete activation functions
-- and composed with the appropriate loss function.
fcnnMnistLoss2
  :: DualMonad d r m
  => MnistData r -> DualNumberVariables d r -> m (DualNumber d r)
fcnnMnistLoss2 (input, target) variables = do
  result <- inline fcnnMnist2 logisticAct softMaxActV input variables
  lossCrossEntropyV target result

-- | The neural network applied to concrete activation functions
-- and composed with the appropriate loss function, using fused
-- softMax and cross entropy as the loss function.
fcnnMnistLossFused2
  :: DualMonad d r m
  => MnistData r -> DualNumberVariables d r -> m (DualNumber d r)
fcnnMnistLossFused2 (input, target) variables = do
  result <- inline fcnnMnist2 logisticAct return input variables
  lossSoftMaxCrossEntropyV target result

fcnnMnistLossFusedRelu2
  :: DualMonad d r m
  => MnistData r -> DualNumberVariables d r -> m (DualNumber d r)
fcnnMnistLossFusedRelu2 (input, target) variables = do
  result <- inline fcnnMnist2 reluAct return input variables
  lossSoftMaxCrossEntropyV target result

-- | A function testing the neural network given testing set of inputs
-- and the trained parameters.
fcnnMnistTest2
  :: forall r. IsScalar 'DModeGradient r
  => [MnistData r] -> Domains r -> r
fcnnMnistTest2 inputs parameters =
  let matchesLabels :: MnistData r -> Bool
      matchesLabels (glyph, label) =
        let nn = inline (fcnnMnist2 @'DModeGradient)
                        logisticAct softMaxActV glyph
            value = primalValue nn parameters
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)


{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}

module Advection1D
where

import qualified Basis
import           Data.List
import qualified Element
import           Flux
import           Mesh
import           Numeric.LinearAlgebra
--import           Numeric.LinearAlgebra.HMatrix
import           Quadrature            (integrate)
import qualified ShapeFcns

-- Synonym for numeric/fractional constraints
type FrElNuFi a = (Fractional a,Element a,Numeric a,Field a)

-- Element stiffness integrand
stiffnessIntegrand :: (Element.Element e,ShapeFcns.ShapeFcn s,Basis.Basis b,FrElNuFi a) => e a -> s b -> [a] -> Matrix a
--stiffnessIntegrand elem shpFcn xi = fromLists [[
--  ShapeFcns.n shpFcn xi i [0] * atIndex (Element.dndx elem shpFcn xi j) 0 * Element.computeJacobianDet elem shpFcn xi
--  | i <- idxRange] | j <- idxRange]
stiffnessIntegrand elem shpFcn xi = fromLists [[
  (atIndex (Element.dndx elem shpFcn xi i) 0 * atIndex (Element.dndx elem shpFcn xi j) 0 + ShapeFcns.n shpFcn xi i [0] * ShapeFcns.n shpFcn xi j [0]) * Element.computeJacobianDet elem shpFcn xi
  | i <- idxRange] | j <- idxRange]
  where
    idxRange = [0..(Element.getNumNodes elem - 1)]

-- Element mass integrand
massIntegrand :: (Element.Element e,ShapeFcns.ShapeFcn s,Basis.Basis b,FrElNuFi a) => e a -> s b -> [a] -> Matrix a
massIntegrand elem shpFcn xi = fromLists [[
  ShapeFcns.n shpFcn xi i [0] * ShapeFcns.n shpFcn xi j [0] * Element.computeJacobianDet elem shpFcn xi
  | i <- idxRange] | j <- idxRange]
  where
    idxRange = [0..(ShapeFcns.getShapeFcnOrder shpFcn)]

-- Computes the element flux matrix assuming an upwind flux.
-- Assumes that advSpeed is positive.
elemFluxMatrix :: FrElNuFi a => a -> Matrix a
elemFluxMatrix advSpeed = scale advSpeed $ fromLists [[-1.0, 0.0, 0.0],[0.0, 0.0, 1.0]]

-- Element mass and stiffness matrices
-- Must multiply by the determinant of the Jacobian here!!!
elemMatrices :: (Element.Element e,ShapeFcns.ShapeFcn s,Basis.Basis b,FrElNuFi a) => e a -> s b -> Int -> (Matrix a,Matrix a)
elemMatrices elem shpFcn ngpts = (stiffnessMat, massMat)
  where
    stiffnessMat = integrate 1 ngpts (stiffnessIntegrand elem shpFcn)
    massMat      = integrate 1 ngpts (massIntegrand elem shpFcn)

-- Function which computes the dimensions of the global matrices
globalMatrixDimension :: Int -> Int
globalMatrixDimension numNodes = 2 * numNodes - 2

-- Function which takes a local element matrix along with an element and transforms it to global
localMatToGlobal :: (Element.Element e,FrElNuFi a) => e a -> Matrix a -> [((Int, Int), a)]
localMatToGlobal elem elemMat = concat [[((Mesh.globalNodeNum en i, Mesh.globalNodeNum en j), atIndex elemMat (i,j) ) | i <- [0..dim-1]] | j <- [0..dim-1]]
  where
    en = Element.getElementNumber elem
    dim = (fst . size) elemMat

-- Function which does the same as the above, but with contributions due to fluxes.
localFluxMatToGlobal :: (Element.Element e,FrElNuFi a) => e a -> Matrix a -> [((Int, Int), a)]
localFluxMatToGlobal elem elemMat = concat [[((Mesh.globalNodeNum en i, Mesh.globalNodeNum en j - 1), atIndex elemMat (i,j) ) | i <- [0..dim-1]] | j <- [0..dim]]
  where
    en = Element.getElementNumber elem
    dim = (fst . size) elemMat

-- Function which assembles global stiffness and mass matrices
-- This function takes a mesh and returns a tuple which corresponds to the
-- global stiffness and global mass matrices, respectively.
-- THIS ONLY WORKS WITH THE DISCONTINUOUS MESH...  IT DOES NOT SUM DUPLICATE ENTRIES.
assembleGlobalMatrices :: (Element.Element e, ShapeFcns.ShapeFcn s,Basis.Basis b, FrElNuFi a) => Mesh e a -> s b -> Int -> a -> [[((Int,Int),a)]]
assembleGlobalMatrices grid shpFcn ngpts advSpd = [globalK, globalM, globalF]
  where
    elems = Mesh.getMeshElements grid
    nnodes = length (concatMap Element.getElementNodes elems)
    elemMats = map (\elem -> elemMatrices elem shpFcn ngpts) elems
    elemK = map fst elemMats
    elemM = map snd elemMats
    elemF = elemFluxMatrix advSpd
    globalK = concat (zipWith localMatToGlobal elems elemK)
    globalM = concat (zipWith localMatToGlobal elems elemM)
    globalF = filter (\entry -> (snd . fst) entry >= 0 && (snd . fst) entry < nnodes) (concatMap (`localFluxMatToGlobal` elemF) elems)

-- This function computes the right hand side of the semi-discrete statement.
-- udot_j = (M_ij)^(-1) ( a K_ij - F_ij ) u_j.
-- This is what needs to be integrated point-by-point.
uDot :: (Floating a,Field a,Num (Vector a)) => Matrix a -> Matrix a -> Matrix a -> a -> Vector a -> Vector a
uDot globalM globalK globalF advSpd u = inv globalM #> ((scale advSpd globalK - globalF) #> u)

-- This function advances the solution one step in time.
advance :: (Floating a,Field a,Num (Vector a)) => Matrix a -> Matrix a -> Matrix a -> a -> Vector a -> a -> Vector a
advance globalM globalK globalF advSpd u dt = scale dt (uDot globalM globalK globalF advSpd u - u)

-- This function applies boundary conditions to the semi-discrete system.
applyBCs _ = []

-- This function computes the solution as a function of time.
{-
solveAdvection numElements icFunc tf = do {
  linBasis <- Lagrange 1;
  shpFcn   <- TensorProduct linBasis;
  grid     <- genDMesh 0.0 1.0 numElements StructElem;
  }
-}


-- Exercisee/Ch2/Pt1/Ex7.hs
module Ex7 where

-- Exercise 7
-- Add data types and parsers to support the full numeric tower of Scheme
-- numeric types. Haskell has built-in types to represent many of these; check
-- the Prelude. For others, you can define compound types that represent eg. a
-- `Rational` as a numerator and denominator, or a `Complex` as a real and an
-- imaginary part (each itself a `Real`).

import Data.Char                            (digitToInt, isSpace)
import Data.Complex
import Data.Ratio
import Control.Monad
import Numeric
import System.Environment

import Text.ParserCombinators.Parsec hiding (spaces)

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal -- (a b c ... . z)
             | Number LispNum
             | Character String
             | String String
             | Bool Bool
             deriving Show -- debug purposes

data LispNum = Complex (Complex Float)
             | Real Float
             | Rational (Ratio Integer)
             | Integer Integer
             deriving (Eq, Show)

parseCharacter :: Parser LispVal
parseCharacter = do
  string "#\\"
  character <- many $ satisfy $ not . isSpace
  return $ Character character

parseString :: Parser LispVal
parseString = do
  char '"'
  -- order is important here; need to try parsing an escape sequence first because
  -- otherwise we fail when we reach the second character of the escape
  -- sequence
  x <- many $ choice $ escChars ++ [nonQuote]
  char '"'
  return $ String x
  where
    nonQuote = noneOf "\""
    -- taken from here:
    -- https://en.wikipedia.org/wiki/Escape_sequences_in_C#Table_of_escape_sequences
    -- (excluded characters by oct/hex codes)
    escChars = [ char '\\' >> char x | x <- "abfnrtv\\'\"?" ]

parseBool :: Parser LispVal
parseBool = do
  try (char '#')
  v <- oneOf "tf"
  return $ case v of
    't' -> Bool True
    'f' -> Bool False

parseAtom :: Parser LispVal
parseAtom = do
              first <- letter <|> symbol
              rest <- many (letter <|> digit <|> symbol)
              return $ Atom (first:rest)

parseNumber :: Parser LispVal
parseNumber = do
  n <- (try parseBin) <|> (try parseOct) <|> (try parseHex) <|> (try parseFlt)
    <|> (try parseRat) <|> (try parseCmp) <|> parseDec
  return $ Number n
  where
    parseBin = do
      string "#b"
      binStr <- many1 $ oneOf "01"
      -- cribbed from http://stackoverflow.com/a/26961027/1893155
      let binVal = foldl (\acc x -> acc * 2 + digitToInt x) 0 binStr
      return $ Integer (toInteger binVal)
    parseOct = do
      string "#o"
      octStr <- many1 octDigit
      let octVal = fst $ (readOct octStr) !! 0
      return $ Integer octVal
    parseDec = parseDecNoPre <|> parseDecPre
    parseDecNoPre = do
      decStr <- many1 digit
      return $ (Integer . read) decStr
    parseDecPre = do
      string "#d"
      decStr <- many1 digit
      return $ (Integer . read) decStr
    parseHex = do
      string "#x"
      hexStr <- many1 hexDigit
      let hexVal = fst $ (readHex hexStr) !! 0
      return $ Integer hexVal
    parseFlt = do
      whole <- many1 digit
      string "."
      mantissa <- many1 digit
      let fltVal = fst $ (readFloat (whole ++ "." ++ mantissa) !! 0)
      return $ Real $ fltVal
    parseRat = do
      m <- many1 digit
      many space
      string "/"
      many space
      n <- many1 digit
      let (m', n') = (fst $ (readDec m) !! 0, fst $ (readDec n) !! 0)
      return $ Rational $ m' % n'
    parseCmp = do
      (Real a) <- (try parseDec) <|> parseFlt
      many space
      string "+"
      many space
      (Real b) <- (try parseDec) <|> parseFlt
      char 'i'
      return $ Complex $ a :+ b

-- TODO : I don't like all these `try`s -- better way to do this?
parseExpr :: Parser LispVal
parseExpr = (try parseNumber)
         <|> (try parseString)
         <|> (try parseCharacter)
         <|> (try parseBool)
         <|> (try parseAtom)

symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~"

spaces :: Parser ()
spaces = skipMany1 space

readExpr :: String -> String
readExpr input = case parse parseExpr "lisp" input of
  Left err -> "No match: " ++ show err
  Right val -> "Found value"

main :: IO ()
main = do
  (expr:_) <- getArgs
  putStrLn (readExpr expr)