{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
import           Control.Monad
import           Control.Monad.Random
import           Criterion.Main
import           Data.Default
import           Data.List                           (foldl')
import           Data.Singletons
import           Data.Singletons.Prelude.List
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static        as SA

import           Grenade

import           Grenade.Layers.Internal.Convolution
import           Grenade.Layers.Internal.Pooling

import           Numeric.LinearAlgebra

ff :: SpecNet
ff = specFullyConnected 2 40 |=> specTanh1D 40 |=> netSpecInner |=> specFullyConnected 20 30 |=> specRelu1D 30 |=> specFullyConnected 30 20 |=> specRelu1D 20 |=> specFullyConnected 20 10 |=> specRelu1D 10 |=> specFullyConnected 10 1 |=> specLogit1D 1 |=> specNil1D 1
  where netSpecInner = specFullyConnected 40 30 |=> specRelu1D 30 |=> specFullyConnected 30 20 |=> specNil1D 20

hugeFf :: SpecNet
hugeFf = specFullyConnected 2 4000 |=> specRelu1D 4000 |=> specFullyConnected 4000 1 |=> specLogit1D 1 |=> specNil1D 1

hugeBigFf :: SpecNet
hugeBigFf =
      specFullyConnected 2 4500
  |=> specTanh1D 4500
  |=> specFullyConnected 4500 750
  |=> specRelu1D 750
  |=> specFullyConnected 750 150
  |=> specRelu1D 150
  |=> specFullyConnected 150 150
  |=> specRelu1D 150
  |=> specFullyConnected 150 30
  |=> specRelu1D 30
  |=> specFullyConnected 30 20
  |=> specRelu1D 20
  |=> specFullyConnected 20 10
  |=> specRelu1D 10
  |=> specFullyConnected 10 1
  |=> specLogit1D 1
  |=> specNil1D 1

main :: IO ()
main = do
  putStrLn $ "Benchmarking with type: " ++ nameF
  x :: S ('D2 60 60) <- randomOfShape
  y :: S ('D3 60 60 1) <- randomOfShape
  SpecConcreteNetwork1D1D netFF <- networkFromSpecificationWith def ff
  SpecConcreteNetwork1D1D netHuge <- networkFromSpecificationWith def hugeFf
  SpecConcreteNetwork1D1D netHugeBig <- networkFromSpecificationWith def hugeBigFf
  SpecConcreteNetwork1D1D netHugeBlas <- networkFromSpecificationWith (NetworkInitSettings UniformInit BLAS Nothing) hugeFf
  SpecConcreteNetwork1D1D netHugeBigBlas <- networkFromSpecificationWith (NetworkInitSettings UniformInit BLAS Nothing) hugeBigFf
  defaultMain
    [ -- bgroup
    --     "im2col"
    --     [ bench "im2col 3x4" $ whnf (im2col 2 2 1 1) ((3 >< 4) [1 ..])
    --     , bench "im2col 28x28" $ whnf (im2col 5 5 1 1) ((28 >< 28) [1 ..])
    --     , bench "im2col 100x100" $ whnf (im2col 10 10 1 1) ((100 >< 100) [1 ..])
    --     ]
    -- , bgroup
    --     "col2im"
    --     [ bench "col2im 3x4" $ whnf (col2im 2 2 1 1 3 4) ((6 >< 4) [1 ..])
    --     , bench "col2im 28x28" $ whnf (col2im 5 5 1 1 28 28) ((576 >< 25) [1 ..])
    --     , bench "col2im 100x100" $ whnf (col2im 10 10 1 1 100 100) ((8281 >< 100) [1 ..])
    --     ]
    -- , bgroup
    --     "poolfw"
    --     [ bench "poolforwards 3x4" $ whnf (poolForward 1 3 4 2 2 1 1) ((3 >< 4) [1 ..])
    --     , bench "poolforwards 28x28" $ whnf (poolForward 1 28 28 5 5 1 1) ((28 >< 28) [1 ..])
    --     , bench "poolforwards 100x100" $ whnf (poolForward 1 100 100 10 10 1 1) ((100 >< 100) [1 ..])
    --     ]
    -- , bgroup
    --     "poolbw"
    --     [ bench "poolbackwards 3x4" $ whnf (poolBackward 1 3 4 2 2 1 1 ((3 >< 4) [1 ..])) ((2 >< 3) [1 ..])
    --     , bench "poolbackwards 28x28" $ whnf (poolBackward 1 28 28 5 5 1 1 ((28 >< 28) [1 ..])) ((24 >< 24) [1 ..])
    --     , bench "poolbackwards 100x100" $ whnf (poolBackward 1 100 100 10 10 1 1 ((100 >< 100) [1 ..])) ((91 >< 91) [1 ..])
    --     ]
    -- , bgroup
    --     "padcrop"
    --     [ bench "pad 2D 60x60" $ whnf (testRun2D Pad) x
    --     , bench "pad 3D 60x60" $ whnf (testRun3D Pad) y
    --     , bench "crop 2D 60x60" $ whnf (testRun2D' Crop) x
    --     , bench "crop 3D 60x60" $ whnf (testRun3D' Crop) y
    --     ]
    -- , bgroup
    --     "feedforward SGD"
    --     [ bench "ANN 1000 training steps" $ nfIO $ netTrain netFF defSGD 1000
    --     , bench "ANN 10000 training steps" $ nfIO $ netTrain netFF defSGD 10000
    --     , bench "ANN Huge 100 train steps" $ nfIO $ netTrain netHuge defSGD 100
    --     , bench "ANN Huge 1000 train steps" $ nfIO $ netTrain netHuge defSGD 1000
    --     ]
    -- ,
      bgroup
        "feedforward Adam"
        [ bench "ANN 100   train steps    Network: Huge          Backend: HMatrix" $ nfIO $ netTrain netHuge        defAdam 100
        , bench "ANN 100   train steps    Network: Huge          Backend: BLAS   " $ nfIO $ netTrain netHugeBlas    defAdam 100
        , bench "ANN 100   train steps    Network: Huge-Big      Backend: HMatrix" $ nfIO $ netTrain netHugeBig     defAdam 100
        , bench "ANN 100   train steps    Network: Huge-Big      Backend: BLAS   " $ nfIO $ netTrain netHugeBigBlas defAdam 100
        , bench "ANN 1000  train steps    Network: Huge          Backend: HMatrix" $ nfIO $ netTrain netHuge        defAdam 1000
        , bench "ANN 1000  train steps    Network: Huge          Backend: BLAS   " $ nfIO $ netTrain netHugeBlas    defAdam 1000
        , bench "ANN 1000  train steps    Network: Huge-Big      Backend: HMatrix" $ nfIO $ netTrain netHugeBig     defAdam 1000
        , bench "ANN 1000  train steps    Network: Huge-Big      Backend: BLAS   " $ nfIO $ netTrain netHugeBigBlas defAdam 1000
        , bench "ANN 1000  train steps    Network: Feedforward   Backend: HMatrix" $ nfIO $ netTrain netFF       defAdam 1000
        , bench "ANN 10000 train steps    Network: Feedforward   Backend: HMatrix" $ nfIO $ netTrain netFF       defAdam 10000
        ]
    ]
  putStrLn $ "Benchmarked with type: " ++ nameF

testRun2D :: Pad 1 1 1 1 -> S ('D2 60 60) -> S ('D2 62 62)
testRun2D = snd ... runForwards

testRun3D :: Pad 1 1 1 1 -> S ('D3 60 60 1) -> S ('D3 62 62 1)
testRun3D = snd ... runForwards

testRun2D' :: Crop 1 1 1 1 -> S ('D2 60 60) -> S ('D2 58 58)
testRun2D' = snd ... runForwards

testRun3D' :: Crop 1 1 1 1 -> S ('D3 60 60 1) -> S ('D3 58 58 1)
testRun3D' = snd ... runForwards

(...) :: (a -> b) -> (c -> d -> a) -> c -> d -> b
(...) = (.) . (.)

netTrain ::
     (SingI (Last shapes), MonadRandom m, KnownNat len1, KnownNat len2, Head shapes ~ 'D1 len1, Last shapes ~ 'D1 len2)
  => Network layers shapes
  -> Optimizer o
  -> Int
  -> m (Network layers shapes)
netTrain net0 op n = do
  inps <-
    replicateM n $ do
      s <- getRandom
      return $ S1D $ SA.randomVector s SA.Uniform * 2 - 1
  let outs =
        flip map inps $ \(S1D v) ->
          if v `inCircle` (fromRational 0.50, 0.50) || v `inCircle` (fromRational (-0.50), 0.50)
            then S1D $ fromRational 1
            else S1D $ fromRational 0
  let trained = foldl' trainEach net0 (zip inps outs)
  return trained
  where
    trainEach !network (i, o) = train op network i o
    inCircle :: KnownNat n => SA.R n -> (SA.R n, RealNum) -> Bool
    v `inCircle` (o, r) = SA.norm_2 (v - o) <= r


{-# LANGUAGE BangPatterns, DeriveDataTypeable, DeriveGeneric #-}

-- |
-- Module    : Statistics.Resampling
-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Resampling statistics.

module Statistics.Resampling
    (
      Resample(..)
    , jackknife
    , jackknifeMean
    , jackknifeVariance
    , jackknifeVarianceUnb
    , jackknifeStdDev
    , resample
    , estimate
    , splitGen
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Concurrent (forkIO, newChan, readChan, writeChan)
import Control.Monad (forM_, liftM, replicateM, replicateM_)
import Data.Binary (Binary(..))
import Data.Data (Data, Typeable)
import Data.Vector.Algorithms.Intro (sort)
import Data.Vector.Binary ()
import Data.Vector.Generic (unsafeFreeze)
import Data.Word (Word32)
import GHC.Conc (numCapabilities)
import GHC.Generics (Generic)
import Numeric.Sum (Summation(..), kbn)
import Statistics.Function (indices)
import Statistics.Sample (mean, stdDev, variance, varianceUnbiased)
import Statistics.Types (Estimator(..), Sample)
import System.Random.MWC (GenIO, initialize, uniform, uniformVector)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU

-- | A resample drawn randomly, with replacement, from a set of data
-- points.  Distinct from a normal array to make it harder for your
-- humble author's brain to go wrong.
newtype Resample = Resample {
      fromResample :: U.Vector Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON Resample
instance ToJSON Resample

instance Binary Resample where
    put = put . fromResample
    get = fmap Resample get

-- | /O(e*r*s)/ Resample a data set repeatedly, with replacement,
-- computing each estimate over the resampled data.
--
-- This function is expensive; it has to do work proportional to
-- /e*r*s/, where /e/ is the number of estimation functions, /r/ is
-- the number of resamples to compute, and /s/ is the number of
-- original samples.
--
-- To improve performance, this function will make use of all
-- available CPUs.  At least with GHC 7.0, parallel performance seems
-- best if the parallel garbage collector is disabled (RTS option
-- @-qg@).
resample :: GenIO
         -> [Estimator]         -- ^ Estimation functions.
         -> Int                 -- ^ Number of resamples to compute.
         -> Sample              -- ^ Original sample.
         -> IO [Resample]
resample gen ests numResamples samples = do
  let !numSamples = U.length samples
      ixs = scanl (+) 0 $
            zipWith (+) (replicate numCapabilities q)
                        (replicate r 1 ++ repeat 0)
          where (q,r) = numResamples `quotRem` numCapabilities
  results <- mapM (const (MU.new numResamples)) ests
  done <- newChan
  gens <- splitGen numCapabilities gen
  forM_ (zip3 ixs (tail ixs) gens) $ \ (start,!end,gen') -> do
    forkIO $ do
      let loop k ers | k >= end = writeChan done ()
                     | otherwise = do
            re <- U.replicateM numSamples $ do
                    r <- uniform gen'
                    return (U.unsafeIndex samples (r `mod` numSamples))
            forM_ ers $ \(est,arr) ->
                MU.write arr k . est $ re
            loop (k+1) ers
      loop start (zip ests' results)
  replicateM_ numCapabilities $ readChan done
  mapM_ sort results
  mapM (liftM Resample . unsafeFreeze) results
 where
  ests' = map estimate ests

-- | Run an 'Estimator' over a sample.
estimate :: Estimator -> Sample -> Double
estimate Mean             = mean
estimate Variance         = variance
estimate VarianceUnbiased = varianceUnbiased
estimate StdDev           = stdDev
estimate (Function est) = est

-- | /O(n) or O(n^2)/ Compute a statistical estimate repeatedly over a
-- sample, each time omitting a successive element.
jackknife :: Estimator -> Sample -> U.Vector Double
jackknife Mean sample             = jackknifeMean sample
jackknife Variance sample         = jackknifeVariance sample
jackknife VarianceUnbiased sample = jackknifeVarianceUnb sample
jackknife StdDev sample = jackknifeStdDev sample
jackknife (Function est) sample
  | G.length sample == 1 = singletonErr "jackknife"
  | otherwise            = U.map f . indices $ sample
  where f i = est (dropAt i sample)

-- | /O(n)/ Compute the jackknife mean of a sample.
jackknifeMean :: Sample -> U.Vector Double
jackknifeMean samp
  | len == 1  = singletonErr "jackknifeMean"
  | otherwise = G.map (/l) $ G.zipWith (+) (pfxSumL samp) (pfxSumR samp)
  where
    l   = fromIntegral (len - 1)
    len = G.length samp

-- | /O(n)/ Compute the jackknife variance of a sample with a
-- correction factor @c@, so we can get either the regular or
-- \"unbiased\" variance.
jackknifeVariance_ :: Double -> Sample -> U.Vector Double
jackknifeVariance_ c samp
  | len == 1  = singletonErr "jackknifeVariance"
  | otherwise = G.zipWith4 go als ars bls brs
  where
    als = pfxSumL . G.map goa $ samp
    ars = pfxSumR . G.map goa $ samp
    goa x = v * v where v = x - m
    bls = pfxSumL . G.map (subtract m) $ samp
    brs = pfxSumR . G.map (subtract m) $ samp
    m = mean samp
    n = fromIntegral len
    go al ar bl br = (al + ar - (b * b) / q) / (q - c)
      where b = bl + br
            q = n - 1
    len = G.length samp

-- | /O(n)/ Compute the unbiased jackknife variance of a sample.
jackknifeVarianceUnb :: Sample -> U.Vector Double
jackknifeVarianceUnb = jackknifeVariance_ 1

-- | /O(n)/ Compute the jackknife variance of a sample.
jackknifeVariance :: Sample -> U.Vector Double
jackknifeVariance = jackknifeVariance_ 0

-- | /O(n)/ Compute the jackknife standard deviation of a sample.
jackknifeStdDev :: Sample -> U.Vector Double
jackknifeStdDev = G.map sqrt . jackknifeVarianceUnb

pfxSumL :: U.Vector Double -> U.Vector Double
pfxSumL = G.map kbn . G.scanl add zero

pfxSumR :: U.Vector Double -> U.Vector Double
pfxSumR = G.tail . G.map kbn . G.scanr (flip add) zero

-- | Drop the /k/th element of a vector.
dropAt :: U.Unbox e => Int -> U.Vector e -> U.Vector e
dropAt n v = U.slice 0 n v U.++ U.slice (n+1) (U.length v - n - 1) v

singletonErr :: String -> a
singletonErr func = error $
                    "Statistics.Resampling." ++ func ++ ": singleton input"

-- | Split a generator into several that can run independently.
splitGen :: Int -> GenIO -> IO [GenIO]
splitGen n gen
  | n <= 0    = return []
  | otherwise =
  fmap (gen:) . replicateM (n-1) $
  initialize =<< (uniformVector gen 256 :: IO (U.Vector Word32))


{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}

-----------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Involutive rings
----------------------------------------------------------------------------
module Linear.Conjugate
  ( Conjugate(..)
  , TrivialConjugate
  ) where

import Data.Complex hiding (conjugate)
import Data.Int
import Data.Word
import Foreign.C.Types (CFloat, CDouble)

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

-- | An involutive ring
class Num a => Conjugate a where
  -- | Conjugate a value. This defaults to the trivial involution.
  --
  -- >>> conjugate (1 :+ 2)
  -- 1.0 :+ (-2.0)
  --
  -- >>> conjugate 1
  -- 1
  conjugate :: a -> a
#ifndef HLINT
  default conjugate :: TrivialConjugate a => a -> a
  conjugate = id
#endif

-- | Requires and provides a default definition such that
--
-- @
-- 'conjugate' = 'id'
-- @
class Conjugate a => TrivialConjugate a

instance Conjugate Integer
instance Conjugate Int
instance Conjugate Int64
instance Conjugate Int32
instance Conjugate Int16
instance Conjugate Int8
instance Conjugate Word
instance Conjugate Word64
instance Conjugate Word32
instance Conjugate Word16
instance Conjugate Word8
instance Conjugate Double
instance Conjugate Float
instance Conjugate CFloat
instance Conjugate CDouble

instance (Conjugate a, RealFloat a) => Conjugate (Complex a) where
  {-# SPECIALIZE instance Conjugate (Complex Float) #-}
  {-# SPECIALIZE instance Conjugate (Complex Double) #-}
  conjugate (a :+ b) = conjugate a :+ negate b

instance TrivialConjugate Integer
instance TrivialConjugate Int
instance TrivialConjugate Int64
instance TrivialConjugate Int32
instance TrivialConjugate Int16
instance TrivialConjugate Int8
instance TrivialConjugate Word
instance TrivialConjugate Word64
instance TrivialConjugate Word32
instance TrivialConjugate Word16
instance TrivialConjugate Word8
instance TrivialConjugate Double
instance TrivialConjugate Float
instance TrivialConjugate CFloat
instance TrivialConjugate CDouble


module MachineLearning.NeuralNetwork.TopologyTest
(
  tests
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Approx
import Test.HUnit.Plus
import qualified Numeric.LinearAlgebra as LA
import MachineLearning.NeuralNetwork.Topology
import qualified MachineLearning.NeuralNetwork.TopologyMaker as TM

nnt = TM.makeTopology TM.ASigmoid TM.LLogistic 15 2 [10]

flattenTest = do
  theta <- initializeThetaIO nnt
  let theta' = flatten $ unflatten nnt theta
      norm = LA.norm_2 (theta - theta')
  assertApproxEqual "flatten" 1e-10 0 norm

tests = [ testGroup "flatten" [
            testCase "flatten" flattenTest
            ]
        ]


module Jumps (findJumps,getJumps,flt2dbl) where

import Prelude hiding (filter,length,replicate,take,(++),splitAt)

import Data.Function (on)
import Data.Ord (comparing)
import Data.Foldable (foldMap')
import Data.Bifunctor (bimap)
import Data.Vector.Unboxed (filter,slice,length,convert,splitAt)
import Data.Vector (replicate,take,force,modify,(++),maxIndexBy)
import Data.Vector.Mutable (write)

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

import Numeric.Sum (sumVector,kbn)
import Statistics.Sample (mean)
import GHC.Float (float2Double)

data Partition = Partition
    { start :: Int
    , num :: Int
    , level :: Double
    , _loss :: Double
    , _split :: Int
    , diff :: Double
    , _levels :: (Double,Double)
    , _losses :: (Double,Double)
    } deriving Show

data Split = Split
    { _num :: Int
    , _lvls :: (Double,Double)
    , _lsses :: (Double,Double)
    , loss :: Double
    }

instance Eq Split where
    (==) = (==) `on` loss

instance Ord Split where
    (<=) = (<=) `on` loss

instance Semigroup Split where
    (<>) = min

instance Monoid Split where
    mempty = Split 0 (0,0) (0,0) (1 / 0) -- 1 / 0 = Infinity

mkPart :: U.Vector Double -> Int -> Double -> Double -> Partition
mkPart xs a lvl lss
    | m <= 1 = Partition a m lvl lss a 0 (0,0) (0,0)
    | otherwise =
        case flip foldMap' [1..m-1] $ \i ->
            let filt = filter (not . isNaN)
                (x1,x2) = bimap filt filt $ splitAt i xs
                m1 = mean x1
                m2 = mean x2
                s1 = sumVector kbn $ U.map (\x -> (x - m1) ^ (2 :: Int)) x1
                s2 = sumVector kbn $ U.map (\x -> (x - m2) ^ (2 :: Int)) x2
            in Split i (m1,m2) (s1,s2) (s1 + s2)
        of Split n lvls lses lss' -> Partition a m lvl lss n (lss - lss') lvls lses
    where m = length xs

splitPart :: U.Vector Double -> Partition -> (Partition,Partition)
splitPart inData (Partition a m _ _ n _ (lvl1,lvl2) (lss1,lss2)) =
    ( mkPart (slice a n inData) a lvl1 lss1
    , mkPart (slice (a+n) (m-n) inData) (a+n) lvl2 lss2
    )

findJumps :: U.Vector Double -> V.Vector Partition
findJumps inData =
    let noNaN = filter (not . isNaN) inData
        lev = mean noNaN
        los = sumVector kbn $ U.map (\x -> (x - lev) ^ (2 :: Int)) noNaN
        vec = replicate 50 $ mkPart inData 0 lev los
    in loop vec 1 0 (diff $ vec V.! 0)
    where
        loop :: V.Vector Partition -> Int -> Int -> Double -> V.Vector Partition
        loop parts k splitInd lossDiff
            | lossDiff <= 120 = force $ take k parts
            | k `mod` 50 == 0 = loop (parts ++ take 50 parts) k splitInd lossDiff
            | otherwise =
                let (x,y) = splitPart inData $ parts V.! splitInd
                    parts' = modify (\v -> write v splitInd x >> write v k y) parts
                    splitInd' = maxIndexBy (comparing diff) parts'
                    lossDiff' = diff $ parts' V.! splitInd'
                in loop parts' (k + 1) splitInd' lossDiff'

getJumps :: V.Vector Partition -> U.Vector (Double,Int,Int)
getJumps = convert . (V.map $ (,,) <$> level <*> start <*> num)

flt2dbl :: U.Vector Float -> U.Vector Double
flt2dbl = U.map float2Double


{-# language BangPatterns #-}
{-# language ConstraintKinds #-}
{-# language DefaultSignatures #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language RankNTypes #-}
{-# language RebindableSyntax #-}
{-# language TypeFamilies #-}
module Prob where

import Prelude hiding (Functor(..), Applicative(..), Monad(..), sequence, (=<<), mapM, mapM_)
import qualified Prelude

import GHC.Exts (Constraint)

import Control.DeepSeq (NFData(..), force)
import Control.Parallel.Strategies (withStrategy, rpar, rseq, parList, parListChunk)
import Control.Monad.State (State)

import Data.List (groupBy, sort, sortBy)
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.Random                          as R
import qualified Data.Random.Distribution.Categorical as R

import Numeric.SpecFunctions (choose)

type QQ = Rational

data Event a = Event !a !QQ deriving (Eq, Ord)

eventToPair :: Event a -> (a, QQ)
eventToPair (Event a p) = (a, p)
{-# inline eventToPair #-}

mapEvent :: (a -> b) -> Event a -> Event b
mapEvent f (Event a p) = Event (f a) p
{-# inline mapEvent #-}

instance NFData a => NFData (Event a) where
    rnf (Event a _) = rnf a

instance Show a => Show (Event a) where
    show (Event a p) = show a ++ ": " ++ show p -- (realToFrac p :: Float)

newtype Prob a = Prob { density :: [Event a] } deriving (Eq, Ord, Show, NFData)

events :: Prob a -> [Event a]
events (Prob es) = es

eventValues :: Prob a -> [a]
eventValues prob = [a | Event a _ <- events prob]

eventProbs :: Prob a -> [QQ]
eventProbs prob = [p | Event _ p <- events prob]

class NoConstr a
instance NoConstr a

class ConstrMonad m where
    type Constr m :: * -> Constraint
    type Constr m = NoConstr

    fmap :: Constr m b => (a -> b) -> m a -> m b
    fmap f ma = do
        a <- ma
        return (f a)

    return :: Constr m a => a -> m a
    (>>=)  :: Constr m b => m a -> (a -> m b) -> m b

    (|>>=|) :: (Constr m a, Constr m b) => m a -> (a -> m b) -> m b
    (|>>=|) = (>>=)

    default return :: Prelude.Monad m => a -> m a
    return = Prelude.return
    default (>>=) :: Prelude.Monad m => m a -> (a -> m b) -> m b
    ma >>= f = ma Prelude.>>= f

infixl 1 >>=

liftA2 :: (ConstrMonad m, Constr m a, Constr m b, Constr m d) => (a -> b -> d) -> m a -> m b -> m d
liftA2 f ma mb = do
    a <- ma
    b <- mb
    return (f a b)
{-# inline liftA2 #-}

liftA3 :: (ConstrMonad m, Constr m a, Constr m b, Constr m d, Constr m e) => (a -> b -> d -> e) -> m a -> m b -> m d -> m e
liftA3 f ma mb md = do
    a <- ma
    b <- mb
    d <- md
    return (f a b d)
{-# inlinable liftA3 #-}

(=<<)  :: (ConstrMonad m, Constr m a, Constr m b) => (a -> m b) -> m a -> m b
f =<< ma = ma >>= f
infixr 1 =<<
{-# inline (=<<) #-}

instance ConstrMonad Maybe
instance ConstrMonad IO
instance ConstrMonad (State s)

instance Prelude.Monad m => ConstrMonad (R.RVarT m)
instance Prelude.Monad m => ProbMonad (R.RVarT m) where
    fromEventsSorted evts = R.categoricalT [(realToFrac p :: Double, a) | Event a p <- evts]

weightedEvents :: Real b => [(a, b)] -> [Event a]
weightedEvents es = [Event a (realToFrac w / s) | (a,w) <- es]
  where
    s = realToFrac $ sum [w | (_,w) <- es]

rnfList :: [a] -> ()
rnfList []     = ()
rnfList (a:as) = a `seq` rnfList as

forceList :: [a] -> [a]
forceList as = rnfList as `seq` as

fmapProb :: Ord b => (a -> b) -> Prob a -> Prob b
fmapProb f (Prob [Event a p]) = Prob [Event (f a) p]
fmapProb f (Prob evts)        =
    Prob $ forceList $ group1 $ sortBy (\(Event b _) (Event b' _) -> compare b b') [Event (f a) p | Event a p <- evts]
  where
    group1 :: Ord a => [Event a] -> [Event a]
    group1 []     = []
    group1 (e:es) = group e es

    group :: Ord a => Event a -> [Event a] -> [Event a]
    group e@(Event a p) ees =
      case ees of
        [] -> [e]
        e'@(Event a' p') : es
          | a == a'   -> group (Event a (p + p')) es
          | otherwise -> e : group e' es

{-# inlinable fmapProb #-}

bindProbWithStrat :: Ord b => (forall c. [c] -> [c]) -> Prob a -> (a -> Prob b) -> Prob b
bindProbWithStrat evalList (Prob [Event a _]) f = Prob $ evalList $ events (f a)
bindProbWithStrat evalList (Prob evts) f =
    Prob $ forceList $ group1 $ merge1 $ evalList $ totalProb [Event (f a) p | Event a p <- evts]
  where
    totalProb :: [Event (Prob a)] -> [[Event a]]
    totalProb ess = [[Event b (p*p') | Event b p' <- es] | Event (Prob es) p <- ess]

    merge1 :: Ord a => [[Event a]] -> [Event a]
    merge1 []       = []
    merge1 (es:ess) = merge es ess

    merge :: Ord a => [Event a] -> [[Event a]] -> [Event a]
    merge es []             = es
    merge es (es':[])       = merge2 es es'
    merge es (es':es'':ess) = merge2 (merge2 es es') (merge es'' ess)

    merge2 :: Ord a => [Event a] -> [Event a] -> [Event a]
    merge2 as         []         = as
    merge2 []         bs         = bs
    merge2 aas@(a:as) bbs@(b:bs)
      | a <= b    = a : merge2 as bbs
      | otherwise = b : merge2 aas bs

    group1 :: Ord a => [Event a] -> [Event a]
    group1 []     = []
    group1 (e:es) = group e es

    group :: Ord a => Event a -> [Event a] -> [Event a]
    group e@(Event a p) ees =
      case ees of
        [] -> [e]
        e'@(Event a' p') : es
          | a == a'   -> group (Event a (p + p')) es
          | otherwise -> e : group e' es

{-# inlinable bindProbWithStrat #-}

instance ConstrMonad Prob where
    type Constr Prob = Ord

    fmap = fmapProb
    {-# inline fmap #-}

    return a = Prob [Event a 1]
    {-# inline return #-}

    (>>=) = bindProbWithStrat forceList
    {-# inline (>>=) #-}

    -- (|>>=|) :: (Ord a, Ord b) => Prob a -> (a -> Prob b) -> Prob b
    (|>>=|) = bindProbWithStrat (withStrategy $ parList rseq)
    {-# inline (|>>=|) #-}

(|=<<|) :: (Ord a, Ord b) => (a -> Prob b) -> Prob a -> Prob b
f |=<<| ma = ma |>>=| f
{-# inline (|=<<|) #-}

ifThenElse :: Bool -> a -> a -> a
ifThenElse True  t _ = t
ifThenElse False _ f = f
{-# inline ifThenElse #-}

fail :: a
fail = error "fail called"

(>>) :: (ConstrMonad m, Constr m (), Constr m b) => m a -> m b -> m b
ma >> mb = fmap (\_ -> ()) ma >>= \() -> mb
infixl 1 >>
{-# inline (>>) #-}

(<<) :: (ConstrMonad m, Constr m (), Constr m b) => m b -> m a -> m b
mb << ma = ma >> mb
infixr 1 <<
{-# inline (<<) #-}

sequence :: (ConstrMonad m, Constr m a, Constr m [a]) => [m a] -> m [a]
sequence []       = return []
sequence (ma:mas) = liftA2 (:) ma (sequence mas)
{-# inlinable sequence #-}

mapM :: (ConstrMonad m, Constr m b, Constr m [b]) => (a -> m b) -> [a] -> m [b]
mapM f []     = return []
mapM f (a:as) = liftA2 (:) (f a) (mapM f as)
{-# inlinable mapM #-}

mapM_ :: (ConstrMonad m, Constr m ()) => (a -> m ()) -> [a] -> m ()
mapM_ f []     = return ()
mapM_ f (a:as) = f a >> mapM_ f as
{-# inlinable mapM_ #-}

forM :: (ConstrMonad m, Constr m b, Constr m [b]) => [a] -> (a -> m b) -> m [b]
forM as f = mapM f as
{-# inline forM #-}

forM_ :: (ConstrMonad m, Constr m ()) => [a] -> (a -> m ()) -> m ()
forM_ as f = mapM_ f as
{-# inline forM_ #-}

sequence1 :: (ConstrMonad m, Constr m a, Constr m [a], Constr m (NonEmpty a)) => NonEmpty (m a) -> m (NonEmpty a)
sequence1 (ma:|mas) = liftA2 (:|) ma (sequence mas)
{-# inline sequence1 #-}

mapM1 :: (ConstrMonad m, Constr m b, Constr m [b], Constr m (NonEmpty b)) => (a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM1 f (a:|as) = liftA2 (:|) (f a) (mapM f as)
{-# inline mapM1 #-}

forM1 :: (ConstrMonad m, Constr m b, Constr m [b], Constr m (NonEmpty b)) => NonEmpty a -> (a -> m b) -> m (NonEmpty b)
forM1 as f = mapM1 f as
{-# inline forM1 #-}

liftFoldl' :: (ConstrMonad m, Constr m a, Constr m b) => (b -> a -> b) -> m b -> [m a] -> m b
liftFoldl' _ !z []     = z
liftFoldl' f !z (p:ps) = liftFoldl' f (liftA2 f z p) ps

liftFoldr :: (ConstrMonad m, Constr m a, Constr m b) => (a -> b -> b) -> m b -> [m a] -> m b
liftFoldr f z = foldr (liftA2 f) z

liftFold :: (ConstrMonad m, Constr m a, Monoid a) => [m a] -> m a
liftFold = liftFoldl' mappend (return mempty)

liftSum :: (ConstrMonad m, Constr m a, Num a) => [m a] -> m a
liftSum = liftFoldl' (+) (return 0)

class ConstrMonad m => ProbMonad m where
    fromEventsSorted :: Constr m a => [Event a] -> m a

instance ProbMonad Prob where
    fromEventsSorted = Prob

uniformly :: (ProbMonad m, Constr m a, Ord a) => [a] -> m a
uniformly = uniformlySorted . sort

uniformlySorted :: (ProbMonad m, Constr m a) => [a] -> m a
uniformlySorted as = fromEventsSorted [Event a (1/n) | a <- as]
  where
    n = fromIntegral (length as)

probTrue :: Prob Bool -> QQ
probTrue prob =
    case events prob of
        [Event False _, Event True p] -> p
        [Event True p]                -> p
        [Event False q]               -> 1 - q
        []                            -> error "empty Prob!"
        _                             -> error "unnormalized Prob!"

probFalse :: Prob Bool -> QQ
probFalse prob = 1 - probTrue prob

eventProb :: Eq a => a -> Prob a -> QQ
eventProb e prob =
    case [p | Event a p <- events prob, a == e] of
      []  -> 0
      [p] -> p
      _   -> error "eventProb: prob is not normalized"

probOf :: (a -> Bool) -> Prob a -> QQ
probOf test prob = sum [p | Event a p <- events prob, test a]

given :: (a -> Bool) -> Prob a -> Prob a
given hyp prob =
    Prob [Event a (p / probHyp) | Event a p <- events prob, hyp a]
  where
    probHyp = probOf hyp prob

bernoulli :: QQ -> Prob Bool
bernoulli p = Prob [Event False (1-p), Event True p]

binomial :: Int -> QQ -> Prob Int
binomial 0 _ = return 0
binomial n p =
    case p of
      0 -> force $ return 0
      1 -> force $ return n
      _ -> force $ Prob [Event k (binomialProb n p k) | k <- [0..n]]

binomialProb :: Int -> QQ -> Int -> QQ
binomialProb n p k
  | k < 0 || k > n = 0
  | n == 0         = 1
  | otherwise      = realToFrac (choose n k) * p^k * (1-p)^fromIntegral (n-k)

distribution :: Ord a => Prob a -> [Event a]
distribution (Prob evts) = scanl1 sumEvents evts
  where
    sumEvents (Event _ p) (Event a' p') = Event a' (p + p')

revDistribution :: Ord a => Prob a -> [Event a]
revDistribution (Prob evts) = scanr1 sumEvents evts
  where
    sumEvents (Event a p) (Event _ p') = Event a (p + p')

summary :: Prob QQ -> IO ()
summary p =
  let mu    = "µ=" ++ show (realToFrac (mean p) :: Double)
      sigma = "σ=" ++ show (sqrt $ realToFrac (variance p) :: Double)
  in
      putStrLn $ mu ++ ", " ++ sigma

summaryInt :: Prob Int -> IO ()
summaryInt = summary . fmap fromIntegral

mean :: Prob QQ -> QQ
mean prob = sum [k * p | Event k p <- events prob]

variance :: Prob QQ -> QQ
variance prob = sum [k^2 * p | Event k p <- events prob] - mean prob^2

maxEvent :: Ord a => Prob a -> a
maxEvent (Prob es) = maximum [a | Event a _ <- es]


{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
module BenchMnistTools where

import Prelude

import           Control.Arrow ((***))
import           Control.DeepSeq (NFData)
import           Criterion.Main
import           Data.List.Index (imap)
import           Data.Proxy (Proxy (Proxy))
import qualified Data.Vector.Generic as V
import qualified Numeric.LinearAlgebra as HM
import           System.Random

import HordeAd
import HordeAd.Tool.MnistTools

mnistTrainBench2 :: forall r. ( NFData r, HasDelta r
                              , UniformRange r )
                 => String -> Int -> [MnistData r] -> Int -> Int
                 -> r
                 -> Benchmark
mnistTrainBench2 extraPrefix chunkLength xs widthHidden widthHidden2 gamma = do
  let nParams0 = fcnnMnistLen0 widthHidden widthHidden2
      params0Init = V.unfoldrExactN nParams0 (uniformR (-0.5, 0.5))
                    $ mkStdGen 33
      f = fcnnMnistLoss0 widthHidden widthHidden2
      chunk = take chunkLength xs
      grad c = fst $ sgd gamma f c (params0Init, V.empty, V.empty, V.empty)
      name = "" ++ extraPrefix
             ++ unwords [ "s" ++ show nParams0, "v0"
                        , "m0" ++ "=" ++ show nParams0 ]
  bench name $ nf grad chunk

mnistTestBench2
  :: forall r. (UniformRange r, HasDelta r)
  => String -> Int -> [MnistData r] -> Int -> Int -> Benchmark
mnistTestBench2 extraPrefix chunkLength xs widthHidden widthHidden2 = do
  let nParams0 = fcnnMnistLen0 widthHidden widthHidden2
      params0Init = V.unfoldrExactN nParams0 (uniformR (-0.5, 0.5))
                    $ mkStdGen 33
      chunk = take chunkLength xs
      score c = fcnnMnistTest0 (Proxy @r) widthHidden widthHidden2 c params0Init
      name = "test " ++ extraPrefix
             ++ unwords [ "s" ++ show nParams0, "v0"
                        , "m0" ++ "=" ++ show nParams0 ]
  bench name $ whnf score chunk

mnistTrainBGroup2 :: [MnistData Double] -> Int -> Benchmark
mnistTrainBGroup2 xs0 chunkLength =
  env (return $ take chunkLength xs0) $
  \ xs ->
  bgroup ("2-hidden-layer MNIST nn with samples: " ++ show chunkLength)
    [ mnistTestBench2 "30|10 " chunkLength xs 30 10  -- toy width
    , mnistTrainBench2 "30|10 " chunkLength xs 30 10 0.02
    , mnistTestBench2 "300|100 " chunkLength xs 300 100  -- ordinary width
    , mnistTrainBench2 "300|100 " chunkLength xs 300 100 0.02
    , mnistTestBench2 "500|150 " chunkLength xs 500 150  -- another common size
    , mnistTrainBench2 "500|150 " chunkLength xs 500 150 0.02
    ]

mnistTrainBGroup2500 :: [MnistData Double] -> Int -> Benchmark
mnistTrainBGroup2500 xs0 chunkLength =
  env (return (xs0, map (V.map realToFrac *** V.map realToFrac)
                    $ take chunkLength xs0)) $
  \ ~(xs, xsFloat) ->
  bgroup ("huge 2-hidden-layer MNIST nn with samples: " ++ show chunkLength)
    [ mnistTestBench2 "2500|750 " chunkLength xs 2500 750
        -- probably mostly wasted
    , mnistTrainBench2 "2500|750 " chunkLength xs 2500 750 0.02
    , mnistTestBench2 "(Float) 2500|750 " chunkLength xsFloat 2500 750
        -- Float test
    , mnistTrainBench2 "(Float) 2500|750 " chunkLength xsFloat 2500 750
                       (0.02 :: Float)
    ]

mnistTrainBench2V :: String -> Int -> [MnistData Double]
                  -> Int -> Int -> Double
                  -> Benchmark
mnistTrainBench2V extraPrefix chunkLength xs widthHidden widthHidden2 gamma = do
  let (nParams0, nParams1, _, _) = fcnnMnistLen1 widthHidden widthHidden2
      params0Init = HM.randomVector 33 HM.Uniform nParams0 - HM.scalar 0.5
      params1Init = V.fromList $
        imap (\i nPV -> HM.randomVector (33 + nPV + i) HM.Uniform nPV
                        - HM.scalar 0.5)
             nParams1
      f = fcnnMnistLoss1 widthHidden widthHidden2
      chunk = take chunkLength xs
      grad c = fst $ sgd gamma f c (params0Init, params1Init, V.empty, V.empty)
      totalParams = nParams0 + sum nParams1
      name = "" ++ extraPrefix
             ++ unwords [ "s" ++ show nParams0, "v" ++ show (length nParams1)
                        , "m0" ++ "=" ++ show totalParams ]
  bench name $ nf grad chunk

mnistTestBench2V :: String -> Int -> [MnistData Double] -> Int -> Int
                 -> Benchmark
mnistTestBench2V extraPrefix chunkLength xs widthHidden widthHidden2 = do
  let (nParams0, nParams1, _, _) = fcnnMnistLen1 widthHidden widthHidden2
      params0Init = HM.randomVector 33 HM.Uniform nParams0 - HM.scalar 0.5
      params1Init = V.fromList $
        imap (\i nPV -> HM.randomVector (33 + nPV + i) HM.Uniform nPV
                        - HM.scalar 0.5)
             nParams1
      chunk = take chunkLength xs
      score c = fcnnMnistTest1 widthHidden widthHidden2 c
                           (params0Init, params1Init)
      totalParams = nParams0 + sum nParams1
      name = "test " ++ extraPrefix
             ++ unwords [ "s" ++ show nParams0, "v" ++ show (length nParams1)
                        , "m0" ++ "=" ++ show totalParams ]
  bench name $ whnf score chunk

mnistTrainBGroup2V :: [MnistData Double] -> Int -> Benchmark
mnistTrainBGroup2V xs0 chunkLength =
  env (return $ take chunkLength xs0) $
  \ xs ->
  bgroup ("2-hidden-layer V MNIST nn with samples: " ++ show chunkLength)
    [ mnistTestBench2V "30|10 " chunkLength xs 30 10  -- toy width
    , mnistTrainBench2V "30|10 " chunkLength xs 30 10 0.02
    , mnistTestBench2V "300|100 " chunkLength xs 300 100  -- ordinary width
    , mnistTrainBench2V "300|100 " chunkLength xs 300 100 0.02
    , mnistTestBench2V "500|150 " chunkLength xs 500 150  -- another common size
    , mnistTrainBench2V "500|150 " chunkLength xs 500 150 0.02
    ]

mnistTrainBench2L :: String -> Int -> [MnistData Double] -> Int -> Int
                  -> Double
                  -> Benchmark
mnistTrainBench2L extraPrefix chunkLength xs widthHidden widthHidden2 gamma = do
  let ((nParams0, nParams1, nParams2, _), totalParams, _reach, parameters0) =
        initializerFixed 33 0.5 (fcnnMnistLen2 widthHidden widthHidden2)
      -- Using the fused version to benchmark against the manual gradient
      -- from backprop that uses it at least in its forward pass,
      -- not againts the derived gradients that are definitively slower.
      f = fcnnMnistLossFused2
      chunk = take chunkLength xs
      grad c = fst $ sgd gamma f c parameters0
      name = "" ++ extraPrefix
             ++ unwords [ "s" ++ show nParams0, "v" ++ show nParams1
                        , "m" ++ show nParams2
                          ++ "=" ++ show totalParams ]
  bench name $ nf grad chunk

mnistTestBench2L :: String -> Int -> [MnistData Double] -> Int -> Int
                 -> Benchmark
mnistTestBench2L extraPrefix chunkLength xs widthHidden widthHidden2 = do
  let ((nParams0, nParams1, nParams2, _), totalParams, _reach, parameters0) =
        initializerFixed 33 0.5 (fcnnMnistLen2 widthHidden widthHidden2)
      chunk = take chunkLength xs
      score c = fcnnMnistTest2 c parameters0
      name = "test " ++ extraPrefix
             ++ unwords [ "s" ++ show nParams0, "v" ++ show nParams1
                        , "m" ++ show nParams2
                          ++ "=" ++ show totalParams ]
  bench name $ whnf score chunk

mnistTrainBGroup2L :: [MnistData Double] -> Int -> Benchmark
mnistTrainBGroup2L xs0 chunkLength =
  env (return $ take chunkLength xs0) $
  \ xs ->
  bgroup ("2-hidden-layer L MNIST nn with samples: " ++ show chunkLength)
    [ mnistTestBench2L "30|10 " chunkLength xs 30 10  -- toy width
    , mnistTrainBench2L "30|10 " chunkLength xs 30 10 0.02
    , mnistTestBench2L "300|100 " chunkLength xs 300 100  -- ordinary width
    , mnistTrainBench2L "300|100 " chunkLength xs 300 100 0.02
    , mnistTestBench2L "500|150 " chunkLength xs 500 150  -- another common size
    , mnistTrainBench2L "500|150 " chunkLength xs 500 150 0.02
    ]


{-# LANGUAGE BangPatterns #-}
module Game.Physics.Integration ( integrate
                                ) where

import Prelewd

import Control.Monad.ST
import Data.Packed.Development ( at' )
import Data.Packed.ST
import Data.Tuple
import Foreign.Storable ( Storable )
import qualified Data.Vector as D -- oh god this'll get confusing.
import Game.Physics.Types
import Numeric.GSL
import Numeric.LinearAlgebra hiding ( i )

-- | Converts a single physics object into its corresponding matrix
--   representation. This works for physics objects in any dimension.
--
--   > PObject (s_x s_y s_z)
--   >         (v_x v_y v_z)
--   >
--   > [ s_x v_x ]T
--   > | s_y v_y |
--   > [ s_z v_z ]
asMatrix :: PhysicsObject -> Matrix Double
asMatrix p = fromRows [ s, v ]
    where
        s = pos p
        v = vel p

-- | Converts a list of physics objects into their corresponding combined
--   matrix by stacking their 'asMatrix' representations vertically.
bigMat :: [PhysicsObject] -> Matrix Double
bigMat = fromColumns . concatMap (toColumns . asMatrix)

-- | This function might be in Haskell, but it's just as dangerous as C.
memcpy :: Storable t
       => STVector s t -- ^ dest
       -> Int -- ^ destination offset
       -> Vector t -- ^ src
       -> Int -- ^ source offset
       -> Int -- ^ number of elements to copy
       -> ST s ()
memcpy dst dstoff src srcoff n = for 0 (<n) (+1) $
                                    \idx -> unsafeWriteVector dst (dstoff+idx) (src `at'` (srcoff+idx))
{-# INLINE memcpy #-}

-- | An emulation of a C for loop.
--
--   for (int i = first clause; second clause; third clause)
--       { loop body }
for :: Monad m
    => a -- First clause.
    -> (a -> Bool) -- Second clause.
    -> (a -> a) -- Third clause.
    -> (a -> m b) -- Loop body. Return value is ignored for convenience.
    -> m ()
for i0 condition inc body = go i0
    where
        go !i | condition i = body i >> go (inc i)
              | otherwise   = return ()
        {-# INLINE go #-}
{-# INLINE for #-}

integrate :: D.Vector PhysicsObject -> Double -> D.Vector PhysicsObject
integrate ps _ | D.length ps == 0 = ps
integrate ps t = D.map (\((PObject _ _ im f), [s', v']) -> PObject s' v' im f)
               . D.zip ps
               . D.map toRows
               . D.fromList -- vector fusion can't handle concatenation.
               . concat
               . toBlocksEvery 2 dims
               . reshape width
               . (! 1) . toRows  -- second row.
               . odeSolveV RKf45  -- Is there a better choice of integrator?
                           t      -- Not really sure about this. odeSolve uses t/100.
                           1.0e-5 -- arbitrary. Lower it if simulation blows
                                  -- up due to inaccuracy. Highten it for
                                  -- inceased performance.
                           1.0e-5 -- same as above.
                           integrator
                           Nothing -- TODO: Can we supply a jacobian one day?
                                   --       Is it recommended, or totally unnecessary?
                                   --       I have no idea. What the fuck is a jacobian?
                           initialConditions
                           $ fromList [0.0, t]
    where
        initialConditions :: Vector Double
        initialConditions = flatten . bigMat $ D.toList ps

        -- integrates by shifting all values over one to the left, and querying
        -- all the force functions for new accelerations. This is written
        -- imperatively since it rapes the GC if done functionally, and takes
        -- 2x as long as a result.
        --
        -- msrc matrix layout (after a (reshape width))
        --
        -- [ s_x_0 s_y_0 s_z_0   s_x_1 s_y_1 s_z_1  ... ]
        -- [ v_x_0 v_y_0 v_z_0   v_x_1 v_y_1 v_z_1  ... ]
        --
        integrator :: Double -> Vector Double -> Vector Double
        integrator dt vsrc = runSTVector $ do
                vdst <- newUndefinedVector (dim vsrc)
                memcpy vdst width vsrc 0 width

                let xs = (subVector 0     width vsrc)
                    vs = (subVector width width vsrc)
                    objectAt i = let j = i*dims
                                     p = ps D.! i
                                  in PObject (subVector j dims xs)
                                             (subVector j dims vs)
                                             (invMass p)
                                             (netForce p)

                for (0, 0) ((< len) . fst) ((+1) *** (+dims)) $ \(i, j) ->
                    let p  = ps D.! i
                        im = invMass p
                        f  = netForce p objectAt dt
                     in for 0 (<dims) (+1) $ \k ->
                            unsafeWriteVector vdst (j+k) (im * (f `at'` k))

                return vdst

        -- dimensions of each point/vector.
        dims = dim . pos $ ps D.! 0

        -- the number of particles we're integrating.
        len = D.length ps

        -- width of every row in the matrix.
        width = dims*len


{-# LANGUAGE PackageImports,BangPatterns,StandaloneDeriving,TypeFamilies,FlexibleContexts,FlexibleInstances,TypeOperators,ExistentialQuantification,NoMonomorphismRestriction,DeriveTraversable,DeriveFunctor,DeriveFoldable #-}
module Kalman where
import Space.Class hiding(dim)
import Solver.RungeKutta
import qualified Space.Class as M
import Data.Typeable
import Data.Distributive
import Data.Functor.Compose
import Linear.Matrix hiding(trace)
import qualified Linear.Metric as Metric
import Vectorization
import Control.Applicative
import Data.Maybe
import Data.Foldable (Foldable,foldl')
import qualified  Data.Foldable as F
import Data.Traversable (Traversable,fmapDefault,mapAccumL)
import qualified Data.Traversable as T
import Linear.V0
import Linear.V1
import Foreign.Storable
import Foreign.FStorable
import Foreign.Cholesky
import qualified Numeric.LinearAlgebra as M
import qualified Data.Packed.Matrix as M
import Local
import Linear.Vector
import CartesianProduct
import Control.Monad.Trans.State

-- Covariance Version

fullM
  :: (RealFloat a, FStorable (Local b),
      Vectorize (Local b), Vectorize (Local b1), Space b,
      Space b1, Cholesky a, M.Field a) =>
     (b a -> b a)
     -> Local b (Local b a)
     -> (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> b1 a
     -> (b a, Local b (Local b a))
     -> (b a, Local b (Local b a))
fullM f r h q  m s= execState  ( do
      modify (prediction f r)
      modify (measure h q m)
      --modify (smoothing f r s)
      --modify (prediction f r)
      )  s

full
  :: (RealFloat a, FStorable (Local b2), FStorable (Local b),
      Vectorize (Local b), Vectorize (Local b1), Space b2, Space b,
      Space b1, Cholesky a, M.Field a) =>
     (b2 a -> b a)
     -> Local b (Local b a)
     -> (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> b1 a
     -> (b2 a, Local b2 (Local b2 a))
     -> (b a, Local b (Local b a))
full f r h q  m s= sm
  where sp = prediction f r s
        sm = measure h q m sp
        ss = smoothing f r s sm
        ssp = prediction f r ss
        ssm = measure h q m ssp

predictionM f = modify . prediction f
measureM h q  =  modify . measure  h q
smoothingM f r  = modify . smoothing f r

prediction
  :: (M.Field a , Vectorize (Local b) ,RealFloat a,
      FStorable (Local b1), Space b1, Space b, Cholesky a) =>
     (b1 a -> b a)
     -> Local b (Local b a)
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a))
prediction f r (x ,p)= (mean ,pxx )
    where
        (mean,pxx,_,_) = sigmaSamplingCovGain f r (x,p)

measure
  :: (RealFloat a, FStorable (Local b),
      Vectorize (Local b1), Space b, Space b1, Cholesky a, M.Field a) =>
     (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> b1 a
     -> (b a, Local b (Local b a))
     -> (b a, Local b (Local b a))
measure h q m (x ,p) = ( xb  , pb )
    where
        (mmean,pxx,pxy,k) =sigmaSamplingCovGain h q (x,p)
        xb = x |+| ( k !*  ( m |-| mmean ) )
        pb = p !-! ( k !*! pxx !*! distribute k )

smoothing
  :: (RealFloat a, FStorable (Local b), Vectorize (Local b1),
      Space b, Space b1, Cholesky a,
      M.Field a) =>
     (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> (b a, Local b (Local b a))
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a))
smoothing f r (x,p) (x1,p1) = (xb ,pb)
    where
        (xmean ,pxx,pxy,k) = sigmaSamplingCovGain f r (x,p)
        xb = x |+| (k !* ( x1 |-| xmean ) )
        pb = p !+! (k !*! ( p1 !-! pxx ) !*! distribute k )

sigmaSamplingCovGain
  :: (RealFloat a,M.Field a ,Vectorize (Local b),
      FStorable (Local b1), Space b1, Space b, Cholesky a) =>
     (b1 a -> b a)
     ->  (Local b (Local b a))
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a), Local b (Local b1 a),Local b1 (Local b a ))
sigmaSamplingCovGain f r (x,p) = (mean , pxx,pxy,k)
    where
        spoints = (sigmaPoints1 x p)
        points = fmap f spoints
        mean = sigmaMean1 points
        pxx = sigmaCov mean points !+! r
        pxy = sigmaCovXY x mean spoints  points
        k =  distribute pxy  !*! inverseM pxx

sigmaSamplingCov
  :: (RealFloat a, Show (Local b1 (Local b1 a)),
      FStorable (Local b1), Space b1, Space b, Cholesky a) =>
     (b1 a -> b a)
     ->  (Local b (Local b a))
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a), Local b (Local b1 a))

sigmaSamplingCov f r (x,p) = (mean , pxx,pxy)
    where
        spoints = (sigmaPoints1 x p)
        points = fmap f spoints
        mean = sigmaMean1 points
        pxx = sigmaCov mean points !+! r
        pxy = sigmaCovXY x mean spoints  points

inverseM = fromLists .  M.toLists . M.inv  . M.fromLists . toLists

{-
measureDelayed h m (x,p) (xlag,plag) =  (xnew,pxk)
    where
        -- Xlag
        sxlag = sigmaPoints1 xlag plag
        -- X
        sx = sigmaPoints1 x p
        -- XXlag
        pxxlag = sigmaCovXY x xlag sx sxlag
        -- Xk
        xk = x :|: xlag
        pxk = prodSym p plag pxxlag
        sxk = sigmaPoints1 xk pxk
        -- Y
        (sy,ym,py) = sigmaTransform h xlag plag
        pxlagy = sigmaCovXY xk ym sxk sy
        k = (distribute pxlagy) !*! (inverseM py)
        xnew = xk |+| (k !* (m |-| ym) )
-}

prodSym x y xy = prod x y xy (distribute xy)

prod x y xy yx = (liftA2 (:|:) x xy ) :|: (liftA2 (:|:) yx  y)

-- Helpers

square :: (Num a, R f) => f a -> f (f a)
square v = mult v v

squareT v = (distribute v) !*! v

mult x m = fmap (x ^*) m

sigmaMean
  :: (RealFloat a, Traversable t, FStorable t, Space b) =>
       t (b a) -> b a -> b a
sigmaMean p p0= case sigmaMean' p p0 of
                  Just x -> x
                  Nothing -> error $ "sigmaMean nonConvergence"

sigmaMean1 (V1 p0 :|: p )= sigmaMean p p0

sigmaMean'
  :: (RealFloat a, Traversable t, FStorable t, Space b) =>
       t (b a) -> b a -> Maybe (b a)
sigmaMean' p p0 = go 1000 p0
    where
        go ! i ! mik
            | i == 0 = Nothing
            | Metric.quadrance m > 1.0e-14 = go (i-1) mik1
            | otherwise = Just mik1
            where
                mik1 = mik |+| m
                s = F.foldr1 (^+^) $  fmap (|-|mik) p
                m =  s ^/ n
                n = 2 * (fromIntegral $ fsize p)

sigmaCovXY xmean ymean samplex sampley  =  covxy
    where
        xstd = fmap (|-| xmean) samplex
        ystd = fmap (|-| ymean) sampley
        covxy =  fmap (^/ 2) $ F.foldr1 (!+!) (liftA2 mult xstd ystd)

sigmaCov mean sample = cov
    where
        std  = fmap (|-| mean) sample
        cov  =  fmap (^/2) $ F.foldr1 (!+!) (fmap square std)

sigmaPoints1 x p = V1 x :|: points
  where
    m = case potrf p of
            Right x -> x
            Left i -> error $ "cholesky factorization error: sigmapoints1 : "
    points = fmap (x|+|) $  m  :|: fmap ((-1) *^) m

sqrtPrediction
  :: (Applicative g1, Applicative g, Traversable g, Traversable g1,
      FStorable g, FStorable g1, FStorable (Local b1), Space b,
      Space b1) =>
     (b Double -> b1 Double)
     -> g1 (Local b1 Double)
     -> (b Double, g (Local b Double))
     -> (b1 Double, Local b1 (Local b1 Double))
sqrtPrediction f r (x,s)= (xmean,sk'')
    where
        (y,w0m,w0c,w1) = weigthsGamma (fromIntegral $ fsize s)
        sp = fmap (x|+|) $ fmap (fmap (*y)) s :|: fmap (fmap (*(-y))) s
        sp' = fmap f sp
        x'  = f x
        xmean = sigmaMean1 $ V1 x' :|: sp'
        std  =fmap (|-| xmean) sp'
        sk' = dgeqrf $ distribute $ (fmap (fmap (*sqrt w1)) std):|: r
        sk'' = distribute $  ch1upJ sk' $ fmap (*sqrt w0c) $ x' |-| xmean

weigthsGamma l =  weigths 1 2 0 l

weigths a b k l= (y,w0m,w0c,w1) where
    lam =  a^2 * (l + k ) - l
    y =  sqrt ( l + lam)
    w0m = lam /(l + lam)
    w0c = w0m  +  (1 - a^2 + b)
    w1 = 1 / (2*(l+ lam))

ch1dnR x = unRight . ch1dn  x
sqrtMeasure
  :: (Show (Local b Double),Show (b Double),Applicative g, Traversable g, FStorable g, FStorable (Local b),
      FStorable (Local b1), Space b1, Space b) =>
     (b1 Double -> b Double)
     -> g (Local b Double)
     -> b Double
     -> (b1 Double, Local b1 (Local b1 Double))
     -> (b1 Double, Local b1 (Local b1 Double))
sqrtMeasure h r m (x,s) =  (x'',s''')
    where
{-
        spx = V1 x :|: (fmap (x|+|) $ fmap (fmap (*y)) s :|: fmap (fmap (*(-y))) s )
        wc = V1 w0c :|: pure w1
        --wm = V1 w0m :|: pure w1
        V1 y' :|: _ =  spy
        xmean = sigmaMean1  spx
        stdx  =  liftA2 (*^) wc $ fmap (|-| xmean ) spx
        spy = fmap h spx
        ymean = sigmaMean1  spy
        stdy  =  fmap (|-| ymean) spy
        pxy =   distribute stdx !*!  stdy
        sk' =  dgeqrf $ distribute $ fmap (fmap (*sqrt w1))  stdy :|: r
        sk''= distribute .fromJust  $   ch1up sk' (fmap (*w0c) $ ymean |-| y')
-}
        (ymean,pxx,pxy) = sigmaSampling h r (x,s)
        k  =  pseudoInverse pxy pxx
        -- x' = xmean |+| (k !* (m |-| y')  )
        x' = x |+| (k !* (m |-| ymean)  )
        u  = pxx !*! distribute k
        s' = distribute $ F.foldl ch1dnR  (distribute s)  u
        (x'',s''',_) = sigmaSampling id V0 (x',s')
{-      spx' = fmap (x' |+| ) $ fmap (fmap (*y))  s' :|: fmap (fmap (*(-y))) s'
        x'' =  sigmaMean1 $ V1 x' :|: spx'
        stdx' =  fmap (|-| x'') spx'
        s'' = dgeqrf $ distribute $ fmap (fmap (*sqrt w1))  stdx'
        s''' = distribute $ fromJust $ ch1up s'' $ fmap (*sqrt w0c) $ x'' |-| x'
-}

pseudoInverse
  :: (Applicative nhrs, Applicative m, Applicative m1,
      Traversable nhrs, Traversable m, Traversable m1, Distributive m1,
      FStorable nhrs, FStorable m, FStorable m1) =>
     nhrs (m1 Double) -> m (m1 Double) -> nhrs (m1 Double)
pseudoInverse y x  = dgelssR (distribute x )  (dgelssR x  y )

sqrtBackward f p (x,s) (x1,s1) = (xb ,sb)
    where
        (xmean ,pxx,pxy) = sigmaSampling f p (x,s)
        k =  pseudoInverse pxy pxx
        xb = x |+| (k !* (x1 |-| xmean ))
        s1pxx = distribute $ F.foldl (\y-> ch1dnR y)  (distribute s1 )  (distribute pxx )
        u  = s1pxx !*! distribute k
        sb = distribute $ F.foldl  ch1upJ  (distribute s)  u

sigmaSampling h r (x,s) = (ymean,sxx,sxy)
    where
        (y,w0m,w0c,w1) = weigthsGamma (fromIntegral $ fsize s)
        spx = V1 x :|: (fmap (x|+|) $ fmap (fmap (*y)) s :|: fmap (fmap (*(-y))) s )
        wc = V1 w0c :|: pure w1
        --wm = V1 w0m :|: pure w1
        V1 y' :|: _ =  spy
        xmean = sigmaMean1  spx
        stdx  =  liftA2 (*^) wc $ fmap (|-| xmean ) spx
        spy = fmap h spx
        ymean = sigmaMean1  spy
        stdy  =  fmap (|-| ymean) spy
        sxy =   distribute stdx !*!  stdy
        sk' =  dgeqrf $ distribute $ fmap (fmap (*sqrt w1))  stdy :|: r
        sxx = distribute $ ch1upJ sk' (fmap (*w0c) $ ymean |-| y')

ch1upJ x = fromJust . ch1up x
dgelssR x = unRight . dgelss x

unRight (Right x) = x
unRight (Left x) = error $ "no right: "


module Utils where

import           Control.Monad         (ap, liftM, liftM2)
import           Data.Char             (isNumber)
import           Data.List             (genericLength)
import qualified Data.Vector           as V
import qualified Data.Vector.Unboxed   as U
import           Numeric.LinearAlgebra (Matrix, cmap, cols, dropColumns,
                                        flatten, fromRows, loadMatrix,
                                        takeColumns, toList, toRows)
import qualified Numeric.LinearAlgebra as L
import           System.FilePath       (splitFileName)
import           System.Posix          (fileExist)

usage :: String
usage = "Usage: ./init_present <file> <Graph nbd size> <PCA nbd size>"

errorMsg1 :: String
errorMsg1 = "You did not enter two positive integers.\n"

errorMsg2 :: String
errorMsg2 = "You did not enter two positive integer arguments.\n"

graphMsg1 :: String
graphMsg1 = "The graph is not fully connected. Please enter larger integers!"

fileMsg :: String
fileMsg = "You have not provided an existent file. Please try again."

getFile :: FilePath -> Bool -> FilePath
getFile file pCond = if pCond
                       then file
                       else error fileMsg

fileArgs :: [String] -> IO Bool
fileArgs = (=<<) fileExist . return . head

intCheck :: Foldable t => t Char -> Bool
intCheck x = (not . null $ x) && foldr ((&&) . isNumber) True x

gtZero :: [String] -> Bool
gtZero = all ((> 0) . (\x -> read x :: Int))

procMNIST :: [Double] -> Matrix Double -> Matrix Double
procMNIST selection mat = images
  where
    images = fromRows [x | (x, c) <- zip xs cs
                         , or $ sequenceA eqfns c]
    xs = toRows . cmap (\x -> 255 - x :: Double) . mSel takeColumns $ mat
    cs = toList $ flatten . mSel dropColumns $ mat
    eqfns = [(== x) | x <- selection]
    mSel f x = f ((+ (-1)) $ cols x) x

checkArgs :: [String] -> (Int,Int)
checkArgs args = if ((== (3 :: Integer)) . genericLength) args
                   then let val = tail args
                        in if all intCheck val && gtZero val
                             then head . ap zip tail . Prelude.map read $ val -- yeah...
                             else error (errorMsg1 ++ usage)
                   else error (errorMsg2 ++ usage)

checkFile :: [FilePath] -> IO FilePath
checkFile args = if not . null $ args
                   then liftM2 getFile <$> return . head <*> fileArgs $ args
                   else error usage

getReqMatrix :: FilePath -> IO (Matrix Double)
getReqMatrix file = if (snd . splitFileName $ file) == "mnist.txt"
                      then liftM (procMNIST [4]) $ loadMatrix file
                      else loadMatrix file

convertDataset :: Matrix Double -> V.Vector (U.Vector Double)
convertDataset = V.map (U.fromList . L.toList) . V.fromList . L.toRows

rowSize :: Matrix Double -> Int
rowSize = L.rows


import Prelude hiding ((<>))
import Numeric.LinearAlgebra hiding (size, step)
import qualified System.Clock as Clock
import Data.Csv (FromNamedRecord, ToNamedRecord, DefaultOrdered)
import Data.Csv.Incremental
import GHC.Generics (Generic)
import qualified Data.ByteString.Lazy as LBS
import Numeric.Sundials.CVode.ODE
import Control.Exception
import System.IO.Unsafe

type Time = Double

getTime :: IO Time
getTime = do
  t <- Clock.getTime Clock.Monotonic
  let ns = realToFrac $ Clock.toNanoSecs t
  return $ ns / 10 ^ (9 :: Int)

timed :: IO a -> IO Time
timed t = do
  start <- getTime
  !_    <- t
  end   <- getTime
  return (end-start)

data Measurement t = Measurement
  { nts :: !Int -- the number of time steps in the output
  , step :: !Double -- the size of a single time step
  , time :: !t -- the total solving time (if known)
  }
  deriving Generic

instance FromNamedRecord (Measurement Time)
instance ToNamedRecord (Measurement Time)
instance DefaultOrdered (Measurement Time)

measure :: Measurement () -> IO (Measurement Time)
measure m@Measurement{..} = do
  let
    rhs _t y =
      let
        mu = 1000
        y1 = y ! 0
        y2 = y ! 1
      in
        fromList [y2, mu * (1 - y1*y1) * y2 - y1]
    y0 = fromList [2,0]
    !times = fromList . take (nts+1) $ iterate (+ step) 0
  time <- (timed . evaluate)
    (odeSolveV BDF Nothing 1e-2 1e-2 rhs y0 times :: Matrix Double)
  return m{time}

main :: IO ()
main = do
  measurements <- sequence $ do
    nts <- [2]
    step <- [1e5, 2e5 .. 2e9]
    let m0 = Measurement { time = (), .. }
    return . unsafeInterleaveIO $ measure m0
  LBS.writeFile "timings.csv" . encodeDefaultOrderedByName . mconcat $ encodeNamedRecord <$> measurements


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric  #-}
module FokkerPlanck.Histogram where

import           Control.DeepSeq
import           Data.Array.Repa     as R
import           Data.Binary
import           Data.Complex
import           Data.List           as L
import           Data.Vector.Unboxed as VU
-- import           GHC.Generics        (Generic)
import           Text.Printf
import           Types

data Histogram a = Histogram
  { histogramSize :: ![Int]
  , histogramNum  :: !Int
  , histogramVec  :: !(Vector a)
  } -- deriving (Generic)

instance (NFData a) => NFData (Histogram a) where
  rnf (Histogram !size !n !vec) = ()

instance (Binary a, Unbox a) => Binary (Histogram a) where
  put (Histogram size n vec) = do
    put size
    put n
    put . VU.toList $ vec
  get = do
    size <- get
    n <- get
    xs <- get
    return $! Histogram size n . VU.fromList $ xs

instance Eq (Histogram a) where
  (==) (Histogram size1 _ _) (Histogram size2 _ _) = size1 == size2

instance (Show a, Unbox a ) =>  Show (Histogram a) where
  show hist@(Histogram size n vec) =
    printf
      "Histogram size: %s, normalization term = %d\n%s\n"
      (showRepaShape size)
      n
      (show $ VU.take 10 vec)

{-# INLINE showRepaShape #-}
showRepaShape :: [Int] -> String
showRepaShape xs = "(Z" L.++ go (L.reverse xs)
  where
    go [] = ")"
    go (x:xs) = printf " :. %d%s" x (go xs)

{-# INLINE addHistogram #-}
addHistogram :: (Num a, Unbox a) => Histogram a -> Histogram a -> Histogram a
addHistogram (Histogram size1 !n1 !vec1) (Histogram size2 !n2 !vec2) =
  if size1 == size2
    then Histogram size1 (n1 + n2) (VU.zipWith (+) vec1 vec2)
    else error $
         "addHistogram: sizes are not equal.\n" L.++ show size1 L.++ " /= " L.++
         show size2

{-# INLINE addHistogramUnsafe #-}
addHistogramUnsafe :: (Num a, Unbox a) => Histogram a -> Histogram a -> Histogram a
addHistogramUnsafe (Histogram size1 !n1 !vec1) (Histogram _ !n2 !vec2) =
  Histogram size1 (n1 + n2) (VU.zipWith (+) vec1 vec2)

{-# INLINE getHistogramVec #-}
getHistogramVec ::
     (Unbox a, Fractional a) => Histogram a -> VU.Vector a
getHistogramVec (Histogram size n vec)
  | n == 0 = error "getHistogramArr: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getHistogramArr: size mismatch\nSize = " L.++ show size L.++ " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise = vec

{-# INLINE getNormalizedHistogramVec #-}
getNormalizedHistogramVec :: (Unbox a, Fractional a) => Histogram a -> VU.Vector a
getNormalizedHistogramVec (Histogram size n vec)
  | n == 0 = error "getNormalizedHistogramVec: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getNormalizedHistogramVec: size mismatch\nSize = " L.++ show size L.++
    " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise = VU.map (/ fromIntegral n) vec

{-# INLINE getHistogramArr #-}
getHistogramArr ::
     (Unbox a, Fractional a, Shape d) => Histogram a -> R.Array U d a
getHistogramArr (Histogram size n vec)
  | n == 0 = error "getHistogramArr: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getHistogramArr: size mismatch\nSize = " L.++ show size L.++ " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise = fromUnboxed (shapeOfList size) vec

{-# INLINE getNormalizedHistogramArr #-}
getNormalizedHistogramArr ::
     (Unbox a, Fractional a, Shape d) => Histogram a -> R.Array U d a
getNormalizedHistogramArr (Histogram size n vec)
  | n == 0 = error "getNormalizedHistogramArr: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getNormalizedHistogramArr: size mismatch\nSize = " L.++ show size L.++
    " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise = fromUnboxed (shapeOfList size) . VU.map (/ fromIntegral n) $ vec

{-# INLINE getNormalizedHistogramArrSink #-}
getNormalizedHistogramArrSink ::
     (Shape d) => Histogram (Complex Double) -> R.Array U d (Complex Double)
getNormalizedHistogramArrSink (Histogram size n vec)
  | n == 0 = error "getNormalizedHistogramArr: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getNormalizedHistogramArr: size mismatch\nSize = " L.++ show size L.++
    " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise =
    let normalizedVec =
          VU.map (\x -> (1 - (magnitude x) / fromIntegral n) :+ 0) $ vec
     in fromUnboxed (shapeOfList size) . VU.map (/ (VU.sum normalizedVec)) $
        normalizedVec

{-# INLINE getNormalizedHistogramArrSink' #-}
getNormalizedHistogramArrSink' ::
     (Shape d) => Histogram (Complex Double) -> R.Array U d (Complex Double)
getNormalizedHistogramArrSink' (Histogram size n vec)
  | n == 0 = error "getNormalizedHistogramArr: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getNormalizedHistogramArr: size mismatch\nSize = " L.++ show size L.++
    " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise =
    let normalizedVec =
          VU.map
            (\x ->
               let m = magnitude x
                in if m == 0
                     then fromIntegral n :+ 0
                     else (fromIntegral n / m) :+ 0) $
          vec
     in fromUnboxed (shapeOfList size) -- . VU.map (/ (VU.sum normalizedVec)) $
        normalizedVec

{-# INLINE getNormalizedHistogramArrR2Z2T0S0 #-}
getNormalizedHistogramArrR2Z2T0S0 ::
     Histogram (Complex Double) -> R.Array U DIM5 (Complex Double)
getNormalizedHistogramArrR2Z2T0S0 (Histogram size n vec)
  | n == 0 = error "getNormalizedHistogramArr: n == 0."
  | (L.product size) /= (VU.length vec) =
    error $
    "getNormalizedHistogramArr: size mismatch\nSize = " L.++ show size L.++
    " = " L.++
    (show . L.product $ size) L.++
    "\nVector length = " L.++
    (show . VU.length $ vec)
  | otherwise =
    let arr = fromUnboxed (shapeOfList size) . VU.map (\x ->  magnitude x / fromIntegral n) $ vec
        maxArr = R.foldS max 0 arr
     in computeS . R.traverse2 arr maxArr const $ \f fMax idx@(Z :. a :. b :. c :. d :. r) ->
          (f (Z :. a :. b :. c :. d :. 1) - f idx) :+ 0

{-# INLINE emptyHistogram #-}
emptyHistogram :: (Unbox a) => [Int] -> a -> Histogram a
emptyHistogram size = Histogram size 0 . VU.replicate (L.product size)

{-# INLINE mapHistogram #-}
mapHistogram :: (Unbox a, Unbox b) => (a -> b) -> Histogram a -> Histogram b
mapHistogram f (Histogram size n vec) = Histogram size n (VU.map f vec)


module Data.SQP
       ( Problem (..)
       , optimize
       ) where

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

import Debug.Trace

data Problem = Problem
               { _cost :: Vector Double -> Double
               , _approxCost :: Vector Double
                             -> (Matrix Double, Vector Double, Double)
               , _trueIneqs :: Vector Double -> Vector Double
               , _approxAffineIneqs :: Vector Double
                                    -> (Matrix Double, Vector Double)
               , _eqs :: Vector Double -> Vector Double
               , _approxEqs :: Vector Double -> (Matrix Double, Vector Double)
               , _numVariables :: Int
               , _numIneqs :: Int
               , _numEqs :: Int
               }

stepAcceptanceThreshold :: Double
stepAcceptanceThreshold = 0.25

trustShrinkFactor :: Double
trustShrinkFactor = 0.1

trustExpandFactor :: Double
trustExpandFactor = 1.5

constraintPenaltyScalingFactor :: Double
constraintPenaltyScalingFactor = 10

minTrustSize :: Double
minTrustSize = 1e-4

minModelImprove :: Double
minModelImprove = 1e-4

minModelImproveRatio :: Double
minModelImproveRatio = negate $ read "Infinity"

constraintSatisfactionThreshold :: Double
constraintSatisfactionThreshold = 1e-4

initTrustSize :: Double
initTrustSize = 0.1

initPenalty :: Double
initPenalty = 10

evalMerit :: Problem -> Double -> Vector Double -> Double
evalMerit problem penaltyParam x =
  let cost = _cost problem x
      penalty =
        penaltyParam * sum (map (max 0.0) $ toList $ _trueIneqs problem x) +
        penaltyParam * sum (map abs $ toList $ _eqs problem x)
  in  cost + penalty

optimize :: Problem
         -> Vector Double -- xInitial
         -> (Vector Double, Double)
optimize problem xInit =
  findSuitableConstraintPenalty problem xInit initTrustSize initPenalty

findSuitableConstraintPenalty :: Problem
                              -> Vector Double -- x
                              -> Double -- trust region size
                              -> Double -- constraint penalty
                              -- xNew, trueMerit
                              -> (Vector Double, Double)
findSuitableConstraintPenalty problem x trustSize penaltyParam =
  let (xNew, trueMerit, newTrustSize) =
        reconvexify problem x trustSize penaltyParam
      ineqsSatisfied = all (<= constraintSatisfactionThreshold) $
                         toList $ _trueIneqs problem xNew
      eqsSatisfied = all ((<= constraintSatisfactionThreshold) . abs) $
                       toList $ _eqs problem xNew
      constraintsSatisfied = ineqsSatisfied && eqsSatisfied
  in  if constraintsSatisfied
      then (xNew, trueMerit)
      else let penalty' = penaltyParam * constraintPenaltyScalingFactor
           in  findSuitableConstraintPenalty
                 problem xNew newTrustSize penalty'

evalQuadratic :: (Matrix Double, Vector Double, Double)
              -> Vector Double
              -> Double
evalQuadratic (m, b, c) x =
  (x `dot` ((0.5 * (m #> x)) + b)) + c

evalLinear :: (Matrix Double, Vector Double)
           -> Vector Double
           -> Vector Double
evalLinear (m, b) x = (m #> x) + b

reconvexify :: Problem
            -> Vector Double -- x
            -> Double -- trust region size
            -> Double -- constraint penalty
            -- xNew, newTrueMerit, newTrustSize
            -> (Vector Double, Double, Double)
reconvexify problem x trustSize penaltyParam =
  let convexCost = _approxCost problem x
      convexIneqs = _approxAffineIneqs problem x
      convexEqs = _approxEqs problem x
      trueMerit = evalMerit problem penaltyParam x
      trustResult = findSuitableTrustStep problem x convexCost
                      convexIneqs convexEqs trueMerit trustSize
                      penaltyParam
  in  case trustResult of
        Reconvexify xNew newMerit newTrustSize ->
          reconvexify problem xNew newTrustSize penaltyParam
        Finished xNew newMerit newTrustSize ->
          (xNew, newMerit, newTrustSize)

data TrustStepResult = -- x merit trustSize
                       Reconvexify (Vector Double) Double Double
                     | Finished (Vector Double) Double Double
  deriving (Show)

findSuitableTrustStep :: Problem
                      -> Vector Double -- x
                      -- convexified cost
                      -> (Matrix Double, Vector Double, Double)
                      -- convexified inequality constraints
                      -> (Matrix Double, Vector Double)
                      -- convexified equality constraints
                      -> (Matrix Double, Vector Double)
                      -> Double -- old true merit
                      -> Double -- trust region size
                      -> Double -- constraint penalty parameter
                      -- xNew, newTrueMerit, newTrustSize
                      -> TrustStepResult
findSuitableTrustStep
  problem x convexCost convexIneq convexEq oldTrueMerit trustSize penaltyParam =
  let trustStep =
        trustRegionStep problem x convexCost convexIneq convexEq oldTrueMerit
          trustSize penaltyParam
  in  case trace (show trustStep) trustStep of
        Reject -> let newTrustSize = trustSize * trustShrinkFactor
                  in  findSuitableTrustStep problem x convexCost convexIneq
                        convexEq oldTrueMerit newTrustSize penaltyParam
        Accept xNew newTrueMerit ->
          let newTrustSize = trustSize * trustExpandFactor
          in  Reconvexify xNew newTrueMerit newTrustSize
        Converged xNew newTrueMerit ->
          let newTrustSize =
                max trustSize $ (minTrustSize / trustShrinkFactor) * 1.5
          in  Finished xNew newTrueMerit newTrustSize

data IterationResult = Reject
                     -- xNew newMerit
                     | Accept (Vector Double) Double
                     | Converged (Vector Double) Double
  deriving (Show)

trustRegionStep :: Problem
                -> Vector Double -- x
                -- convexified cost
                -> (Matrix Double, Vector Double, Double)
                -- convexified inequality constraints
                -> (Matrix Double, Vector Double)
                -- convexified equality constraints
                -> (Matrix Double, Vector Double)
                -> Double -- old true merit
                -> Double -- trust region size
                -> Double -- constraint penalty parameter
                -> IterationResult
trustRegionStep
  problem x convexCost convexIneq convexEq oldTrueMerit trustSize penaltyParam =
    let (xNew, modelMerit) = solveQuadraticSubproblem problem x convexCost
                               convexIneq convexEq trustSize penaltyParam
        trueMerit =  evalMerit problem penaltyParam xNew
        trueImprove = oldTrueMerit - trueMerit
        modelImprove = oldTrueMerit - modelMerit
    in  if modelImprove < -1e-2 -- why so high? our convexification?
        then error $ "Model improvement got worse: " ++
               show (modelImprove, modelMerit, oldTrueMerit, trueMerit)
        else
          if trustSize < minTrustSize ||
             modelImprove < minModelImprove ||
             modelImprove / oldTrueMerit < minModelImproveRatio
          then Converged x oldTrueMerit
          else
            if trueImprove > 0.0 &&
               trueImprove / modelImprove > stepAcceptanceThreshold
            then Accept xNew trueMerit
            else Reject

solveQuadraticSubproblem :: Problem
                         -> Vector Double
                         -- convex cost
                         -> (Matrix Double, Vector Double, Double)
                         -- convexified inequality constraints
                         -> (Matrix Double, Vector Double)
                         -- convexified equality constraints
                         -> (Matrix Double, Vector Double)
                         -> Double -- trust region size
                         -> Double -- constraint penalty parameter
                         -> (Vector Double, Double) -- new x, model merit
solveQuadraticSubproblem
  problem x (costMatrix, costVector, costConstant)
  (approxIneqMatrix, approxIneqVector) (approxEqMatrix, approxEqVector)
  trustSize penaltyParam =
  -- We approximate each nonlinear inequality as |ax + b|^+. For each
  -- of these, we introduce a new optimization variable t (i.e., a
  -- slack variable) that comes with two inequalities:
  --
  -- 0 <= t
  --
  -- ax + b <= t
  --
  -- Each nonlinear equality as |ax + b|. For each of these, two new
  -- optimization variables s and t with following properties:
  --
  -- -s <= 0
  -- -t <= 0
  -- s - t = ax + b
  let numIneqs = _numIneqs problem
      numEqs = _numEqs problem
      numVariables = _numVariables problem

      costMatrixWithSlacks =
        diagBlock [ costMatrix
                  , konst 0.0 (numIneqs, numIneqs)
                  , konst 0.0 (2 * numEqs, 2 * numEqs)]

      -- LET'S MAKE COST SPD
      (l, v) = eigSH costMatrixWithSlacks
      positiveEigVals = cmap (max 1e-12) l
      augmentedCostMatrix = v <> diag positiveEigVals <> tr v

      augmentedCostVector = vjoin [costVector
                                  , konst 1.0 numIneqs
                                  , konst 1.0 $ 2 * numEqs]

      augmentedIneqMatrix =
        ((scalar penaltyParam * approxIneqMatrix) |||
         negate (ident numIneqs) |||
         konst 0.0 (numIneqs, 2 * numEqs)) ===
         (konst 0.0 (numIneqs + 2 * numEqs, numVariables) |||
          negate (ident $ numIneqs + 2 * numEqs))
      augmentedIneqVector =
        vjoin [ scalar penaltyParam * approxIneqVector
              , konst 0.0 $ numIneqs + 2 * numEqs ]

      ineqMatrixWithTrustConstraints =
        augmentedIneqMatrix ===
        (ident numVariables ||| konst 0.0 (numVariables, numIneqs + 2 * numEqs)) ===
        (negate (ident numVariables) ||| konst 0.0 (numVariables, numIneqs + 2 * numEqs))
      ineqVectorWithTrustConstraints =
        vjoin [ augmentedIneqVector
              , negate (x + scalar trustSize)
              , x - scalar trustSize ]

      augmentedEqMatrix =
        scalar penaltyParam * approxEqMatrix |||
        konst 0.0 (numEqs, numIneqs) |||
        negate (ident numEqs) |||
        ident numEqs
      augmentedEqVector = scalar penaltyParam * approxEqVector

      -- inequalities are given as ax + b <= 0, but quadprog++ wants
      -- a'x + b' >= 0
      result = solveQuadProg
        (augmentedCostMatrix, augmentedCostVector)
        (Just (augmentedEqMatrix, augmentedEqVector)) $
        Just ((-ineqMatrixWithTrustConstraints),
              (-ineqVectorWithTrustConstraints))

  in  case result of
        Left e -> error $ show e
        Right (xNewAugmented, newMeritHomogenous) ->
          let xNew = subVector 0 numVariables xNewAugmented
              newMerit = newMeritHomogenous + costConstant

              -- Compute the merit of the approximation without regularization
              cost = evalQuadratic (costMatrix, costVector, costConstant) xNew
              ineqViolations =
                evalLinear (approxIneqMatrix, approxIneqVector) xNew
              ineqPenalty =
                penaltyParam * (sum $ map (max 0.0) $ toList ineqViolations)
              eqViolations = evalLinear (approxEqMatrix, approxEqVector) xNew
              eqPenalty =
                penaltyParam * (sum $ map abs $ toList eqViolations)
              meritNoReg = cost + ineqPenalty + eqPenalty

              -- DEBUG
              meritError = abs $ newMerit - meritNoReg
          in  if meritError / (abs meritNoReg) > 0.5
              then error $ "merits don't match!" ++ show (meritNoReg, newMerit)
              else (xNew, meritNoReg)


{-# LANGUAGE TemplateHaskell #-}

module Scheme.Types where
import Data.DeriveTH
import Data.Maybe
import Data.IORef

import Control.Monad.Error

import qualified Data.Complex as C (Complex(..), realPart)
import Data.Ratio (Ratio, (%), numerator, denominator)
import Data.Vector as V (Vector, foldl)

import Text.Parsec.Error
import Text.Parsec.Pos

import System.IO
import DeriveGetMaybe

type ThrowsErrorIO = ErrorT LispError IO

data Env = Env
    { readOnly :: [(String, LispVal)]
    , definitions :: IORef [(String, IORef LispVal)] }
    deriving (Eq)

data LispVal
    = Atom String
    | Bool Bool
    | Num LispNum
    | Character Char
    | String String
    | List [LispVal]
    | DottedList [LispVal] LispVal
    | Vector (V.Vector LispVal)
    | Func [String] (Maybe String) [LispVal] Env
    | PrimitiveFunc String
    | Port Handle
    deriving (Eq)

data LispNum
    = Int Integer
    | Rational (Ratio Integer)
    | Double Double
    | Complex (C.Complex Double)
    deriving (Show, Eq)

data LispError
    = ParserError ParseError
    | NumArgs Int [LispVal]
    | BadExpr String LispVal
    | TypeMismatch String String
    | UnboundVar String String
    | Default String

instance Show LispVal where
    show (Atom s)            = s
    show (Bool s)            = if s then "#t" else "#f"
    show (Num n)             = show n
    show (Character c)       = "\\#" ++ [c]
    show (String s)          = "\"" ++ s ++ "\""
    show (PrimitiveFunc f)   = "<primitive : " ++ f ++ ">"
    show (List vs)           = "(" ++ unwords (map show vs) ++ ")"
    show (DottedList vs dot) = "(" ++ unwords (map show vs)
                                   ++ " . " ++ show dot ++ ")"
    show (Vector vs)         = "#(" ++ V.foldl (\acc e -> show e ++ " " ++ acc) "" vs ++ ")"
    show (Func params vararg body _) =
        "(lambda (" ++ unwords params
        ++ maybe "" (" . "++) vararg ++ ") " ++ unwords (map show body)++ ")"
    show (Port h)            = "<port: " ++ show h ++ ">"

$( derive makeIs ''LispVal )
$( derive makeIs ''LispNum )
$( derive makeGet ''LispVal )

getInt :: LispNum -> Maybe Integer
getInt (Int n) = Just n
getInt _ = Nothing

getExact :: LispNum -> Maybe (Ratio Integer)
getExact (Int n) = Just (n % 1)
getExact (Rational n) = Just n
getExact _ = Nothing

getAsDouble :: LispNum -> Double
getAsDouble (Int n) = fromInteger n
getAsDouble (Rational n) = fromInteger (numerator n) / fromInteger (denominator n)
getAsDouble (Complex n) = C.realPart n
getAsDouble (Double n) = n

getAsComplex :: LispNum -> C.Complex Double
getAsComplex (Complex n) = n
getAsComplex n = getAsDouble n C.:+  0

---------AHHHHHHHHH!!!!!!!!
--the horror

--engulfs

instance Fractional LispNum where
    a / b = if isComplex a && isComplex b
            then Complex $ getAsComplex a / getAsComplex b
            else Double $ getAsDouble a / getAsDouble b
    recip a = (Complex $ 1 C.:+ 0) / a
    fromRational = Rational

instance Num LispNum where
    (Int a) + (Int b) = Int $ a + b
    (Rational a) + (Rational b) = Rational $ a + b
    (Double a) + (Double b) = Double $ a + b
    (Complex a) + (Complex  b) = Complex $ a + b
    a + b = uncurry (+) $ coerceNums a b

    (Int a) * (Int b) = Int $ a * b
    (Rational a) * (Rational b) = Rational $ a * b
    (Double a) * (Double b) = Double $ a * b
    (Complex a) * (Complex  b) = Complex $ a * b
    a * b = uncurry (*) $ coerceNums a b

    a - b = (\(a', b') -> a' + negate b') (coerceNums a b)
    abs (Int n) = Int           $ abs n
    abs (Rational n) = Rational $ abs n
    abs (Double n) = Double     $ abs n
    abs (Complex n) = Complex   $ abs n
    negate (Int n) = Int           $ negate n
    negate (Rational n) = Rational $ negate n
    negate (Double n) = Double     $ negate n
    negate (Complex n) = Complex   $ negate n
    signum (Int n) = Int           $ signum n
    signum (Rational n) = Rational $ signum n
    signum (Double n) = Double     $ signum n
    signum (Complex n) = Complex   $ signum n
    fromInteger = Int

coerceNums :: LispNum -> LispNum -> (LispNum, LispNum)
coerceNums al bl = fromJust $ msum
    [ do
         a <- getInt al
         b <- getInt bl
         return (Int a, Int b)
    , do
         a <- getExact al
         b <- getExact bl
         return (Rational a, Rational b)
    , if isComplex al && isComplex bl
        then Just (Complex $ getAsComplex al, Complex $ getAsComplex bl)
        else Nothing
    , Just (Double $ getAsDouble al, Double $ getAsDouble bl) ]

instance Show LispError where
         show (ParserError e)       = show e
         show (NumArgs n vs)        = "expected " ++ show n ++ " args, found " ++ show (length vs) ++ "\n\t" ++ unwords (map show vs)
         show (TypeMismatch want v) = "type mismatch. expected " ++ want ++ ", found " ++ show v
         show (BadExpr msg expr)    = msg ++ " : " ++ show expr
         show (UnboundVar msg var)  = msg ++ " : " ++ var
         show (Default msg)         = "error : " ++ msg

instance Error LispError where
    noMsg = Default "mysterious failure experienced."
    strMsg = Default

nullEnv readOnlys = newIORef [] >>= \e -> return (Env readOnlys e)

setVar :: Env -> String -> LispVal -> ThrowsErrorIO ()
setVar e@(Env readOnlys env) var lisp = liftIO (readIORef env) >>= \envList ->
       maybe (defineVar e var lisp)
             (const $ throwError $ Default "setting an existing variable") $ lookup var envList

getVar (Env readOnlys env) var = maybe findInDefinitions return (lookup var readOnlys)
    where findInDefinitions = liftIO (readIORef env) >>= maybe
                    (throwError $ Default $ "undefined atom \"" ++ var ++ "\"" )
                    (liftIO . readIORef) . lookup var

defineVar (Env readOnlys env) var lisp =
    do envList <- liftIO $ readIORef env
       maybe (return ()) (throwError . Default . ("redefining read-only "++) . show) (lookup var readOnlys)
       maybe (liftIO $ newIORef lisp >>= \wrapped -> writeIORef env ((var, wrapped):envList))
             (\oldLisp -> liftIO $ writeIORef oldLisp lisp) $ lookup var envList

showIOThrows :: (MonadIO m, Show e) => ThrowsErrorIO e -> m String
showIOThrows = (liftIO . runErrorT) >=> return . either show show

-- introspection!
errToLisp :: LispError -> LispVal
errToLisp (ParserError p) =
    List [Atom "parse-error"
        , List [Atom "source-name", String . sourceName $ errorPos p]
        , List [Atom "source-line", Num . Int . fromIntegral . sourceLine $ errorPos p]
        , List [Atom "source-column", Num . Int . fromIntegral . sourceColumn $ errorPos p]
        , List [Atom "message", String $ showErrorMessages "or" "unknown parse error"
                                    "expecting" "unexpected" "end of input"
                                    (errorMessages p)]]
errToLisp err = String . show $ err


{-# language TypeOperators #-}
{-# language FlexibleInstances #-}
{-# language UndecidableInstances #-}
{-# language MonoLocalBinds #-}
module Data.Functor.Comm
  ( Comm
  , CommBind
  ) where

import Control.Applicative (ZipList, Const)
import Control.Applicative.Backwards
import Control.Commutative
import Control.Comonad
import Control.Comonad.Env
import Control.Comonad.Traced
import Control.Monad.Trans.Identity
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Reader
import Control.Monad.Trans.Writer.Lazy as Lazy
import Control.Monad.Trans.Writer.Strict as Strict
import Data.Complex
import Data.Functor.Bind
import Data.Functor.Compose
import Data.Functor.Identity
import Data.Functor.Product as Functor
import Data.Functor.Reverse
import Data.Map.Lazy
import Data.Monoid as Monoid
import Data.Proxy
import Data.Semigroup as Semigroup
import Data.Semigroup.Commutative
import Data.Tagged
import GHC.Generics as Generics
import Linear
import Linear.Plucker
import Linear.V

-- |
-- @
-- ('<.>') = 'flip' ('<..>')
-- @
class Apply f => Comm f

instance (Comm f, Comm g) => Comm (Compose f g)
instance (Comm m, CommutativeSemigroup w) => Comm (Strict.WriterT w m)
instance (Comm m, CommutativeSemigroup w) => Comm (Lazy.WriterT w m)
instance CommutativeSemigroup a => Comm ((,) a)

-- instance (CommutativeSemigroup a, CommutativeSemigroup b) => Comm ((,,) a b)
-- instance (CommutativeSemigroup a, CommutativeSemigroup b, CommutativeSemigroup c) => Comm ((,,,) a b c)
-- instance (CommutativeSemigroup a, CommutativeSemigroup b, CommutativeSemigroup c, CommutativeSemigroup d) => Comm ((,,,,) a b c d)

instance Comm ((->) e)
instance Comm m => Comm (ReaderT e m)
instance Comm V4
instance Comm V3
instance Comm V2
instance Comm Linear.V1
instance Comm V0
instance Comm Plucker
instance Dim n => Comm (V n)
instance Comm Complex
instance Comm Quaternion
instance Comm Identity
instance Comm m => Comm (IdentityT m)
instance (Comm m, Comm n) => Comm (Functor.Product m n)
instance CommutativeMonad m => Comm (MaybeT m)
instance Comm m => Comm (Reverse m)
instance Comm m => Comm (Backwards m)
instance CommutativeSemigroup a => Comm (Const a)
instance Comm (Tagged a)
instance Comm Proxy
-- instance Comm m => Comm (Kleisli m a)
instance Comm (Cokleisli w a)
instance Comm ZipList
instance Comm Option -- NB: Not CommutativeSemigroup!
-- semigroupoids missing instances:
instance Comm f => Comm (Alt f)
instance Comm Semigroup.Last -- NB: Not CommutativeSemigroup!
instance Comm Semigroup.First -- NB: Not CommutativeSemigroup!
instance Comm Monoid.Last -- NB: Not CommutativeSemigroup!
instance Comm Monoid.First -- NB: Not CommutativeSemigroup!
instance Comm Monoid.Product -- NB: Not CommutativeSemigroup!
instance Comm Monoid.Sum -- NB: Not CommutativeSemigroup!

instance Comm m => Comm (M1 i c m)
instance (Comm f, Comm g) => Comm (f :.: g)
instance (Comm f, Comm g) => Comm (f :*: g)
instance CommutativeSemigroup c => Comm (K1 i c) -- missing in base
instance Comm f => Comm (Rec1 f)
instance Comm Par1
instance Comm U1
instance (Comm w, CommutativeSemigroup e) => Comm (EnvT e w)
instance Comm w => Comm (TracedT e w)
-- instance Comm f => Comm (Ap f) -- ghc 8.6

class    (Comm m, Bind m) => CommBind m
instance (Comm m, Bind m) => CommBind m

-- instances that don't extend to Commutative
instance Ord k => Comm (Map k)
instance Comm Generics.V1


import Data.List (intercalate)
import Numeric.LinearAlgebra.Data ((|>))
import Data.Array (elems, bounds)
import Raytracer.Geometry (cube)
import Raytracer.Camera (Camera(Camera), Point(Point), calculate_rays, fire_rays)

-- This makes a crude ASCII image
make_shitty_image array = intercalate "\n" $ fmap (\list -> "." ++ list ++ ".") $ takes (x+1) $ fmap pixel $ elems array
	where
	(x, y) = let ((Point x0 y0), (Point x1 y1)) = bounds array; in (x1 - x0, y1 - y0)
	pixel (Just _) = '#'
	pixel Nothing = ' '
	takes n list = map fst $ scanl (\acc v -> splitAt v $ snd acc) (splitAt n list) $ replicate y n

main = putStrLn $ make_shitty_image $ fire_rays test_rays test_cube
	where
	wres = 160
	hres = 30
	test_camera = Camera 4 4 wres hres Nothing (3 |> [-3, -6, -3]) (3 |> [1,1,1])
	test_rays = calculate_rays test_camera
	test_cube = cube (3 |> [2, 0, 0]) (3 |> [0, 2, 0]) (3 |> [0, 0, 2]) (3 |> [0, 0, 0])