{-# LANGUAGE DataKinds, TypeFamilies #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
-- | Scalar-based implementation of fully connected neutral network
-- for classification of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistFcnnScalar where

import Prelude

import           Control.Exception (assert)
import           Data.Proxy (Proxy)
import qualified Data.Strict.Vector as Data.Vector
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, var0)
import HordeAd.Tool.MnistData

-- | Compute the output of a neuron, without applying activation function,
-- from trainable inputs in @xs@ and parameters (the bias and weights)
-- at @variables@ starting at @offset@. Useful for neurons in the middle
-- of the network, receiving inputs from other neurons.
sumTrainableInputs
  :: forall d r m. DualMonad d r m
  => Data.Vector.Vector (DualNumber d r) -> Int -> DualNumberVariables d r
  -> m (DualNumber d r)
sumTrainableInputs xs offset variables = do
  let bias = var0 variables offset
      f :: DualNumber d r -> Int -> DualNumber d r -> DualNumber d r
      f !acc i u =
        let v = var0 variables (offset + 1 + i)
        in acc + u * v
  returnLet $ V.ifoldl' f bias xs

-- | Compute the output of a neuron, without applying activation function,
-- from constant data in @xs@ and parameters (the bias and weights)
-- at @variables@ starting at @offset@. Useful for neurons at the bottom
-- of the network, tasked with ingesting the data.
sumConstantData
  :: forall d r m. DualMonad d r m
  => Vector r -> Int -> DualNumberVariables d r -> m (DualNumber d r)
sumConstantData xs offset variables = do
  let bias = var0 variables offset
      f :: DualNumber d r -> Int -> r -> DualNumber d r
      f !acc i r =
        let v = var0 variables (offset + 1 + i)
        in acc + scale r v
  returnLet $ V.ifoldl' f bias xs

hiddenLayerMnist
  :: forall d r m. DualMonad d r m
  => (DualNumber d r -> m (DualNumber d r)) -> Vector r
  -> DualNumberVariables d r -> Int
  -> m (Data.Vector.Vector (DualNumber d r))
hiddenLayerMnist factivation input variables width = do
  let nWeightsAndBias = V.length input + 1
      f :: Int -> m (DualNumber d r)
      f i = do
        outSum <- sumConstantData input (i * nWeightsAndBias) variables
        factivation outSum
  V.generateM width f

middleLayerMnist
  :: forall d r m. DualMonad d r m
  => (DualNumber d r -> m (DualNumber d r))
  -> Data.Vector.Vector (DualNumber d r)
  -> Int -> DualNumberVariables d r -> Int
  -> m (Data.Vector.Vector (DualNumber d r))
middleLayerMnist factivation hiddenVec offset variables width = do
  let nWeightsAndBias = V.length hiddenVec + 1
      f :: Int -> m (DualNumber d r)
      f i = do
        outSum <- sumTrainableInputs hiddenVec
                                     (offset + i * nWeightsAndBias)
                                     variables
        factivation outSum
  V.generateM width f

outputLayerMnist
  :: forall d r m. DualMonad d r m
  => (Data.Vector.Vector (DualNumber d r)
      -> m (Data.Vector.Vector (DualNumber d r)))
  -> Data.Vector.Vector (DualNumber d r) -> Int
  -> DualNumberVariables d r -> Int
  -> m (Data.Vector.Vector (DualNumber d r))
outputLayerMnist factivation hiddenVec offset variables width = do
  let nWeightsAndBias = V.length hiddenVec + 1
      f :: Int -> m (DualNumber d r)
      f i = sumTrainableInputs hiddenVec
                               (offset + i * nWeightsAndBias)
                               variables
  vOfSums <- V.generateM width f
  factivation vOfSums

fcnnMnistLen0 :: Int -> Int -> Int
fcnnMnistLen0 widthHidden widthHidden2 =
  widthHidden * (sizeMnistGlyph + 1)
  + widthHidden2 * (widthHidden + 1)
  + sizeMnistLabel * (widthHidden2 + 1)

-- | 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 widths of the hidden layers are @widthHidden@ and @widthHidden2@
-- and from these, the @fcnnMnistLen2@ function computes the number
-- of scalar dual number parameters (variables) to be given to the program.
fcnnMnist0 :: forall d r m. DualMonad d r m
           => (DualNumber d r -> m (DualNumber d r))
           -> (Data.Vector.Vector (DualNumber d r)
               -> m (Data.Vector.Vector (DualNumber d r)))
           -> Int
           -> Int
           -> Vector r
           -> DualNumberVariables d r
           -> m (Data.Vector.Vector (DualNumber d r))
fcnnMnist0 factivationHidden factivationOutput widthHidden widthHidden2
           input variables = do
  let !_A = assert (sizeMnistGlyph == V.length input) ()
  layer1 <- inline hiddenLayerMnist factivationHidden input
                                    variables widthHidden
  let offsetMiddle = widthHidden * (sizeMnistGlyph + 1)
  layer2 <- inline middleLayerMnist factivationHidden layer1
                                    offsetMiddle variables widthHidden2
  let offsetOutput = offsetMiddle + widthHidden2 * (widthHidden + 1)
  inline outputLayerMnist factivationOutput layer2
                          offsetOutput variables sizeMnistLabel

-- | The neural network applied to concrete activation functions
-- and composed with the appropriate loss function.
fcnnMnistLoss0
  :: DualMonad d r m
  => Int -> Int -> MnistData r -> DualNumberVariables d r
  -> m (DualNumber d r)
fcnnMnistLoss0 widthHidden widthHidden2 (input, target) variables = do
  result <- inline fcnnMnist0 logisticAct softMaxAct
                              widthHidden widthHidden2 input variables
  lossCrossEntropy target result

-- | A function testing the neural network given testing set of inputs
-- and the trained parameters.
--
-- The proxy argument is needed only for the (spurious) SPECIALIZE pragma,
-- becuase I can't write @SPECIALIZE fcnnMnistTest0 \@Double@.
fcnnMnistTest0 :: forall r. IsScalar 'DModeGradient r
           => Proxy r -> Int -> Int -> [MnistData r] -> Domain0 r
           -> r
fcnnMnistTest0 _ widthHidden widthHidden2 inputs params0 =
  let matchesLabels :: MnistData r -> Bool
      matchesLabels (glyph, label) =
        let nn = inline (fcnnMnist0 @'DModeGradient) logisticAct softMaxAct
                                        widthHidden widthHidden2 glyph
            value = V.map (\(D r _) -> r)
                    $ primalValueGeneral nn (params0, V.empty, V.empty, V.empty)
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)
{-# SPECIALIZE fcnnMnistTest0 :: Proxy Double -> Int -> Int -> [MnistData Double] -> Domain0 Double -> Double #-}


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}

import           Codec.Picture.Bitmap
import           Codec.Picture.Types
import           Control.Monad.Random
import           Data.Serialize
import           Data.Word (Word8)
import           GHC.TypeLits
import           Grenade
import           Options.Applicative
import qualified Data.ByteString as B
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra.Static as SA

type ZDim      = 5
type TotalDim  = 2+5
type OutDim    = 3
type Hidden    = 100

-- type BatchSize = 500 * 500

type CppnNet = Network '[ FullyConnected TotalDim Hidden , Tanh
                        , FullyConnected Hidden   Hidden , Softmax
                        , FullyConnected Hidden   Hidden , Tanh
                        , FullyConnected Hidden   OutDim , Logit
                        ]
                       '[ 'D1 TotalDim
                        , 'D1 Hidden   , 'D1 Hidden
                        , 'D1 Hidden   , 'D1 Hidden
                        , 'D1 Hidden   , 'D1 Hidden
                        , 'D1 OutDim   , 'D1 OutDim
                        ]

randomNet :: MonadRandom m => m CppnNet
randomNet = randomNetwork

buildInputs :: Int
            -> Int
            -> S ('D1 ZDim)
            -> [ S ('D1 TotalDim) ]
buildInputs h w (S1D z') = vects
  where
    z      = V.toList (SA.unwrap z')
    start  = -1
    end    = 1
    range  = end - start
    -- TODO: Presently only works for a square.
    step   = range / (fromIntegral w - 1)
    pts    = [ start, start + step .. end ]
    points = [ [x, y] ++ z | x <- pts, y <- pts ]

    vects  = map (S1D . SA.vector) points

netForward :: CppnNet
           -> Int
           -> Int
           -> IO (CppnNet, Image PixelRGB8)
netForward net width height = do

  z <- do
    seed <- getRandom
    return $ S1D (SA.randomVector seed SA.Uniform :: SA.R ZDim)

  let inputs       = buildInputs width height z
      outputs      = map (runNet net) inputs
      extractPixel = toPixel . extract

      pixels :: V.Vector (PixelBaseComponent PixelRGB8)
      pixels = V.fromList (concatMap extractPixel outputs)

      img :: Image PixelRGB8
      img = Image width height pixels

  return $ (net, img)

toPixel :: (Double, Double, Double)
        -> [Word8]
toPixel (r,g,b) = [ fromIntegral $ round (r * 255)
                  , fromIntegral $ round (g * 255)
                  , fromIntegral $ round (b * 255) ]

extract :: S ('D1 OutDim) -> (Double, Double, Double)
extract (S1D a) = let [r,g,b] = V.toList (SA.unwrap a)
                    in (r,g,b)

netLoad :: FilePath -> IO CppnNet
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $ runGet (get :: Get CppnNet) modelData

data CppnOpts = CppnOpts Int Int (Maybe FilePath) (Maybe FilePath)

opts :: Parser CppnOpts
opts =
  CppnOpts <$> option auto (long "width"  <> short 'w' <> value 500)
           <*> option auto (long "height" <> short 'h' <> value 500)
           <*> optional (strOption (long "load"))
           <*> optional (strOption (long "save"))

main :: IO ()
main = do
  CppnOpts width height load save <- execParser (info (opts <**> helper) idm)

  net0 <- case load of
    Just loadFile -> netLoad loadFile
    Nothing -> randomNet

  (net, image) <- netForward net0 width height

  writeBitmap "test.png" image

  case save of
    Just saveFile -> B.writeFile saveFile $ runPut (put net)
    Nothing -> return ()

  putStrLn $ "Done!"


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Bitraversable
import           Data.Default
import           Data.Finite
import           Data.IDX
import           Data.List.Split
import           Data.Time.Clock
import           Data.Traversable
import           Data.Tuple
import           Data.Type.Product
import           Learn.Neural
import           Numeric.BLAS.HMatrix
import           Numeric.LinearAlgebra.Static
import           Text.Printf
import qualified Data.Vector                     as V
import qualified Data.Vector.Generic             as VG
import qualified Data.Vector.Unboxed             as VU
import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWC

loadMNIST
    :: FilePath
    -> FilePath
    -> IO (Maybe [(HM '[784], HM '[10])])
loadMNIST fpI fpL = runMaybeT $ do
    i <- MaybeT          $ decodeIDXFile       fpI
    l <- MaybeT          $ decodeIDXLabelsFile fpL
    d <- MaybeT . return $ labeledIntData l i
    r <- MaybeT . return $ for d (bitraverse mkImage mkLabel . swap)
    liftIO . evaluate $ force r
  where
    mkImage :: VU.Vector Int -> Maybe (HM '[784])
    mkImage = fmap HM . create . VG.convert . VG.map (\i -> fromIntegral i / 255)
    mkLabel :: Int -> Maybe (HM '[10])
    mkLabel = fmap (oneHot . only) . packFinite . fromIntegral

main :: IO ()
main = MWC.withSystemRandom $ \g -> do
    Just train <- loadMNIST "data/train-images-idx3-ubyte" "data/train-labels-idx1-ubyte"
    Just test  <- loadMNIST "data/t10k-images-idx3-ubyte"  "data/t10k-labels-idx1-ubyte"
    putStrLn "Loaded data."
    net0 :: Network 'FeedForward HM ( '[784] :~ FullyConnected )
                                   '[ '[300] :~ LogitMap
                                    , '[300] :~ FullyConnected
                                    , '[100] :~ LogitMap
                                    , '[100] :~ FullyConnected
                                    , '[10 ] :~ SoftMax '[10]
                                    ]
                                    '[10] <- initDefNet g
    let dout = alongNet net0 $ Nothing
                           :&% Just 0.2
                           :&% Nothing
                           :&% Just 0.2
                           :&% Nothing
                           :&% DOExt
    flip evalStateT net0 . forM_ [1..] $ \e -> do
      train' <- liftIO . fmap V.toList $ MWC.uniformShuffle (V.fromList train) g
      liftIO $ printf "[Epoch %d]\n" (e :: Int)

      forM_ ([1..] `zip` chunksOf batch train') $ \(b, chnk) -> StateT $ \n0 -> do
        printf "(Batch %d)\n" (b :: Int)

        t0 <- getCurrentTime
        -- n' <- evaluate $ optimizeList_ (bimap only_ only_ <$> chnk) n0
        --                                -- (sgdOptimizer rate netOpPure crossEntropy)
        --                                (adamOptimizer def netOpPure crossEntropy)
        n' <- optimizeListM_ (bimap only_ only_ <$> chnk) n0
                             (adamOptimizerM def (netOpDOPure dout g) crossEntropy)
        t1 <- getCurrentTime
        printf "Trained on %d points in %s.\n" batch (show (t1 `diffUTCTime` t0))

        let trainScore = testNetList maxTest (someNet n') chnk
            testScore  = testNetList maxTest (someNet n') test
        printf "Training error:   %.2f%%\n" ((1 - trainScore) * 100)
        printf "Validation error: %.2f%%\n" ((1 - testScore ) * 100)

        return ((), n')
  where
    -- rate :: Double
    -- rate  = 0.02
    batch :: Int
    batch = 2500


{-# LANGUAGE TemplateHaskell #-}

module Types where

import Lens.Micro.TH
import Lens.Micro
import qualified Statistics.Distribution.Normal as Stat
import qualified Statistics.Distribution as Stat

data Bandit d = Bandit d

instance (Stat.Mean d, Stat.Variance d) => Show (Bandit d) where
  show (Bandit d) = "Bandit {mean = " ++ show (Stat.mean d) ++ ", stdDev = " ++ show (Stat.stdDev d) ++ "}"

distLens :: Lens' (Bandit d) d
distLens = lens (\(Bandit d) -> d) (\_ d -> Bandit d)

meanLens :: Lens' Stat.NormalDistribution Double
meanLens = lens (\d -> Stat.mean d) (\d m -> Stat.normalDistr m (Stat.stdDev d))

-- Statistics about a given bandit.
data BanditStats = BanditStats { _totalReward :: Double
                               , _timesPulled :: Int
                               } deriving (Eq, Show)

makeLenses ''BanditStats

instance Monoid BanditStats where
  a `mappend` b = BanditStats (a ^. totalReward + b ^. totalReward) (a ^. timesPulled + b ^. timesPulled)
  mempty = BanditStats 0.0 0

type BanditDist = Stat.NormalDistribution

data SimState = SimState { _bandits :: [(BanditStats, Bandit BanditDist)] -- ^ The bandits and associated statistics
                         , _roundsCount :: Int -- ^ Number of rounds so far
                         , _regretHist :: [Double]
                         }

makeLenses ''SimState

instance Show SimState where
  show (SimState b count regretHist_) = "SimState { bandits = " ++ show b ++ ", n = " ++ show count ++ ", regret = "++ show regretHist_ ++"}"

type BanditState = (BanditStats, Bandit BanditDist)


-- |
-- Module      : Streamly.Statistics
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : Apache-2.0
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Statistical measures over a stream of data. All operations use numerically
-- stable floating point arithmetic.
--
-- Measurements can be performed over the entire input stream or on a sliding
-- window of fixed or variable size.  Where possible, measures are computed
-- online without buffering the input stream.
--
-- Currently there is no overflow detection.
--
-- References:
--
-- * https://en.wikipedia.org/wiki/Statistics
-- * https://mathworld.wolfram.com/topics/ProbabilityandStatistics.html

-- Resources:
--
-- This may be another useful resource for incremental (non-windowed)
-- computation:
--
-- https://www.researchgate.net/publication/287152055_Incremental_Statistical_Measures
--
-- Sample Statistics
--
-- Terms
--
-- Population: the complete data set from which statistical samples are taken.
--
-- Sample: a subset of the population.
--
-- https://en.wikipedia.org/wiki/Sample_(statistics)
--
-- Estimator:
--
-- Statistical measures can be computed either from the actual population
-- or from samples. Statistical measures computed from the samples provide an
-- estimate of the actual measures of the entire population. Measures computed
-- from samples may not truly reflect the actual measures and may have to be
-- adjusted for biases or errors.
--
-- An "estimator" is a method or function to compute a statistical measure from
-- sampled data. For example, the sample variance is an esitmator of the
-- population variance.
--
-- https://en.wikipedia.org/wiki/Estimator
--
-- Bias:
--
-- The result computed by an estimator may not be centered at the true value as
-- determined by computing the measure for the actual population. Such an
-- estimator is called a biased estimator.  For example, notice how
-- 'sampleVariance' is adjusted for bias.
--
-- https://en.wikipedia.org/wiki/Bias_of_an_estimator
--
-- Consistency:
--
-- https://en.wikipedia.org/wiki/Consistent_estimator

{-# LANGUAGE ScopedTypeVariables #-}
module Streamly.Statistics
    (
    -- * Incremental Folds
    -- | Folds of type @Fold m (a, Maybe a) b@ are incremental sliding window
    -- folds. An input of type @(a, Nothing)@ indicates that the input element
    -- @a@ is being inserted in the window without ejecting an old value
    -- increasing the window size by 1. An input of type @(a, Just a)@
    -- indicates that the first element is being inserted in the window and the
    -- second element is being removed from the window, the window size remains
    -- the same. The window size can only increase and never decrease.
    --
    -- You can compute the statistics over the entire stream using sliding
    -- window folds by keeping the second element of the input tuple as
    -- @Nothing@.
    --
      Window.lmap
    , Window.cumulative

    -- * Summary Statistics
    -- | See https://en.wikipedia.org/wiki/Summary_statistics .

    -- ** Sums
    , Window.length
    , Window.sum
    , Window.sumInt
    , Window.powerSum

    -- ** Location
    -- | See https://en.wikipedia.org/wiki/Location_parameter .
    --
    -- See https://en.wikipedia.org/wiki/Central_tendency .
    , Window.minimum
    , Window.maximum
    , rawMoment
    , rawMomentFrac

    -- Pythagorean means (https://en.wikipedia.org/wiki/Pythagorean_means)
    , mean
    , welfordMean
    , geometricMean
    , harmonicMean

    , quadraticMean

    -- Generalized mean
    , powerMean
    , powerMeanFrac

    -- ** Weighted Means
    -- | Exponential Smoothing.
    , ewma
    , ewmaAfterMean
    , ewmaRampUpSmoothing

    -- ** Spread
    -- | Second order central moment is a statistical measure of dispersion.
    -- The \(k\)th moment about the mean (or \(k\)th central moment) is defined
    -- as:
    --
    -- \(\mu_k = \frac{1}{n}\sum_{i=1}^n {(x_{i}-\mu)}^k\)
    --
    -- See https://mathworld.wolfram.com/CentralMoment.html .
    --
    -- See https://en.wikipedia.org/wiki/Statistical_dispersion .
    , Window.range
    , md
    , variance
    , stdDev

    -- ** Shape
    -- | Third and fourth order central moments are a measure of shape.
    --
    -- See https://en.wikipedia.org/wiki/Shape_parameter .
    --
    -- See https://en.wikipedia.org/wiki/Standardized_moment .
    , skewness
    , kurtosis

    -- XXX Move to Statistics.Sample or Statistics.Estimation module?
    -- ** Estimation
    , sampleVariance
    , sampleStdDev
    , stdErrMean

    -- ** Resampling
    , resample
    , foldResamples
    , jackKnifeMean
    , jackKnifeVariance
    , jackKnifeStdDev

    -- ** Probability Distribution
    , frequency
    , mode

    -- Histograms
    , HistBin (..)
    , binOffsetSize
    , binFromSizeN
    , binFromToN
    , binBoundaries
    , histogram

    -- * Transforms
    , fft
    )
where

import Control.Exception (assert)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Functor.Identity (runIdentity, Identity)
import Data.Bits (Bits(complement, shiftL, shiftR, (.&.), (.|.)))
import Data.Complex (Complex ((:+)))
import Data.Map.Strict (Map, foldrWithKey)
import Foreign.Storable (Storable)
import Streamly.Data.Fold.Tee(Tee(..), toFold)
import Streamly.Internal.Control.Concurrent (MonadAsync)
import Streamly.Internal.Data.Array.Foreign.Type
    (Array, length, toStream, unsafeIndexIO)
import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
import Streamly.Internal.Data.Stream.IsStream (SerialT)
import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import System.Random.MWC (createSystemRandom, uniformRM)

import qualified Streamly.Internal.Data.Array.Foreign as Array
import qualified Streamly.Internal.Data.Array.Foreign.Mut as MA
import qualified Streamly.Internal.Data.Fold as Fold
import qualified Streamly.Internal.Data.Fold.Window as Window
import qualified Streamly.Internal.Data.Stream.IsStream as Stream
import qualified Streamly.Internal.Data.Unfold as Unfold

import Prelude hiding (length, sum, minimum, maximum)

-- TODO: Overflow checks. Would be good if we can directly replace the
-- operations with overflow checked operations.
--
-- See https://hackage.haskell.org/package/safe-numeric
-- See https://hackage.haskell.org/package/safeint
--
-- TODO We have many of these functions in Streamly.Data.Fold as well. Need to
-- think about deduplication.

-------------------------------------------------------------------------------
-- Transforms
-------------------------------------------------------------------------------

-- XXX These utility functions can be moved to streamly-numeric

-- | Test if the given integer value is a power of 2.
{-# INLINE isPower2 #-}
isPower2 :: Int -> Bool
isPower2 n = n .&. (n - 1) == 0

-- | Create a power of 2
--
-- Argument must be less than 64 assuming 64-bit Int size.
--
{-# INLINE _power2 #-}
_power2 :: Int -> Int
_power2 n = shiftL 1 n

-- | Create a bit mask with lower n bits 0 and the rest as 1.
--
-- Argument must be less than 64 assuming 64-bit Int size.
--
{-# INLINE maskLowerN #-}
maskLowerN :: Int -> Int
maskLowerN n = complement (shiftL 1 n - 1)

-- | Compute the base 2 logarithm of the given value.
--
-- Assumes the Int size to be 64-bit.
--
{-# INLINE logBase2 #-}
logBase2 :: Int -> Int
logBase2 v0
    | v0 <= 0   = error $ "logBase2: input must be greater than 0 " ++ show v0
    | otherwise = go 32 0 v0

    where

    go !bits !result !v
        | bits == 0 = result
        | v .&. maskLowerN bits /= 0 =
             go (bits `shiftR` 1) (result .|. bits) (v `shiftR` bits)
        | otherwise = go (bits `shiftR` 1) result v

-- Algo translated from the statistics library.
--
-- XXX We can use a wrapper API that takes an array of Double input instead of
-- array of Complex.
--
-- | Compute fast fourier transform of an array of 'Complex' values.
--
-- Array length must be power of 2.
--
{-# INLINE fft #-}
fft :: MonadIO m => MA.Array (Complex Double) -> m ()
fft marr
    | isPower2 len = bitReverse 0 0
    | otherwise  = error "fft: Array length must be power of 2"

    where

    len = MA.length marr

    halve x = x `shiftR` 1

    twice x = x `shiftL` 1

    inner i j k
        | k <= j  = inner i (j - k) (halve k)
        | otherwise = bitReverse (i + 1) (j + k)

    bitReverse i j
        | i == len - 1 = stage 0 1
        | otherwise = do
            when (i < j) $ MA.unsafeSwapIndices i j marr
            inner i j (halve len)

    log2len = logBase2 len

    stage l !l1
        | l == log2len = return ()
        | otherwise = do
            let !l2 = twice l1
                !e  = -6.283185307179586/fromIntegral l2
                flight j !a | j == l1   = stage (l + 1) l2
                            | otherwise = do
                    let butterfly i | i >= len  = flight (j + 1) (a + e)
                                    | otherwise = do
                            let i1 = i + l1
                            xi1 :+ yi1 <- MA.getIndexUnsafe i1 marr
                            let !c = cos a
                                !s = sin a
                                d  = (c * xi1 - s * yi1) :+ (s * xi1 + c * yi1)
                            ci <- MA.getIndexUnsafe i marr
                            MA.putIndexUnsafe  i1 (ci - d) marr
                            MA.putIndexUnsafe  i (ci + d) marr
                            butterfly (i + l2)
                    butterfly j
            flight 0 0

-------------------------------------------------------------------------------
-- Mean
-------------------------------------------------------------------------------

-- | Arithmetic mean of elements in a sliding window:
--
-- \(\mu = \frac{\sum_{i=1}^n x_{i}}{n}\)
--
-- This is also known as the Simple Moving Average (SMA) when used in the
-- sliding window and Cumulative Moving Avergae (CMA) when used on the entire
-- stream.
--
-- Mean is the same as the first raw moment.
--
-- \(\mu = \mu'_1\)
--
-- >>> mean = rawMoment 1
-- >>> mean = powerMean 1
-- >>> mean = Fold.teeWith (/) sum length
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE mean #-}
mean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
mean = Window.mean

-- | Recompute mean from old mean when an item is removed from the sample.
{-# INLINE _meanSubtract #-}
_meanSubtract :: Fractional a => Int -> a -> a -> a
_meanSubtract n oldMean oldItem =
    let delta = (oldItem - oldMean) / fromIntegral (n - 1)
     in oldMean - delta

-- | Recompute mean from old mean when an item is added to the sample.
{-# INLINE meanAdd #-}
meanAdd :: Fractional a => Int -> a -> a -> a
meanAdd n oldMean newItem =
    let delta = (newItem - oldMean) / fromIntegral (n + 1)
     in oldMean + delta

-- We do not carry rounding errors, therefore, this would be less numerically
-- stable than the kbn mean.
--
-- | Recompute mean from old mean when an item in the sample is replaced.
{-# INLINE meanReplace #-}
meanReplace :: Fractional a => Int -> a -> a -> a -> a
meanReplace n oldMean oldItem newItem =
    let n1 = fromIntegral n
        -- Compute two deltas instead of a single (newItem - oldItem) because
        -- the latter would be too small causing rounding errors.
        delta1 = (newItem - oldMean) / n1
        delta2 = (oldItem - oldMean) / n1
     in (oldMean + delta1) - delta2

-- | Same as 'mean' but uses Welford's algorithm to compute the mean
-- incrementally.
--
-- It maintains a running mean instead of a running sum and adjusts the mean
-- based on a new value.  This is slower than 'mean' because of using the
-- division operation on each step and it is numerically unstable (as of now).
-- The advantage over 'mean' could be no overflow if the numbers are large,
-- because we do not maintain a sum, but that is a highly unlikely corner case.
--
-- /Internal/
{-# INLINE welfordMean #-}
welfordMean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
welfordMean = Fold step initial extract

    where

    initial =
        return
            $ Partial
            $ Tuple'
                (0 :: a)   -- running mean
                (0 :: Int) -- count of items in the window

    step (Tuple' oldMean w) (new, mOld) =
        return
            $ Partial
            $ case mOld of
                Nothing -> Tuple' (meanAdd w oldMean new) (w + 1)
                Just old -> Tuple' (meanReplace w oldMean old new) w

    extract (Tuple' x _) = return x

-------------------------------------------------------------------------------
-- Moments
-------------------------------------------------------------------------------

-- XXX We may have chances of overflow if the powers are high or the numbers
-- are large. A limited mitigation could be to use welford style avg
-- computation. Do we need an overflow detection?
--
-- | Raw moment is the moment about 0. The \(k\)th raw moment is defined as:
--
-- \(\mu'_k = \frac{\sum_{i=1}^n x_{i}^k}{n}\)
--
-- >>> rawMoment k = Fold.teeWith (/) (powerSum p) length
--
-- See https://en.wikipedia.org/wiki/Moment_(mathematics) .
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE rawMoment #-}
rawMoment :: (Monad m, Fractional a) => Int -> Fold m (a, Maybe a) a
rawMoment k = Fold.teeWith (/) (Window.powerSum k) Window.length

-- | Like 'rawMoment' but powers can be negative or fractional. This is
-- slower than 'rawMoment' for positive intergal powers.
--
-- >>> rawMomentFrac p = Fold.teeWith (/) (powerSumFrac p) length
--
{-# INLINE rawMomentFrac #-}
rawMomentFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
rawMomentFrac k = Fold.teeWith (/) (Window.powerSumFrac k) Window.length

-- XXX Overflow can happen when large powers or large numbers are used. We can
-- keep a running mean instead of running sum but that won't mitigate the
-- overflow possibility by much. The overflow can still happen when computing
-- the mean incrementally.

-- | The \(k\)th power mean of numbers \(x_1, x_2, \ldots, x_n\) is:
--
-- \(M_k = \left( \frac{1}{n} \sum_{i=1}^n x_i^k \right)^{\frac{1}{k}}\)
--
-- \(powerMean(k) = (rawMoment(k))^\frac{1}{k}\)
--
-- >>> powerMean k = (** (1 / fromIntegral k)) <$> rawMoment k
--
-- All other means can be expressed in terms of power mean. It is also known as
-- the generalized mean.
--
-- See https://en.wikipedia.org/wiki/Generalized_mean
--
{-# INLINE powerMean #-}
powerMean :: (Monad m, Floating a) => Int -> Fold m (a, Maybe a) a
powerMean k = (** (1 / fromIntegral k)) <$> rawMoment k

-- | Like 'powerMean' but powers can be negative or fractional. This is
-- slower than 'powerMean' for positive intergal powers.
--
-- >>> powerMeanFrac k = (** (1 / k)) <$> rawMomentFrac k
--
{-# INLINE powerMeanFrac #-}
powerMeanFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
powerMeanFrac k = (** (1 / k)) <$> rawMomentFrac k

-- | The harmonic mean of the positive numbers \(x_1, x_2, \ldots, x_n\) is
-- defined as:
--
-- \(HM = \frac{n}{\frac1{x_1} + \frac1{x_2} + \cdots + \frac1{x_n}}\)
--
-- \(HM = \left(\frac{\sum\limits_{i=1}^n x_i^{-1}}{n}\right)^{-1}\)
--
-- >>> harmonicMean = Fold.teeWith (/) length (lmap recip sum)
-- >>> harmonicMean = powerMeanFrac (-1)
--
-- See https://en.wikipedia.org/wiki/Harmonic_mean .
--
{-# INLINE harmonicMean #-}
harmonicMean :: (Monad m, Fractional a) => Fold m (a, Maybe a) a
harmonicMean = Fold.teeWith (/) Window.length (Window.lmap recip Window.sum)

-- | Geometric mean, defined as:
--
-- \(GM = \sqrt[n]{x_1 x_2 \cdots x_n}\)
--
-- \(GM = \left(\prod_{i=1}^n x_i\right)^\frac{1}{n}\)
--
-- or, equivalently, as the arithmetic mean in log space:
--
-- \(GM = e ^{{\frac{\sum_{i=1}^{n}\ln a_i}{n}}}\)
--
-- >>> geometricMean = exp <$> lmap log mean
--
-- See https://en.wikipedia.org/wiki/Geometric_mean .
{-# INLINE geometricMean #-}
geometricMean :: (Monad m, Floating a) => Fold m (a, Maybe a) a
geometricMean = exp <$> Window.lmap log mean

-- | The quadratic mean or root mean square (rms) of the numbers
-- \(x_1, x_2, \ldots, x_n\) is defined as:
--
-- \(RMS = \sqrt{ \frac{1}{n} \left( x_1^2 + x_2^2 + \cdots + x_n^2 \right) }.\)
--
-- >>> quadraticMean = powerMean 2
--
-- See https://en.wikipedia.org/wiki/Root_mean_square .
--
{-# INLINE quadraticMean #-}
quadraticMean :: (Monad m, Floating a) => Fold m (a, Maybe a) a
quadraticMean = powerMean 2

-------------------------------------------------------------------------------
-- Weighted Means
-------------------------------------------------------------------------------

-- XXX Is this numerically stable? We can use the kbn summation here.
-- | ewmaStep smoothing-factor old-value new-value
{-# INLINE ewmaStep #-}
ewmaStep :: Double -> Double -> Double -> Double
ewmaStep k x0 x1 = (1 - k) * x0 + k * x1

-- XXX Compute this in a sliding window?
--
-- | @ewma smoothingFactor@.
--
-- @ewma@ of an empty stream is 0.
--
-- Exponential weighted moving average, \(s_n\), of \(n\) values,
-- \(x_1,\ldots,x_n\), is defined recursively as:
--
-- \(\begin{align} s_0& = x_0\\ s_n & = \alpha x_{n} + (1-\alpha)s_{n-1},\quad n>0 \end{align}\)
--
-- If we expand the recursive term it becomes an exponential series:
--
-- \(s_n = \alpha \left[x_n + (1-\alpha)x_{n-1} + (1-\alpha)^2 x_{n-2} + \cdots + (1-\alpha)^{n-1} x_1 \right] + (1-\alpha)^n x_0\)
--
-- where \(\alpha\), the smoothing factor, is in the range \(0 <\alpha < 1\).
-- More the value of \(\alpha\), the more weight is given to newer values.  As
-- a special case if it is 0 then the weighted sum would always be the same as
-- the oldest value, if it is 1 then the sum would always be the same as the
-- newest value.
--
-- See https://en.wikipedia.org/wiki/Moving_average
--
-- See https://en.wikipedia.org/wiki/Exponential_smoothing
--
{-# INLINE ewma #-}
ewma :: Monad m => Double -> Fold m Double Double
ewma k = extract <$> Fold.foldl' step (Tuple' 0 1)

    where

    step (Tuple' x0 k1) x = Tuple' (ewmaStep k1 x0 x) k

    extract (Tuple' x _) = x

-- XXX It can perhaps perform better if implemented as a custom fold?
--
-- | @ewma n k@ is like 'ewma' but uses the mean of the first @n@ values and
-- then uses that as the initial value for the @ewma@ of the rest of the
-- values.
--
-- This can be used to reduce the effect of volatility of the initial value
-- when k is too small.
--
{-# INLINE ewmaAfterMean #-}
ewmaAfterMean :: Monad m => Int -> Double -> Fold m Double Double
ewmaAfterMean n k =
    Fold.concatMap (\i -> (Fold.foldl' (ewmaStep k) i)) (Fold.take n Fold.mean)

-- | @ewma n k@ is like 'ewma' but uses 1 as the initial smoothing factor and
-- then exponentially smooths it to @k@ using @n@ as the smoothing factor.
--
-- This is significantly faster than 'ewmaAfterMean'.
--
{-# INLINE ewmaRampUpSmoothing #-}
ewmaRampUpSmoothing :: Monad m => Double -> Double -> Fold m Double Double
ewmaRampUpSmoothing n k1 = extract <$> Fold.foldl' step initial

    where

    initial = Tuple' 0 1

    step (Tuple' x0 k0) x1 =
        let x = ewmaStep k0 x0 x1
            k = ewmaStep n k0 k1
        in Tuple' x k

    extract (Tuple' x _) = x

-------------------------------------------------------------------------------
-- Spread/Dispersion
-------------------------------------------------------------------------------

-- | @md n@ computes the mean absolute deviation (or mean deviation) in a
-- sliding window of last @n@ elements in the stream.
--
-- The mean absolute deviation of the numbers \(x_1, x_2, \ldots, x_n\) is:
--
-- \(MD = \frac{1}{n}\sum_{i=1}^n |x_i-\mu|\)
--
-- Note: It is expensive to compute MD in a sliding window. We need to
-- maintain a ring buffer of last n elements and maintain a running mean, when
-- the result is extracted we need to compute the difference of all elements
-- from the mean and get the average. Using standard deviation may be
-- computationally cheaper.
--
-- See https://en.wikipedia.org/wiki/Average_absolute_deviation .
--
-- /Pre-release/
{-# INLINE md #-}
md ::  MonadIO m => Fold m ((Double, Maybe Double), m (MA.Array Double)) Double
md =
    Fold.rmapM computeMD
        $ Fold.tee (Fold.lmap fst mean) (Fold.lmap snd Fold.last)

    where

    computeMD (mn, rng) =
        case rng of
            Just action -> do
                arr <- action
                Stream.fold Fold.mean
                    $ Stream.map (\a -> abs (mn - a))
                    $ Stream.unfold MA.read arr
            Nothing -> return 0.0

-- | The variance \(\sigma^2\) of a population of \(n\) equally likely values
-- is defined as the average of the squares of deviations from the mean
-- \(\mu\). In other words, second moment about the mean:
--
-- \(\sigma^2 = \frac{1}{n}\sum_{i=1}^n {(x_{i}-\mu)}^2\)
--
-- \(\sigma^2 = rawMoment(2) - \mu^2\)
--
-- \(\mu_2 = -(\mu'_1)^2 + \mu'_2\)
--
-- Note that the variance would be biased if applied to estimate the population
-- variance from a sample of the population. See 'sampleVariance'.
--
-- See https://en.wikipedia.org/wiki/Variance.
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE variance #-}
variance :: (Monad m, Fractional a) => Fold m (a, Maybe a) a
variance = Fold.teeWith (\p2 m -> p2 - m ^ (2 :: Int)) (rawMoment 2) mean

-- | Standard deviation \(\sigma\) is the square root of 'variance'.
--
-- This is the population standard deviation or uncorrected sample standard
-- deviation.
--
-- >>> stdDev = sqrt <$> variance
--
-- See https://en.wikipedia.org/wiki/Standard_deviation .
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE stdDev #-}
stdDev :: (Monad m, Floating a) => Fold m (a, Maybe a) a
stdDev = sqrt <$> variance

-- | Skewness \(\gamma\) is the standardized third central moment defined as:
--
-- \(\tilde{\mu}_3 = \frac{\mu_3}{\sigma^3}\)
--
-- The third central moment can be computed in terms of raw moments:
--
-- \(\mu_3 = 2(\mu'_1)^3 - 3\mu'_1\mu'_2 + \mu'_3\)
--
-- Substituting \(\mu'_1 = \mu\), and \(\mu'_2 = \mu^2 + \sigma^2\):
--
-- \(\mu_3 = -\mu^3 - 3\mu\sigma^2 + \mu'_3\)
--
-- Skewness is a measure of symmetry of the probability distribution. It is 0
-- for a symmetric distribution, negative for a distribution that is skewed
-- towards left, positive for a distribution skewed towards right.
--
-- For a normal like distribution the median can be found around
-- \(\mu - \frac{\gamma\sigma}{6}\) and the mode can be found around
-- \(\mu - \frac{\gamma \sigma}{2}\).
--
-- See https://en.wikipedia.org/wiki/Skewness .
--
{-# INLINE skewness #-}
skewness :: (Monad m, Floating a) => Fold m (a, Maybe a) a
skewness =
    toFold
        $ (\rm3 sd mu ->
            rm3 / sd ^ (3 :: Int) - 3 * (mu / sd) - (mu / sd) ^ (3 :: Int)
          )
        <$> Tee (rawMoment 3)
        <*> Tee stdDev
        <*> Tee mean

-- XXX We can compute the 2nd, 3rd, 4th raw moments by repeatedly multiplying
-- instead of computing the powers every time.
--
-- | Kurtosis \(\kappa\) is the standardized fourth central moment, defined as:
--
-- \(\tilde{\mu}_4 = \frac{\mu_4}{\sigma^4}\)
--
-- The fourth central moment can be computed in terms of raw moments:
--
-- \(\mu_4 = -3(\mu'_1)^4 + 6(\mu'_1)^2\mu'_2 - 4\mu'_1\mu'_3\ + \mu'_4\)
--
-- Substituting \(\mu'_1 = \mu\), and \(\mu'_2 = \mu^2 + \sigma^2\):
--
-- \(\mu_4 = 3\mu^4 + 6\mu^2\sigma^2 - 4\mu\mu'_3 + \mu'_4\)
--
-- It is always non-negative. It is 0 for a point distribution, low for light
-- tailed (platykurtic) distributions and high for heavy tailed (leptokurtic)
-- distributions.
--
-- \(\kappa >= \gamma^2 + 1\)
--
-- For a normal distribution \(\kappa = 3\sigma^4\).
--
-- See https://en.wikipedia.org/wiki/Kurtosis .
--
{-# INLINE kurtosis #-}
kurtosis :: (Monad m, Floating a) => Fold m (a, Maybe a) a
kurtosis =
    toFold
        $ (\rm4 rm3 sd mu ->
             ( 3 * mu ^ (4 :: Int)
            + 6 * mu ^ (2 :: Int) * sd ^ (2 :: Int)
            - 4 * mu * rm3
            + rm4) / (sd ^ (4 :: Int))
          )
        <$> Tee (rawMoment 4)
        <*> Tee (rawMoment 3)
        <*> Tee stdDev
        <*> Tee mean

-------------------------------------------------------------------------------
-- Estimation
-------------------------------------------------------------------------------

-- | Unbiased sample variance i.e. the variance of a sample corrected to
-- better estimate the variance of the population, defined as:
--
-- \(s^2 = \frac{1}{n - 1}\sum_{i=1}^n {(x_{i}-\mu)}^2\)
--
-- \(s^2 = \frac{n}{n - 1} \times \sigma^2\).
--
-- See https://en.wikipedia.org/wiki/Bessel%27s_correction.
--
{-# INLINE sampleVariance #-}
sampleVariance :: (Monad m, Fractional a) => Fold m (a, Maybe a) a
sampleVariance = Fold.teeWith (\n s2 -> n * s2 / (n - 1)) Window.length variance

-- | Sample standard deviation:
--
-- \(s = \sqrt{sampleVariance}\)
--
-- >>> sampleStdDev = sqrt <$> sampleVariance
--
-- See https://en.wikipedia.org/wiki/Unbiased_estimation_of_standard_deviation
-- .
--
{-# INLINE sampleStdDev #-}
sampleStdDev :: (Monad m, Floating a) => Fold m (a, Maybe a) a
sampleStdDev = sqrt <$> sampleVariance

-- | Standard error of the sample mean (SEM), defined as:
--
-- \( SEM = \frac{sampleStdDev}{\sqrt{n}} \)
--
-- See https://en.wikipedia.org/wiki/Standard_error .
--
-- /Space/: \(\mathcal{O}(1)\)
--
-- /Time/: \(\mathcal{O}(n)\)
{-# INLINE stdErrMean #-}
stdErrMean :: (Monad m, Floating a) => Fold m (a, Maybe a) a
stdErrMean = Fold.teeWith (\sd n -> sd / sqrt n) sampleStdDev Window.length

-------------------------------------------------------------------------------
-- Resampling
-------------------------------------------------------------------------------

{-# INLINE foldArray #-}
foldArray :: Storable a => Fold Identity a b -> Array a -> b
foldArray f = runIdentity . Stream.fold f . toStream

-- XXX Is this numerically stable? Should we keep the rounding error in the sum
-- and take it into account when subtracting?
--
-- | Given an array of @n@ items, compute mean of @(n - 1)@ items at a time,
-- producing a stream of all possible mean values omitting a different item
-- every time.
--
{-# INLINE jackKnifeMean #-}
jackKnifeMean :: (Monad m, Fractional a, Storable a) => Array a -> SerialT m a
jackKnifeMean arr = do
    let len = fromIntegral (length arr - 1)
        s = foldArray Fold.sum arr
     in Stream.map (\b -> (s - b) / len) $ toStream arr

-- | Given an array of @n@ items, compute variance of @(n - 1)@ items at a time,
-- producing a stream of all possible variance values omitting a different item
-- every time.
--
{-# INLINE jackKnifeVariance #-}
jackKnifeVariance :: (Monad m, Fractional a, Storable a) =>
    Array a -> SerialT m a
jackKnifeVariance arr = do
    let len = fromIntegral $ length arr - 1
        foldSums (s, s2) x = (s + x, s2 + x ^ (2 :: Int))
        (sum, sum2) = foldArray (Fold.foldl' foldSums (0.0, 0.0)) arr
        var x = (sum2 - x ^ (2 :: Int)) / len -  ((sum - x) / len) ^ (2::Int)
     in Stream.map var $ toStream arr

-- | Standard deviation computed from 'jackKnifeVariance'.
--
{-# INLINE jackKnifeStdDev #-}
jackKnifeStdDev :: (Monad m, Storable a, Floating a) =>
    Array a -> SerialT m a
jackKnifeStdDev = Stream.map sqrt . jackKnifeVariance

-- XXX This can be made more modular if the replicateM unfold can take count
-- from the seed.
--
-- | Randomly select elements from an array, with replacement, producing
-- a stream of the same size as the original array.
{-# INLINE resample #-}
resample :: (MonadIO m, Storable a) => Unfold m (Array a) a
resample = Unfold step inject

    where

    inject arr = liftIO $ do
        g <- createSystemRandom
        return $ (g, arr, length arr, 0)

    chooseOne g arr len = do
        i <- uniformRM (0, len - 1) g
        unsafeIndexIO i arr

    step (g, arr, len, idx) = liftIO $ do
        if idx >= len
        then return Stop
        else do
            e <- chooseOne g arr len
            return $ Yield e (g, arr, len, idx + 1)

-- | Resample an array multiple times and run the supplied fold on each
-- resampled stream, producing a stream of fold results. The fold is usually an
-- estimator fold.
{-# INLINE foldResamples #-}
foldResamples :: (MonadAsync m, Storable a) =>
       Int          -- ^ Number of resamples to compute.
    -> Array a      -- ^ Original sample.
    -> Fold m a b   -- ^ Estimator fold
    -> SerialT m b
foldResamples n arr fld =
    Stream.replicateM n (Unfold.fold fld resample arr)

-------------------------------------------------------------------------------
-- Probability Distribution
-------------------------------------------------------------------------------

-- | Determine the frequency of each element in the stream.
--
{-# INLINE frequency #-}
frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
frequency = Fold.classifyWith id Fold.length

-- | Find out the most frequently ocurring element in the stream and its
-- frequency.
--
{-# INLINE mode #-}
mode :: (Monad m, Ord a) => Fold m a (Maybe (a, Int))
mode = Fold.rmapM findMax frequency

    where

    fmax k v Nothing = Just (k, v)
    fmax k v old@(Just (_, v1))
        | v > v1 = Just (k, v)
        | otherwise = old

    findMax = return . foldrWithKey fmax Nothing

-------------------------------------------------------------------------------
-- Histograms
-------------------------------------------------------------------------------

-- | @binOffsetSize offset binSize input@. Given an integral input value,
-- return its bin index provided that each bin contains @binSize@ items and the
-- bins are aligned such that the 0 index bin starts at @offset@ from 0. If
-- offset = 0 then the bin with index 0 would have values from 0 to binSize -
-- 1.
--
-- This API does not put a bound on the number of bins, therefore, the number
-- of bins could be potentially large depending on the range of values.
--
{-# INLINE binOffsetSize #-}
binOffsetSize :: Integral a => a -> a -> a -> a
binOffsetSize offset binSize x = (x - offset) `div` binSize

data HistBin a = BelowRange | InRange a | AboveRange deriving (Eq, Show)

instance (Eq a, Ord a) => Ord (HistBin a) where
    compare BelowRange BelowRange = EQ
    compare BelowRange (InRange _) = LT
    compare BelowRange AboveRange = LT

    compare (InRange _) BelowRange = GT
    compare (InRange x) (InRange y)= x `compare` y
    compare (InRange _) AboveRange = LT

    compare AboveRange BelowRange = GT
    compare AboveRange (InRange _) = GT
    compare AboveRange AboveRange = EQ

-- | @binFromSizeN low binSize nbins input@. Classify @input@ into bins
-- specified by a @low@ limit, @binSize@ and @nbins@. Inputs below the lower
-- limit are classified into 'BelowRange' and inputs above the highest bin are
-- classified into 'AboveRange'. 'InRange' inputs are classified into bins
-- starting from bin index 0.
--
{-# INLINE binFromSizeN #-}
binFromSizeN :: Integral a => a -> a -> a -> a -> HistBin a
binFromSizeN low binSize nbins x =
    let high = low + binSize * nbins
     in if x < low
        then BelowRange
        else if x >= high
             then AboveRange
             else InRange ((x - low) `div` binSize)

-- | @binFromToN low high nbins input@. Like @binFromSizeN@ except that a range
-- of lower and higher limit is specified. @binSize@ is computed using the
-- range and @nbins@. @nbins@ is rounded to the range @0 < nbins < (high - low
-- + 1)@. @high >= low@ must hold.
--
{-# INLINE binFromToN #-}
binFromToN :: Integral a => a -> a -> a -> a -> HistBin a
binFromToN low high n x =
    let count = high - low + 1
        n1 = max n 1
        n2 = min n1 count
        binSize = count `div` n2
        nbins =
            if binSize * n2 < count
            then n2 + 1
            else n2
     in assert (high >= low) (binFromSizeN low binSize nbins x)

-- Use binary search to find the bin
--
-- | Classify an input value to bins using the bin boundaries specified in an
-- array.
--
-- /Unimplemented/
--
{-# INLINE binBoundaries #-}
binBoundaries :: -- Integral a =>
    Array.Array a -> a -> HistBin a
binBoundaries = undefined

-- | Given a bin classifier function and a stream of values, generate a
-- histogram map from indices of bins to the number of items in the bin.
--
-- >>> Stream.fold (histogram (binOffsetSize 0 3)) $ Stream.fromList [1..15]
-- fromList [(0,2),(1,3),(2,3),(3,3),(4,3),(5,1)]
--
{-# INLINE histogram #-}
histogram :: (Monad m, Ord k) => (a -> k) -> Fold m a (Map k Int)
histogram bin = Fold.classifyWith bin Fold.length


module Raytracer.Camera where

import Numeric.LinearAlgebra.Data (Vector, (|>), toList)
import Raytracer.Geometry (Ray(Ray), (|*), intersection)
import Numeric.LinearAlgebra (cross)
import Data.Array (listArray)
import Data.Monoid (mappend)
import Data.Tuple (swap)
import Data.Ix (range, Ix, index, inRange)
import Data.Function (on)

-- This is a camera. It's got a bunch of crap
-- So, first are settings. They're size of the surface in our units (w h) then resolution of the surface (w h).
-- Then the next two are the position of the camera, and the vector it's pointing in
-- The maybe is Focal Distance. If it's Nothing, then we do an orthographic representation
--   If it's a number, then it's the number of units back that the focal point is.
-- We assume for now that the camera is always level relative to (1, 0, 1) (No roll)
data Camera = Camera Double Double Int Int (Maybe Double) (Vector Double) (Vector Double) deriving Show

data Point = Point Int Int deriving (Eq, Show)

instance Ord Point where
	compare (Point a b) (Point c d) = mappend (compare b d) (compare a c)

instance Ix Point where
	range ((Point a b), (Point c d)) = (range (b, d)) >>= (\y -> (range (a, c)) >>= (\x -> return $ Point x y))
	index ((Point a b), (Point c d)) (Point e f) = (((c - a) + 1) * (f - b)) + (e - a)
	inRange ((Point a b), (Point c d)) (Point e f) = (inRange (a, c) e) && (inRange (b, d) f)

norm :: Vector Double -> Double
norm = sqrt.sum.(map (^2)).toList

unitize :: Vector Double -> Vector Double
unitize v = (1/(norm v)) |* v

-- This camera is currently orthographic, rather than perspective
calculate_ray (Camera width height wres hres focus pos direction) = (\pos -> Ray (ray_direction pos) pos) . ray_pos
	where
	-- This gives me the width axis of my image
	-- It is acheived by a cross product of my looking direction and a vertical axis
	width_axis = unitize $ cross (3 |>  [0, 1, 0]) direction
	-- This gives me the height axis of my image
	height_axis = unitize $ cross direction width_axis
	-- This function computes a vector based on progress along a vector
	partial_vector vec steps current_step = (on (/) fromIntegral current_step steps) |* vec
	-- This function computes the position of a ray given its place in the matrix
	ray_pos (Point x y) = (partial_vector (width |* width_axis) wres x) + (partial_vector (height |* height_axis) hres y) + pos + (3 |> [-width / 2, -height/2, 0])
	-- This function takes in the point and returns the direction of the ray at that point
	-- If it's perspective it involves subtracting from the focal point. If it's orthographic it's just always direction.
	ray_direction = maybe (const direction) (\len -> (flip (-)) $ pos - (len |* direction)) focus

-- This function, finally, generates an array of wres and hres full of the proper ray at each point
calculate_rays camera = build_array (calculate_ray camera) (Point 1 1, Point wres hres)
	where
	(Camera _ _ wres hres _ _ _) = camera
	-- This function takes a function expecting (x,y) and bounds and builds an array by calling the function at each point
	build_array func bounds = listArray bounds $ map func $ range bounds

fire_ray mesh ray = intersection ray mesh
fire_rays rays mesh = fmap (fire_ray mesh) rays


{-# LANGUAGE UnicodeSyntax, DataKinds, TypeOperators, KindSignatures,
             TypeInType, GADTs, MultiParamTypeClasses, FunctionalDependencies,
             TypeFamilies, AllowAmbiguousTypes, FlexibleInstances,
             UndecidableInstances, InstanceSigs, TypeApplications,
             ScopedTypeVariables, EmptyCase, FlexibleContexts,
             RankNTypes, LambdaCase
#-}

-- Implementation of density matrix interpretation of quantum computation
module LinTrans2 where

import Prelim

import qualified Data.Complex as C
import Data.Complex (Complex(..))
import Control.Monad.State.Lazy
import Data.Singletons
import Data.Maybe (fromJust)
--import Data.Tuple hiding (swap)
import Unsafe.Coerce
import Data.Constraint (Dict(..))
import Data.List ( (\\) )
import qualified Debug.Trace as Tr

-- I don't like the show instance for ℂ
newtype ℂ = ℂ (Complex Double)

liftC :: (Complex Double -> Complex Double) -> ℂ -> ℂ
liftC f (ℂ c) = ℂ $ f c

i :: ℂ
i = ℂ $ 0 :+ 1
conjugate = liftC C.conjugate
instance Num ℂ where
  ℂ m + ℂ n = ℂ $ m + n
  ℂ m - ℂ n = ℂ $ m - n
  ℂ m * ℂ n = ℂ $ m * n
  abs = liftC abs
  signum = liftC signum
  fromInteger = ℂ . fromInteger
instance Fractional ℂ where
  fromRational = ℂ . fromRational
  ℂ m / ℂ n    = ℂ $ m / n
instance Floating ℂ where
  pi    = ℂ $ pi
  exp   = liftC exp
  log   = liftC log
  sin   = liftC sin
  cos   = liftC cos
  asin  = liftC asin
  acos  = liftC acos
  atan  = liftC atan
  sinh  = liftC sinh
  cosh  = liftC cosh
  asinh = liftC asinh
  acosh = liftC acosh
  atanh = liftC atanh
instance Show ℂ where
  show (ℂ (α :+ β)) =   if β == 0 then show α
                        else if α == 0 then show β ++ "i"
                        else show α ++ " + " ++ show β ++ "i"

--------------
-- Matrices --
--------------

-- A matrix Matrix (m,n) f of dimension 2^m × 2^n is a function f from pairs of
-- nats to ℂ
data Matrix = Matrix { dim :: (Int,Int)
                     , idx :: (Int, Int) -> ℂ
                     }
dom :: Matrix -> [Int]
dom (Matrix (m,_) _) = [0..(2^m)-1]

cod :: Matrix -> [Int]
cod (Matrix (_,n) _) = [0..(2^n)-1]

mkMatrix :: (Int,Int) -> [[ℂ]] -> Matrix
mkMatrix (m,n) ls = Matrix (m,n) $ \(i,j) -> (ls !! i) !! j

--------------
-- Printing --
--------------

rows :: Matrix -> [[ℂ]]
rows mat = f <$> dom mat
  where
    f i = (\j -> idx mat (i,j)) <$> cod mat
showRows :: Matrix -> [String]
showRows mat = show <$> rows mat
instance Show Matrix where
  show mat = unlines $ showRows mat

---------------------------------------
-- Primitive matrices and operations --
---------------------------------------

ident :: Int -> Matrix
ident n = Matrix (n,n) $ \(i,j) -> if i == j then 1 else 0

transpose :: Matrix -> Matrix
transpose (Matrix (m,n) f) = Matrix (n,m) $ \(i,j) -> f (j,i)

dot :: Matrix -> Matrix -> ℂ
dot v1 v2 | snd (dim v1) == fst (dim v2) =
            foldl (+) 0 $ (\x -> idx v1 (0,x) * idx v2 (x,0)) <$> dom v2
          | otherwise = error "Mismatched dimensions when taking dot product"

row :: Int -> Matrix -> Matrix
row i (Matrix (_,n) f) = Matrix (0,n) $ \(_,j) -> f(i,j)

col :: Int -> Matrix -> Matrix
col j (Matrix (m,_) f) = Matrix (m,0) $ \(i,_) -> f(i,j)

plusM :: Matrix -> Matrix -> Matrix
plusM (Matrix dim1 f1) (Matrix dim2 f2) | dim1 == dim2 =
       Matrix dim1 $ \(i,j) -> f1(i,j) + f2(i,j)
plusM _ _ | otherwise    = error "Mismatched dimensions when adding matrices"

multM :: Matrix -> Matrix -> Matrix
multM mat1 mat2 | n1 == m2 =
    Matrix (m1,n2) $ \(i,j) -> dot (row i mat1) (col j mat2)
  where
    (m1,n1) = dim mat1
    (m2,n2) = dim mat2
multM _ _ | otherwise = error "Mismatched dimensions when multiplying matrices"

kron :: Matrix -> Matrix -> Matrix
kron (Matrix (m1,n1) f1) (Matrix (m2,n2) f2) =
    Matrix (m1*n1, m2*n2) $ \(i,j) ->
           f1 (i `div` 2^m2, j `div` 2^n2) * f2 (i `mod` 2^m2, j `mod` 2^n2)

zero n = Matrix (n,n) $ \_ -> 0

scale :: ℂ -> Matrix -> Matrix
scale c (Matrix (m,n) f) = Matrix (m,n) $ \(i,j) -> c * f(i,j)

----------------------
-- Quantum Matrices --
----------------------

ket0 = mkMatrix (1,0) [[1],[0]]
ket1 = mkMatrix (1,0) [[0],[1]]
density0 = ket0 * transpose ket0
density1 = ket1 * transpose ket1
newD True  = density1
newD False = density0

hadamard = mkMatrix (1,1) [[1/sqrt 2, 1/sqrt 2]
                          ,[1/sqrt 2, -1/sqrt 2]]

pauliX = mkMatrix (1,1) [[0,1]
                       ,[1,0]]
pauliY = mkMatrix (1,1) [[0,-1]
                       ,[i,0]]
pauliZ = mkMatrix (1,1) [[1,0]
                       ,[0,-1]]

cnot = mkMatrix (2,2) [[1,0,0,0]
                      ,[0,1,0,0]
                      ,[0,0,0,1]
                      ,[0,0,1,0]]

------------------
-- Permutations --
------------------

-- encode and decode modulo a key k
encode :: Int -> [Int] -> Int
encode k [] = 0
encode k (i : ls) = i + k * encode k ls

decode :: Int -> Int -> [Int]
decode k i | i == 0 = []
decode k i | otherwise = i `mod` k : decode k (i `div` k)

fromAssocList :: [(Int,Int)] -> Int -> Int
fromAssocList []           = id
fromAssocList ((a,b) : ls) = assocFun a b . fromAssocList ls
  where
    assocFun a b i = if i == a then b else i

-- Check if two lists are equal up to the permutation specified, i.e.
-- if permuting ls1 by f is equal to ls2
isPermutation :: Eq a => (Int -> Int) -> [a] -> [a] -> Bool
isPermutation f ls1 ls2 = all (\i -> ls1 !! i == ls2 !! f i) [0..len]
  where
    len = length ls1

swapFun :: Int -> (Int -> Int) -> Matrix
swapFun k f = Matrix (k,k) $ \(i,j) ->
    if isPermutation f (decode k i) (decode k j) then 1 else 0

-- swap takes a list of qubit/bit variables and produces a matrix that permutes
-- those qubits.
-- i.e. swap ls |φ_0 ⋯ φ_n⟩ = |φ_f(0) ⋯ ρ_f(n)⟩
-- where f = fromAssocList ls
swap :: Int -> [Int] -> Matrix
swap k ls = swapFun k $ fromAssocList $ zip [0..] ls

-------------------
-- Density Monad --
-------------------

-- A density monad is a nondeterminism state monad:
-- Density -> [Density]
-- An element op of type DensityMonad corresponds to the superoperator
-- \ρ -> ∑ (op ρ)
type DensityMonad = StateT Matrix []

newM :: Bool -> DensityMonad Int
newM b = do
    ρ ← get
    put $ ρ `kron` newD b
    return $ fst (dim ρ)

-- runQ applies the superoperator to the identity density matrix of size 1.
runQ :: DensityMonad a -> [(a,Matrix)]
runQ m = runStateT m (ident 0)

-- getDensity combines the result of runQ into a single density matrix
getDensity :: DensityMonad a -> Matrix
getDensity m = foldr f (zero n) ls
  where
    ls = runQ m
    n = fst . dim . snd $ head ls
    f (_,ρ) ρ0 = ρ + ρ0

-----------------
-- Application --
-----------------

super :: Matrix -> DensityMonad ()
super mat = do ρ ← get
               put $ mat * ρ * transpose mat

applyMatrix :: Matrix -> [Int] -> DensityMonad ()
applyMatrix mat ls = do super $ swap n ls  -- moves the relevant qubits to the front
                        super mat -- applies operation
                        super . transpose $ swap n ls -- moves the qubits back to where they were
  where
    (n,_) = dim mat

branch :: DensityMonad a -> DensityMonad a -> DensityMonad a
branch m0 m1 = StateT $ \ρ -> [runStateT m0 ρ, runStateT m1 ρ] >>= id

measM :: Int -> DensityMonad Bool
measM i = branch (applyMatrix density0 [i] >> return False) (applyMatrix density1 [i] >> return True)

------------------
-- Num instance --
------------------

instance Num Matrix where
  (+) = plusM
  mat1 - mat2 = mat1 `plusM` scale (-1) mat2
  (*) = multM
  abs = undefined
  signum mat = undefined -- trace?
  fromInteger n = ident $ fromInteger n

-------------------
-- Show Instance --
-------------------

instance Show (DensityMonad a) where
  show = show . getDensity

{-

------------------
-- Num instance --
------------------

instance (SingI m, SingI n) => Num (Matrix m n) where
  (+) mat1 mat2 (i,j) = mat1(i,j) + mat2(i,j)
  (-) mat1 mat2 (i,j) = mat1(i,j) - mat2(i,j)
  (*) mat1 mat2 (i,j) = mat1(i,j) * mat2(i,j)
  abs mat (i,j) = abs $ mat(i,j)
  signum mat = undefined -- trace mat
  fromInteger n = matrix . repeat $ fromIntegral n
instance Num Density where
  Density n1 mat1 + Density n2 mat2 =
    withSingI (two `raiseToSNat` n1) $
    withSingI (two `raiseToSNat` n2) $
    case eqSNat n1 n2 of
      Left Dict -> Density n1 $ mat1 + mat2
      Right _   -> error $ "Cannot add mismatched matrices "
                            ++ show mat1 ++ " and " ++ show mat2
  _ - _ = undefined
  _ * _ = undefined
  abs _ = undefined
  signum = undefined
  fromInteger n = undefined

-- Debugging help
m0 = square @Two [1,0,0,0
                 ,0,1,0,0
                 ,0,0,0,1
                 ,0,0,1,0]
rho = Density three $ \case
  (0,0) -> -0.5
  (0,2) ->  0.5
  (3,0) -> -0.5
  (3,2) ->  0.5
  (_,_) -> 0

-}


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Pendulum.Model (
  Pendulum(..),
  mkPendlumSystem,
  stepPendulum
  ) where

import qualified Data.Vector.Sized as V
import           GHC.TypeLits (KnownNat)
import           Numeric.Hamilton
import           Numeric.LinearAlgebra.Static

data Pendulum n = Pendulum
  { _lengthes :: V.Vector n Double
  , _masses   :: V.Vector n Double
  , _config   :: Config n
  }

gravity :: Floating a => a
gravity = 9.8

toFloating :: Floating a => Double -> a
toFloating = fromRational . toRational

duplicate :: [a] -> [a]
duplicate = concatMap (\x -> [x, x])

merge :: [a] -> [a] -> [a]
merge []     ys     = ys
merge xs     []     = xs
merge (x:xs) (y:ys) = x : y : merge xs ys

evenElems :: [a] -> [a]
evenElems []       = []
evenElems (_:[])   = []
evenElems (_:x:xs) = x : evenElems xs

toVector :: forall a m. KnownNat m => [a] -> V.Vector m a
toVector = (\(Just x) -> x) . V.fromList

mkPendlumSystem :: forall m n. (KnownNat m, KnownNat n) => Pendulum n -> System m n
mkPendlumSystem p =
  let lengthes = V.toList (_lengthes p)
      masses   = V.toList (_masses p)
      masses' :: R m
      masses' = vector (duplicate masses)
      coordinates :: Floating a => V.Vector n a -> V.Vector m a
      coordinates v =
        let ths = V.toList v
            ls = map toFloating lengthes
            xs =                scanl1 (+) (zipWith (*) ls (map sin ths))
            ys = (* (-1.0)) <$> scanl1 (+) (zipWith (*) ls (map cos ths))
         in toVector (merge xs ys)
      potential :: Floating a => V.Vector m a -> a
      potential v =
        let ys = evenElems (V.toList v)
            ms = map toFloating masses
         in gravity * foldl1 (+) (zipWith (*) ms ys)
   in mkSystem' masses' coordinates potential

stepPendulum :: forall m n. (KnownNat m, KnownNat n) => Double -> System m n -> Pendulum n -> Pendulum n
stepPendulum dt system pd =
  let phase  = toPhase system (_config pd)
      phase' = stepHam dt system phase
   in pd {_config = fromPhase system phase'}


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
module Statistics.Distribution.Random.Normal (normal) where

import Control.Monad
import Data.Bits
import Data.Word
import qualified Data.Vector.Unboxed as I
import Random.CRI
import Statistics.Distribution.Random.Uniform

-- Copied from mwc-random.

data T = T {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Generate a normally distributed random variate.
--
-- The implementation uses Doornik's modified ziggurat algorithm.
-- Compared to the ziggurat algorithm usually used, this is slower,
-- but generates more independent variates that pass stringent tests
-- of randomness.
normal :: (Source m g Double, Source m g Word32) => g m -> m Double
normal gen = loop
  where
    loop = do
      u  <- (subtract 1 . (*2)) `liftM` uniform gen
      ri <- uniform gen
      let i  = fromIntegral ((ri :: Word32) .&. 127)
          bi = I.unsafeIndex blocks i
          bj = I.unsafeIndex blocks (i+1)
      if abs u < I.unsafeIndex ratios i
        then return $! u * bi
        else if i == 0
        then normalTail (u < 0)
        else do
          let x  = u * bi
              xx = x * x
              d  = exp (-0.5 * (bi * bi - xx))
              e  = exp (-0.5 * (bj * bj - xx))
          c <- uniform gen
          if e + c * (d - e) < 1
            then return x
            else loop
    blocks = let f = exp (-0.5 * r * r)
             in (`I.snoc` 0) . I.cons (v/f) . I.cons r .
                I.unfoldrN 126 go $! T r f
      where
        go (T b g)   = let !u = T h (exp (-0.5 * h * h))
                           h  = sqrt (-2 * log (v / b + g))
                       in Just (h, u)
        v            = 9.91256303526217e-3
    {-# NOINLINE blocks #-}
    r                = 3.442619855899
    ratios           = I.zipWith (/) (I.tail blocks) blocks
    {-# NOINLINE ratios #-}
    normalTail neg  = tailing
      where tailing  = do
              x <- ((/r) . log) `liftM` uniform gen
              y <- log          `liftM` uniform gen
              if y * (-2) < x * x
                then tailing
                else return $! if neg then x - r else r - x
{-# INLINE normal #-}


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import qualified Statistics.Distribution as Stat
import qualified Statistics.Distribution.Normal as Stat
import qualified Statistics.Distribution.Uniform as Stat
import Control.Monad.Primitive
import qualified System.Random.MWC as Rand
import Control.Monad.State.Strict
import Control.Monad.Reader
import Lens.Micro
import Types
import qualified Data.List as List
import Options.Generic

-- Contraints that allow for random number generation
type BanditM m = (MonadIO m, MonadReader (Rand.Gen (PrimState IO)) m)

-- |Generate a list of random bandits using dist to generate their parameters.
randomBandits ::
  (BanditM m, Stat.ContGen d1, Stat.ContGen d2)
  => Int
  -> d1
  -> d2
  -> m [Bandit BanditDist]
randomBandits count avgDist stdDevDist =
  replicateM count randomBandit
  where
    randomBandit = do
      gen <- ask
      mean <- liftIO $ Stat.genContVar avgDist gen
      stdDev <- liftIO $ Stat.genContVar stdDevDist gen
      return $ Bandit (Stat.normalDistr mean (abs stdDev))

-- |Pulls the handle and returns the reward
pullArm ::
  (BanditM m, Stat.ContGen d)
  => Bandit d
  -> m BanditStats
pullArm (Bandit d) = do
  gen <- ask
  reward <- liftIO $ Stat.genContVar d gen
  return $ BanditStats reward 1

simulate ::
  BanditM m
  => SimState -- ^ Starting state
  -> Int      -- ^ Number of rounds
  -> (SimState -> m Int) -- ^ Choosing function
  -> m SimState
simulate simState 0 _ = return simState
simulate simState rounds f = do
  chosenIndex <- f simState
  let Just (chosenStat, chosenBandit) = simState ^? bandits. ix chosenIndex -- Pick a bandit
  newStat <- pullArm chosenBandit
  let newState = simState
                 & (bandits . ix chosenIndex) . _1 .~ (chosenStat `mappend` newStat) -- Update stats
                 & roundsCount %~ (+1) -- Update rounds count
                 & regretHist %~ (\t -> (regret simState) : t)
  simulate newState (rounds - 1) f

-- | Initialize the state but creating the specified number of bandits and pulling each
-- | one once.
initState ::
  BanditM m
  => Int -- Number of bandits
  -> m SimState
initState count = do
  initBandits <- randomBandits count (Stat.uniformDistr 0 50) (Stat.normalDistr 0 5)
  let initStats = replicate count mempty
  return $ SimState (zip initStats initBandits) 0 []

regret :: SimState -> Double
regret simState =
  let allTimesPulled = fromIntegral <$> (simState ^.. bandits . each . _1 . timesPulled)
      allMeans = simState ^.. bandits . each . _2 . distLens . meanLens
      actualExpectedReturn = sum (zipWith (*) allTimesPulled allMeans)
      maxMean = maximum (simState ^.. bandits . each . _2 . distLens . meanLens)
      rounds = fromIntegral $ simState ^. roundsCount
      maxExpectedReturn = maxMean * rounds
  in
    maxExpectedReturn - actualExpectedReturn

type Concrete a = (ReaderT (Rand.Gen (PrimState IO)) IO) a

run ::
  Int
  -> Int
  -> (SimState -> Concrete Int)
  -> Concrete SimState
run rounds count f = do
  startingState <- initState count
  simulate startingState rounds f

-- | Return the index of the largest item as determined by the function
argMaxIndex :: (Ord b, Num b) => [a] -> (a -> b) -> Maybe Int
argMaxIndex as f =
  let magnitudes = f <$> as
      indexAndMagnitude = zip [0..] magnitudes
      biggestToSmallest = (List.sortOn (negative . snd) indexAndMagnitude)
  in biggestToSmallest ^? _head . _1
  where
    negative x = -x -- Swap the sign, so it's sorted biggest to smallest.

-- | Choose the bandit with the highest Upper Confidence Bound
choose :: SimState -> Int
choose simState =
  let banditStats = simState ^.. bandits . each . _1
      curRound = simState ^. roundsCount
      Just maxIndex = argMaxIndex banditStats (ucb curRound)
  in maxIndex

ucb :: Int -> BanditStats -> Double
ucb curRound banditStats =
  let totalTimesPulled = fromIntegral (banditStats ^. timesPulled)
      empiricalMean = (banditStats ^. totalReward) / totalTimesPulled
      t = fromIntegral curRound
      confidence = sqrt ((2.0 * log (1 + (t * ((log t) ^ 2)))) / totalTimesPulled)
      in
    if totalTimesPulled /= 0 then empiricalMean + confidence else maxDouble
  where
    maxDouble = 100000000 -- TODO: What's the actual max double?

-- | Choose a random bandit.
chooseRandom :: BanditM m => SimState -> m Int
chooseRandom simState = do
    let count = fromIntegral $ length $ simState ^. bandits
    let d = Stat.uniformDistr 0 (count - 1)
    gen <- ask
    rand <- liftIO $ Stat.genContVar d gen
    return $ truncate rand

-- | Pick the bandit with the best empirical mean.
chooseBestMean :: SimState -> Int
chooseBestMean simState =
  let banditStats = simState ^.. bandits . each . _1
      Just maxIndex = argMaxIndex banditStats avgReturn
  in
    maxIndex

-- | Pick the bandit with the highest mean reward.
godMode :: SimState -> Int
godMode simState =
  let bandits_ = simState ^.. bandits . each . _2
      Just maxIndex = argMaxIndex bandits_ (\(Bandit d) -> Stat.mean d)
  in
    maxIndex

avgReturn :: BanditStats -> Double
avgReturn (BanditStats reward pulled) = reward / (fromIntegral pulled)

main :: IO ()
main = do
  (opts :: BanditOpts Unwrapped) <-unwrapRecord "Multi-armed bandit simulator."
  rand <- Rand.createSystemRandom
  finalState <- runReaderT (run (optRounds opts) (optBandits opts) (return . godMode)) rand
  _ <- traverse (putStrLn . show) (reverse (finalState ^. regretHist))
  return ()

data BanditOpts w =
  BanditOpts { optRounds :: w ::: Int <?> "Number of rounds to simulate."
             , optBandits :: w ::: Int <?> "Number of bandits."
             -- , optMeanMean :: w ::: Double <?> "Mean of the mean."
             -- , optStdDevMean :: w ::: Double <?> "StdDev of the mean."
             -- , optMeanStdDev :: w ::: Double <?> "Mean of the StdDev."
             -- , optStdDevStdDev :: w ::: Double <?> "StdDev of the StdDev."
             } deriving (Generic)

instance ParseRecord (BanditOpts Wrapped)


-- This file is part of Quipper. Copyright (C) 2011-2014. Please see the
-- file COPYRIGHT for a list of authors, copyright holders, licensing,
-- and other details. All rights reserved.
--
-- ======================================================================

{-# LANGUAGE NoMonomorphismRestriction #-}

-- | This module contains an implementation of the oracle and its
-- automatic lifting to quantum circuits using Template Haskell.
module Algorithms.QLS.TemplateOracle where

import Data.Complex
import Libraries.Auxiliary

import Control.Monad

import QuipperLib.Arith hiding (template_symb_plus_)
import Algorithms.QLS.Utils
import Algorithms.QLS.QDouble
import Algorithms.QLS.RealFunc
import Algorithms.QLS.QSignedInt
import Algorithms.QLS.CircLiftingImport

import Quipper
import Quipper.CircLifting

-- | Lifted version of @'any'@ (local version).
build_circuit
local_any :: (a -> Bool) -> [a] -> Bool
local_any f l = case l of
            []    -> False
            (h:t) -> if (f h) then True else local_any f t

-- | Auxiliary function.
build_circuit
itoxy :: Int -> Int -> Int -> (Int, Int)
itoxy i nx ny =
   if (i <= 0) then (0,0)
   else if (i > (nx-1)*ny + nx*(ny-1)) then (0,0)
   else ((mod (i - 1) (2*nx - 1)) + 1, ceiling ((fromIntegral i)/(2.0*(fromIntegral nx)-1.0)))

-- | The function sin /x/ \/ /x/.
build_circuit
sinc :: Double -> Double
sinc x = if (x /= 0.0) then (sin x) / x else 1.0

-- | Auxiliary function.
build_circuit
edgetoxy :: Int -> Int -> Int -> (Double, Double)
edgetoxy e nx ny =
  let (ex,ey) = itoxy e nx ny in
     if (ex < nx) then ((fromIntegral ex) + 0.5, fromIntegral ey)
     else (fromIntegral (ex - nx + 1), (fromIntegral ey) + 0.5)

-- | Auxiliary function. The inputs are:
--
-- * /y1/ :: 'Int' - Global edge index of row index of desired matrix element;
--
-- * /y2/ :: 'Int' - Global edge index of column index of desired matrix element;
--
-- * /nx/ :: 'Int' - Number of vertices left to right;
--
-- * /ny/ :: 'Int' - Number of vertices top to bottom;
--
-- * /lx/ :: 'Double' - Length of horizontal edges (distance between vertices in /x/ direction);
--
-- * /ly/ :: 'Double' - Length of vertical edges (distance between vertices in /y/ direction);
--
-- * /k/ :: 'Double' - Plane wave wavenumber.
--
-- The output is the matrix element /A/(/y1/, /y2/).
build_circuit
calcmatrixelement :: --
	-- Inputs
	Int 	-- int y1 - Global edge index of row index of desired matrix element
	-> Int 	-- int y2 - Global edge index of column index of desired matrix element
	-> Int 	-- int nx - Number of vertices left to right
	-> Int	-- int ny - Number of vertices top to bottom
	-> Double -- float lx - Length of horizontal edges (distance between vertices in x direction)
	-> Double -- float ly - Length of vertical edges (distance between vertices in y direction)
	-> Double -- float k - Plane wave wavenumber
	-- Outputs
	-> Complex Double -- float A - Matrix element A(y1,y2)
calcmatrixelement y1 y2 nx ny lx ly k =
	let (xg1, yg1) = itoxy y1 nx ny in
        let (xg2, yg2) = itoxy y2 nx ny in
        let b = if ( (y1==y2) && (xg1 >= nx) ) then ly/lx - k*k*lx*ly/3.0
                -- B11 and B22
                else if ( (y1==y2) && (xg1<nx) ) then lx/ly - k*k*lx*ly/3.0
                -- B12 and B21
                else if ( (abs(yg1-yg2) == 1) && (abs(xg1-xg2) == 0) && (xg1<nx) ) then -lx/ly + k*k*lx*ly/12.0
                -- B34 and B43
                else if ( (abs(yg1-yg2)==0) && (abs(xg1-xg2) == 1) && (xg1>=nx) ) then -ly/lx + k*k*lx*ly/12.0
                -- B13
                else if ( (yg1==(yg2+1)) && (xg1==(xg2-nx+1)) && (xg2>=nx) ) then -1.0
                -- B31
                else if ( (yg2==(yg1+1)) && (xg2==(xg1-nx+1)) && (xg1>=nx) ) then -1.0
                -- B14
                else if ( (yg1==(yg2+1)) && (xg1==(xg2-nx)) && (xg1<nx) ) then 1.0
                -- B41
                else if ( (yg2==(yg1+1)) && (xg2==(xg1-nx)) && (xg2<nx) ) then 1.0
                -- B42
                else if ( (yg1==yg2) && (xg1==(xg2+nx)) && (xg1>=nx) ) then -1.0
                -- B24
                else if ( (yg2==yg1) && (xg2==(xg1+nx)) && (xg2>=nx) ) then -1.0
                -- B32
                else if ( (yg1==yg2) && (xg1==(xg2+nx-1)) && (xg2<nx) ) then 1.0
                -- B23
                else if ( (yg2==yg1) && (xg2==(xg1+nx-1)) && (xg1<nx) ) then 1.0
                else -1.0
        in
        let c = if ( (y1==y2) && ( (xg1==nx) || (xg1==(2*nx-1)) ) ) then 0.0 :+ (k*ly)
                else if  ( (y1==y2) && ( ((yg1==1) && (xg1<nx)) || ((yg1==ny) && (xg1<nx)) ) ) then 0.0 :+ (k*lx)
                else 0.0 :+ 0.0
	in (b :+ 0.0) + c

-- | Auxiliary function.
build_circuit
get_edges l = case l of
                [] -> []
                (x:tt) -> case tt of
                   [] -> []
                   (y:t) -> (x,y):(get_edges (y:t))

-- | Auxiliary function.
build_circuit
checkedge :: Int -> [(Double,Double)] -> Int -> Int -> Bool
checkedge e scatteringnodes nx ny =
     let (xi,yi) = edgetoxy e nx ny in
     let test_elt ((x1,y1),(x2, y2)) = xi >= x1 && yi >= y1 && xi <= x2 && yi <= y2 in
     let half = take_half scatteringnodes in
     let test_list = local_any test_elt (get_edges half) in (not test_list)

-- | Oracle /r/.
build_circuit
calcRweights :: Int -> Int -> Int -> Double -> Double -> Double -> Double -> Double -> Complex Double
calcRweights y nx ny lx ly k theta phi =
     let (xc',yc') = edgetoxy y nx ny in
     let xc = (xc'-1.0)*lx - ((fromIntegral nx)-1.0)*lx/2.0 in
     let yc = (yc'-1.0)*ly - ((fromIntegral ny)-1.0)*ly/2.0 in
     let (xg,yg) = itoxy y nx ny in

     if (xg == nx) then

         let i = (mkPolar ly (k*xc*(cos phi)))*
                 (mkPolar 1.0 (k*yc*(sin phi)))*
                 ((sinc (k*ly*(sin phi)/2.0)) :+ 0.0) in

         let r = ( cos(phi) :+ k*lx )*((cos (theta - phi))/lx :+ 0.0) in i * r

     else if (xg==2*nx-1) then

         let i = (mkPolar ly (k*xc*cos(phi)))*
                 (mkPolar 1.0 (k*yc*sin(phi)))*
                 ((sinc (k*ly*sin(phi)/2.0)) :+ 0.0) in

         let r = ( cos(phi) :+ (- k*lx))*((cos (theta - phi))/lx :+ 0.0) in i * r

     else if ( (yg==1) && (xg<nx) ) then

         let i = (mkPolar lx (k*yc*sin(phi)))*
                 (mkPolar 1.0 (k*xc*cos(phi)))*
                 ((sinc (k*lx*(cos phi)/2.0)) :+ 0.0) in

         let r = ( (- sin phi) :+ k*ly )*((cos(theta - phi))/ly :+ 0.0) in i * r

     else if ( (yg==ny) && (xg<nx) ) then

         let i = (mkPolar lx (k*yc*sin(phi)))*
                 (mkPolar 1.0 (k*xc*cos(phi)))*
                 ((sinc (k*lx*(cos phi)/2.0)) :+ 0.0) in

         let r = ( (- sin phi) :+ (- k*ly) )*((cos(theta - phi)/ly) :+ 0.0) in i * r

     else 0.0 :+ 0.0

-- | Auxiliary function for oracle /A/.
build_circuit
convertband :: Int -> Int -> Int -> Int -> Int
convertband y b nx ny =
 let nedges = (nx - 1)*ny + nx*(ny - 1) in
 let (ex,ey) = itoxy y nx ny in
 let x = if ( (ex < nx) && (ey /= 1) ) then
           case b of
                1 -> y-2*nx+1
                2 -> y-nx
                3 -> y-nx+1
                5 -> y
                7 -> y+nx-1
                8 -> y+nx
                9 -> y+2*nx-1
                _ -> -1
         else if ( (ex < nx) && (ey == 1) ) then
           case b of
                5 -> y
                7 -> y+nx-1
                8 -> y+nx
                9 -> y+2*nx-1
                _ -> -1
         else if ( (ex >= nx) && (ex /= nx) && (ex /= 2*nx-1) ) then
           case b of
                    2 -> y-nx
                    3 -> y-nx+1
                    4 -> y-1
                    5 -> y
                    6 -> y+1
                    7 -> y+nx-1
                    8 -> y+nx
                    _ -> -1
         else if ( (ex >= nx) && (ex == nx) ) then
                 case b of
                    3 -> y-nx+1
                    5 -> y
                    6 -> y+1
                    8 -> y+nx
                    _ -> -1
         else if ( (ex >= nx) && (ex == 2*nx-1) ) then
                case b of
                    2 -> y-nx
                    4 -> y-1
                    5 -> y
                    7 -> y+nx-1
                    _ -> -1
         else -1
 in if ( (x < 1) || (x > nedges) ) then -1 else x

-- | Oracle /A/. It is equivalent to the Matlab function
-- 'getBandNodeValues'.
--
-- 'getNodeValuesMoreOutputs v b ...'  outputs the node of the edge
-- connected to vertex v in band b, and a real number parameterized by
-- the 'BoolParam' parameter: the magnitude (PFalse) or the phase
-- (PTrue) of the complex value at the corresponding place in the matrix A.
build_circuit
getNodeValuesMoreOutputs ::
    Int -> Int -> Int -> Int -> [(Double,Double)] -> Double -> Double -> Double -> BoolParam
    -> Int -> (Int, Double)
getNodeValuesMoreOutputs v' b' nx ny scatteringnodes lx ly k argflag maxConnectivity =
   let maxC = getIntFromParam maxConnectivity in
   let b = getIntFromParam b' in
   let nedges = (nx - 1)*ny + nx*(ny - 1) in
   let flag = v' <= nedges in
   let v = (if flag then v' else (v' - nedges)) in
   let nodeDefault = (if flag then v + nedges else v) in
   let valueDefault = 0.0 in
   let indicesDefault = (-1, -1) in
   let isvalid = checkedge v scatteringnodes nx ny
   in
   if ( (not isvalid) && b == 5 ) then

     let indices = (if flag then (v, v + nedges) else (v + nedges, v)) in
     case argflag of
        PTrue  -> (nodeDefault, valueDefault)
        PFalse -> (nodeDefault, 1.0)

   else if ( (not isvalid) && b /= 5) then (nodeDefault, valueDefault)

   else if ((b > maxC + 2) || (b <= 0)) then (nodeDefault, valueDefault)

   else

   let x = convertband v b' nx ny in

   if (x == -1) then (nodeDefault, valueDefault)
   else

   let isvalid = checkedge x scatteringnodes nx ny in

   if isvalid then

     let ax = calcmatrixelement v x nx ny lx ly k in
     let (node, indices) = if flag then (x + nedges, (v, x + nedges))
                           else (x, (v+nedges,x)) in
     let value = case argflag of
                   PTrue -> atan2  (imagPart ax) (realPart ax)
                   PFalse -> magnitude ax
     in (node, value)

   else
     (nodeDefault, valueDefault)

-- | Auxiliary function for oracle /b/. The inputs are:
--
-- * /y/ :: 'Int' - Global edge index.  Note this is the unmarked /y/
-- coordinate, i.e. the coordinate without scattering regions removed;
--
-- * /nx/ :: 'Int' - Number of vertices left to right;
--
-- * /ny/ :: 'Int' - Number of vertices top to bottom;
--
-- * /lx/ :: 'Double' - Length of horizontal edges (distance between vertices in /x/ direction);
--
-- * /ly/ :: 'Double' - Length of vertical edges (distance between vertices in /y/ direction);
--
-- * /k/ :: 'Double' - Plane wave wavenumber;
--
-- * θ :: 'Double' - Direction of wave propagation;
--
-- * /E0/ :: 'Double' - Magnitude of incident plane wave.
--
-- The output is the magnitude of the electric field on edge /y/.
build_circuit
calcincidentfield :: --
	-- Inputs
		Int -- int y - Global edge index.  Note this is the unmarked y coordinate,
			-- i.e. the coordinate without scattering regions removed.
	-> Int -- int nx - Number of vertices left to right
	-> Int --  int ny - Number of vertices top to bottom
	-> Double -- float lx - Length of horizontal edges (distance between vertices in x direction)
	-> Double -- float ly - Length of vertical edges (distance between vertices in y direction)
	-> Double -- float k - Plane wave wavenumber
	-> Double -- float theta - Direction of wave propagation
	-> Double -- float E0 - Magnitude of incident plane wave
	-- Outputs
	-> Complex Double -- complex float e - Magnitude of electric field on edge y
calcincidentfield y nx ny lx ly k theta e0 =
        let (xg, yg) = itoxy y nx ny in
        --Determine whether edge is horizontal or vertical
        let isvertical = xg >= nx in
        let (xvalueTmp, yvalueTmp) = edgetoxy y nx ny in
        let xvalue = xvalueTmp * lx in
        let yvalue = yvalueTmp * ly in
        -- Convert x and y edge coordinates to x and y values and caluculate field
        if isvertical then
          mkPolar (-cos(theta)*e0) ( -k*(xvalue*cos(theta)+yvalue*sin(theta)))
        else
          mkPolar (sin(theta)*e0) ( -k*(xvalue*cos(theta)+yvalue*sin(theta)))

-- | Auxiliary function for oracle /b/.
build_circuit
getconnection :: Int -> Int -> Int -> Int -> Int -> Int
getconnection y i' nx ny maxConnectivity =
  let i = getIntFromParam i' in
  let maxC = getIntFromParam maxConnectivity in
  let (ex,ey) = itoxy y nx ny in
  let x = if ( (ex < nx) && (ey /= 1) ) then
            case i' of
                1 -> y-2*nx+1
                2 -> y-nx
                3 -> y-nx+1
                4 -> y
                5 -> y+nx-1
                6 -> y+nx
                7 -> y+2*nx-1
                _ -> -1
          else if ( (ex < nx) && (ey == 1) ) then
            case i' of
                1 -> y
                2 -> y+nx-1
                3 -> y+nx
                4 -> y+2*nx-1
                _ -> -1
          else if ( (ex >= nx) && (ex /= nx) && (ex /= 2*nx-1) ) then
             case i' of
                1 -> y-nx
                2 -> y-nx+1
                3 -> y-1
                4 -> y
                5 -> y+1
                6 -> y+nx-1
                7 -> y+nx
                _ -> -1
          else if ( (ex >= nx) && (ex == nx) ) then
             case i' of
                1 -> y-nx+1
                2 -> y
                3 -> y+1
                4 -> y+nx
                _ -> -1
          else if ( (ex >= nx) && (ex == 2*nx-1) ) then
             case i' of
                1 -> y-nx
                2 -> y-1
                3 -> y
                4 -> y+nx-1
                _ -> -1
          else -1
  in
  if (i > maxC) then -1
  else if (x > nx*(ny-1)+ny*(nx-1)) then -1
  else x

-- | Auxiliary function to @'template_paramZero'@.
build_circuit
local_loop_with_index_aux :: Int -> Int -> t -> (Int -> t -> t) -> t
local_loop_with_index_aux i n x f =
   case paramMinus n i of
     0 -> x
     _ -> local_loop_with_index_aux (paramSucc i) n (f i x) f

-- | Local version of @'loop_with_index'@, for lifting.
build_circuit
local_loop_with_index :: Int -> t -> (Int -> t -> t) -> t
local_loop_with_index n x f = local_loop_with_index_aux paramZero n x f

-- | Oracle /b/.
build_circuit
getKnownWeights :: Int -> Int -> Int -> [(Double,Double)] -> Double -> Double -> Double -> Double -> Double -> Int -> Complex Double
getKnownWeights y nx ny scatteringnodes lx ly k theta e0 maxConnectivity =
   let makeConnections i connections = let x = getconnection y (paramSucc i) nx ny maxConnectivity in
                                       let t = not $ checkedge x scatteringnodes nx ny in
                                       (x,t):connections
   in
   let calcTang b (c,t) = if t
                          then let matElt = calcmatrixelement y c nx ny lx ly k in
                               let incField = calcincidentfield c nx ny lx ly k theta e0
                               in b - matElt * incField
                          else b
   in
   let connections = local_loop_with_index maxConnectivity [] makeConnections in
   if (not $ checkedge y scatteringnodes nx ny) then 0.0 :+ 0.0
   else foldl calcTang (0.0 :+ 0.0) connections

----------------------------------------------------------------------
----------------------------------------------------------------------
-- Testing functions

test_template_sinc = do
      f <- template_sinc
      r <- qinit (0 :: FDouble)
      f r
      return ()

test_template_itoxy = do
      f <- template_itoxy
      x <- qinit (0 :: FSignedInt)
      y <- qinit (0 :: FSignedInt)
      z <- qinit (0 :: FSignedInt)
      g <- f x
      h <- g y
      k <- h z
      return ()

test_template_edgetoxy = do
      f <- template_edgetoxy
      x <- qinit (0 :: FSignedInt)
      y <- qinit (0 :: FSignedInt)
      z <- qinit (0 :: FSignedInt)
      g <- f x
      h <- g y
      k <- h z
      return ()

test_template_calcRweights = do
      f <- template_calcRweights
      n1 <- qinit (0 :: FSignedInt)
      n2 <- qinit (0 :: FSignedInt)
      n3 <- qinit (0 :: FSignedInt)
      x1 <- qinit (0 :: FDouble)
      x2 <- qinit (0 :: FDouble)
      x3 <- qinit (0 :: FDouble)
      x4 <- qinit (0 :: FDouble)
      x5 <- qinit (0 :: FDouble)
      f1 <- f n1
      f2 <- f1 n2
      f3 <- f2 n3
      g1 <- f3 x1
      g2 <- g1 x2
      g3 <- g2 x3
      g4 <- g3 x4
      g5 <- g4 x5
      return ()

test_template_calcincidentfield = do
   y' <- qinit (0 :: FSignedInt)
   nx' <- qinit (0 :: FSignedInt)
   ny' <- qinit (0 :: FSignedInt)
   lx' <- qinit (0 :: FDouble)
   ly' <- qinit (0 :: FDouble)
   k' <- qinit (0 :: FDouble)
   theta' <- qinit (0 :: FDouble)
   e0' <- qinit (0 :: FDouble)
   f <- template_calcincidentfield
   f1 <- f y'
   f2 <- f1 nx'
   f3 <- f2 ny'
   f4 <- f3 lx'
   f5 <- f4 ly'
   f6 <- f5 k'
   f7 <- f6 theta'
   f8 <- f7 e0'
   return ()

test_template_calcmatrixelement = do
   y1 <- qinit (0 :: FSignedInt)
   y2 <- qinit (0 :: FSignedInt)
   nx <- qinit (0 :: FSignedInt)
   ny <- qinit (0 :: FSignedInt)
   lx <- qinit (0 :: FDouble)
   ly <- qinit (0 :: FDouble)
   k  <- qinit (0 :: FDouble)
   f <- template_calcmatrixelement
   f1 <- f y1
   f2 <- f1 y2
   f3 <- f2 nx
   f4 <- f3 ny
   f5 <- f4 lx
   f6 <- f5 ly
   f7 <- f6 k
   return ()

test_template_getconnection = do
   y <- qinit (0 :: FSignedInt)
   nx <- qinit (0 :: FSignedInt)
   ny <- qinit (0 :: FSignedInt)
   f <- template_getconnection
   f1 <- f y
   f2 <- f1 6
   f3 <- f2 nx
   f4 <- f3 ny
   f5 <- f4 7
   return ()

test_template_checkedge = do
   y <- qinit (0 :: FSignedInt)
   nx <- qinit (0 :: FSignedInt)
   ny <- qinit (0 :: FSignedInt)
   s <- qinit [(0 :: FDouble,0 :: FDouble),(0 :: FDouble,0 :: FDouble)]
   f <- template_checkedge
   f1 <- f y
   f2 <- f1 s
   f3 <- f2 nx
   f4 <- f3 ny
   return ()


-- An abstract data type for performing multiple linear regression,
-- using HMatrix as the backend for matrix computations.
--
-- types:
--   Matrix a (2D-array)
--   Vector a (vector from the standard vector package)
--
-- Both types are dense, immutable, and strict in all elements and are
-- manipulated as whole blocks.
--
-- matrix product is a <> b, where a and b are matrices
-- matrix-vector product is a #> x, where a is a matrix and x a vector
-- vector-vector dot product is x <.> y
-- the (conjugate) transpose is trans a, where a is a matrix or vector
-- the general linear system solver is <\>
--
-- to construct a 4 by 3 matrix:
-- let a = (4><3) [1,2,3,4,5,6,7,8,9,10,11,12] :: Matrix R
--
-- to construct a vector of double-precision floating points:
-- let x = vector [1,2,3]
--
-- to multiply vectors:
-- x * y
--
-- to add two vectors elementwise:
-- x + y
--
-- to compute vector dot product:
-- x <.> y
--
-- to perform matrix-vector product:
-- m  #> x
--
-- to perform matrix-matrix product
-- m <> n
--
-- to sum elements of a vector or matrix:
-- sumelements m
--

module Regression.HMatrix (
    model_features,
    model_outputs,
    model_weights,
    model_predictions,
    model_rss,
    create_model,
    create_features,
    create_weights,
    predict,
    powers,
    normalize,
    newtons_method,
    newtons_method_with_norm,
    cross_validate,
    gradient_descent
) where

import Numeric.LinearAlgebra hiding (magnitude)
import Numeric.LinearAlgebra.HMatrix hiding (magnitude)
import Data.List hiding (transpose)
import Data.Maybe (fromJust)
import Debug.Trace (trace)
import Stat (range, mean, stdev)
import qualified Data.Vector as V

-- A model is:
-- * a feature matrix with n rows and d columns
-- * a vector of length n representing the observed output
-- * a weight vector containing the calculated vector of optimized weights
-- * the predicted outputs using the optimized weights
-- * the
data Model = MO {
    model_features    :: FeatureMatrix,
    model_outputs     :: FeatureVector,
    model_weights     :: WeightVector,
    model_predictions :: FeatureVector,
    model_rss         :: Double
} deriving (Show)

-- A two-dimensional matrix whose columns contain feature values.
data FeatureMatrix = FM {
    fm_name_indexes:: [(String, Int)],
    fm_values :: Matrix Double
} deriving (Show)

-- A vector of values for a named feature.
data FeatureVector = FV {
    fv_name :: String,
    fv_values :: Vector Double
}
instance Show FeatureVector where
    show (FV name values) =
        name ++ " = " ++ show (toList values)

type Feature a = (String, a -> Double)

type Output a = (String, a -> Double)

type Optimizer = FeatureMatrix -> FeatureVector -> WeightVector

-- A vector of weights.  Each weight corresponds to a feature in the
-- feature matrix used to calculate the weights.
data WeightVector = WV {
    wv_name_indexes :: [(String, Int)],
    wv_values :: Vector Double
}

instance Show WeightVector where
    show (WV name_indexes values) =
        concatMap showWeight (zip name_indexes (toList values))
        where showWeight ((name, i), value) = name ++ " = " ++ show value ++ "\n"

-- Creates a model from a list of records, a list of features, and the output
-- accessor function.  This computes the weight vector as well.  This function
-- does not scale the features nor does it add an intercept feature (a feature
-- whose values are all 1).
create_model :: [a] -> [Feature a] -> Output a -> Optimizer -> Model
create_model rows features (output_name, output) optimizer =
    let fmat         = create_features features rows
        nn           = length rows
        observations = FV output_name (fromList (Data.List.map output rows))
        weights      = optimizer fmat observations
        predictions  = predict fmat weights
        residuals    = rss observations predictions
    in MO fmat observations weights predictions residuals

-- Multiplies a feature matrix by a weights vector to obtain a prediction of
-- output.
predict :: FeatureMatrix -> WeightVector -> FeatureVector
predict (FM _ h) (WV _ w) = FV "" (h #> w)

-- Generates the feature matrix, usually denoted as H, with N rows and
-- D features where D is the length of the feature list.
create_features :: [Feature a] -> [a] -> FeatureMatrix
create_features hs inputs =
    let n = length inputs
        d = length hs
        names = (Data.List.map fst hs)
        name_indexes = Prelude.zip names [0..]
        dat = [h(row) | row <- inputs, (_,h) <- hs]
        h = (n><d) dat
    in FM name_indexes h

-- Computes the residual sum of squares from the observed output and
-- the predicted output.
rss :: (Num a) => FeatureVector -> FeatureVector -> Double
rss (FV _ v1) (FV _ v2) = rss' v1 v2

rss' :: Vector Double -> Vector Double -> Double
rss' v1 v2 = diff <.> diff
    where diff = v1 - v2

-- Takes a list of numbers and turns them into a vector of weights.
create_weights :: FeatureMatrix -> [Double] -> WeightVector
create_weights (FM name_indexes h) weights = WV name_indexes (vector weights)

newtons_method :: Optimizer
newtons_method h y = newtons_method_with_norm 0 h y

newtons_method_with_norm :: Matrix Double -> Optimizer
newtons_method_with_norm l (FM n h) (FV o y) = WV n w
    where w = newtons_method_with_norm' l h y

-- Performs newtons method, finding an estimate of the minimum.  This method
-- uses the formula:
-- w_hat = (H^t * H - lambda*I)^-1 * H^T * y
newtons_method_with_norm' :: Matrix Double -> Matrix Double -> Vector Double -> Vector Double
newtons_method_with_norm' l h y = ((inv (hth - mod)) <> th) #> y
    where th = tr h
          hth = th <> h
          (d,_) = size hth
          mod = l * (ident d) :: Matrix Double

-- Performs gradient descent, updating the weights until the residual
-- sum of squares is less than the epsilon value e, at which point the
-- weight matrix is returned.  n is the step size.
gradient_descent :: Double -> Double -> WeightVector -> Optimizer
gradient_descent e n (WV name w) (FM _ f) (FV _ o) =
    let ft      = tr f
        weights = gradient_descent' e n w f ft o
    in WV name weights

gradient_descent' :: Double -> Double -> Vector Double -> Matrix Double ->
    Matrix Double -> Vector Double -> Vector Double
gradient_descent' e n w h ht y =
    let grad = gradient h ht y w -- (-2H^t(y-Hw))
        grad_len = magnitude grad  -- grad RSS(w) == ||2H^t(y-HW)||
    --in if grad_len < e
    in if (trace ("gradient = " ++ show grad_len) grad_len) < e
        then w
        else let delta = cmap (*(-n)) grad -- (2nH^t(y-Hw))
                 w' = w + delta
             in gradient_descent' e n w' h ht y

-- Compute the magnitude of the given vector.
magnitude :: Vector Double -> Double
magnitude vec = sumElements $ cmap (\y -> y^2) vec

-- Calculates the gradient of the residual sum of squares (-2H^t(y-Hw)).
-- This is used to compute the magnitude of the gradient, to see if the
-- function is minimized.  It is also used to update the weights of the
-- features.
gradient :: Matrix Double -> Matrix Double -> Vector Double ->
    Vector Double -> Vector Double
gradient h ht y w =
    let yhat = h #> w
        err = y - yhat
        prod = ht #> err
        grad = cmap (*(-2)) prod
    in grad

-- Given a feature, computes polynomial powers of that feature from
-- 1 (the original feature) up to and including n.
powers :: Feature a -> Int -> [Feature a]
powers (name, f) n = Data.List.map  (\i -> (name ++ (show i), (\a -> (f a)^i))) [1..n]

-- given a list of values and a feature, return a normalized version of that
-- feature where the mean of the feature is subtracted from each feature and
-- the result is divided by the standard deviation of the feature.
normalize :: [a] -> (a -> Double) -> a -> Double
normalize xs f = (/sdxs) . (+(-meanxs)) . f
    where meanxs = mean (Data.List.map f xs)
          sdxs = stdev (Data.List.map f xs)

-- Compute the average k-fold cross-validation error for a given L2 penalty and
-- a given split size k.
-- Note that the rows should be shuffled before they are provided.
cross_validate :: Int -> [a] -> [Feature a] -> Output a  -> Optimizer -> Double
cross_validate k rows features o@(output_name, output) optimizer = kmean
    where ksplits         = ksplit k rows
          training_data   = map fst ksplits
          valid_data      = map snd ksplits
          trained_models  = map (\rows -> create_model rows features o optimizer) training_data
          trained_weights = map model_weights trained_models
          valid_features  = map (create_features features) valid_data
          predictions     = zipWith (\h w -> predict h w) valid_features trained_weights
          observations    = map (\dat -> FV output_name (fromList (Data.List.map output dat))) valid_data
          valid_rss       = zipWith (\o p -> rss o p) observations predictions
          kmean           = mean valid_rss

-- prepare splits of the data for k-means where k is the number of
-- clusters and the number of values in the test set (the second list
-- in the pair) is n/k where n is the size of the input list.
ksplit :: Int -> [a] -> [([a], [a])]
ksplit k xs = map (ksplit' k xs) [0..k-1]

-- prepare a k split ([a], [a]) for the chunk at index i with list xs
-- and with number of chunks k
ksplit' :: Int -> [a] -> Int -> ([a], [a])
ksplit' k xs i = ((take (sz*i) xs) ++ (drop (sz*(i+1)) xs), take sz (drop (sz*i) xs))
    where sz = length xs `div` k