{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}

module Backprop.Learn.Model.Stochastic (
    dropout
  , rreLU
  , injectNoise, applyNoise
  , injectNoiseR, applyNoiseR
  ) where

import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Types
import           Control.Monad.Primitive
import           Data.Bool
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.LinearAlgebra.Static.Vector
import qualified Data.Vector.Storable.Sized            as SVS
import qualified Statistics.Distribution               as Stat
import qualified System.Random.MWC                     as MWC
import qualified System.Random.MWC.Distributions       as MWC

-- | Dropout layer.  Parameterized by dropout percentage (should be between
-- 0 and 1).
--
-- 0 corresponds to no dropout, 1 corresponds to complete dropout of all
-- nodes every time.
dropout
    :: KnownNat n
    => Double
    -> Model 'Nothing 'Nothing (R n) (R n)
dropout r = Func
    { runFunc      = (auto (realToFrac (1 - r)) *)
    , runFuncStoch = \g x -> do
        (x *) . auto . vecR <$> SVS.replicateM (mask g)
    }
  where
    mask :: PrimMonad m => MWC.Gen (PrimState m) -> m Double
    mask = fmap (bool 1 0) . MWC.bernoulli r

-- | Random leaky rectified linear unit
rreLU
    :: (Stat.ContGen d, Stat.Mean d, KnownNat n)
    => d
    -> Model 'Nothing 'Nothing (R n) (R n)
rreLU d = Func
    { runFunc      = vmap' (preLU v)
    , runFuncStoch = \g x -> do
        α <- vecR <$> SVS.replicateM (Stat.genContVar d g)
        pure (zipWithVector preLU (constVar α) x)
    }
  where
    v :: BVar s Double
    v = auto (Stat.mean d)

-- | Inject random noise.  Usually used between neural network layers, or
-- at the very beginning to pre-process input.
--
-- In non-stochastic mode, this adds the mean of the distribution.
injectNoise
    :: (Stat.ContGen d, Stat.Mean d, Fractional a)
    => d
    -> Model 'Nothing 'Nothing a a
injectNoise d = Func
    { runFunc = (realToFrac (Stat.mean d) +)
    , runFuncStoch = \g x -> do
        e <- Stat.genContVar d g
        pure (realToFrac e + x)
    }

-- | 'injectNoise' lifted to 'R'
injectNoiseR
    :: (Stat.ContGen d, Stat.Mean d, KnownNat n)
    => d
    -> Model 'Nothing 'Nothing (R n) (R n)
injectNoiseR d = Func
    { runFunc = (realToFrac (Stat.mean d) +)
    , runFuncStoch  = \g x -> do
        e <- vecR <$> SVS.replicateM (Stat.genContVar d g)
        pure (constVar e + x)
    }

-- | Multply by random noise.  Can be used to implement dropout-like
-- behavior.
--
-- In non-stochastic mode, this scales by the mean of the distribution.
applyNoise
    :: (Stat.ContGen d, Stat.Mean d, Fractional a)
    => d
    -> Model 'Nothing 'Nothing a a
applyNoise d = Func
    { runFunc = (realToFrac (Stat.mean d) *)
    , runFuncStoch  = \g x -> do
        e <- Stat.genContVar d g
        pure (realToFrac e * x)
    }

-- | 'applyNoise' lifted to 'R'
applyNoiseR
    :: (Stat.ContGen d, Stat.Mean d, KnownNat n)
    => d
    -> Model 'Nothing 'Nothing (R n) (R n)
applyNoiseR d = Func
    { runFunc = (realToFrac (Stat.mean d) *)
    , runFuncStoch  = \g x -> do
        e <- vecR <$> SVS.replicateM (Stat.genContVar d g)
        pure (constVar e * x)
    }


{-# LANGUAGE OverloadedStrings #-}

module Statistics.Classification.ROC (
    ROC(..)
  , ROCPoint(..)
  , rocPoint
  , efficientROC
  , roc
  ) where

import Data.List (foldl')

import Data.Csv ((.:), (.=), DefaultOrdered(..), FromNamedRecord(..), ToNamedRecord(..), header, namedRecord)
import qualified Data.Vector as V

import Statistics.Classification.ConfusionMatrix (ConfusionMatrix, trueConfMatrix, truePosRate, falsePosRate)
import Statistics.Classification.Types (ClassificationResult(..), ClassificationScore(..), sortByScore)

data ROC = ROC {
    rocAUC :: !Double
  , rocCurve :: !(V.Vector ROCPoint)
  } deriving (Eq, Show, Read)

data ROCPoint = ROCPoint {
    rocTruePos :: {-# UNPACK #-}!Double
  , rocFalsePos :: {-# UNPACK #-}!Double
  , rocCutoff :: {-# UNPACK #-}!Double
  , rocConfusion :: !ConfusionMatrix
  } deriving (Eq, Show, Read)

rocPoint :: Double -> ConfusionMatrix -> ROCPoint
rocPoint cutoff cm = ROCPoint {
    rocTruePos = truePosRate cm
  , rocFalsePos = falsePosRate cm
  , rocCutoff = cutoff
  , rocConfusion = cm
  }

instance DefaultOrdered ROCPoint where
  headerOrder _ = header [
      "true_positive_rate"
    , "false_positive_rate"
    , "cutoff"
    ] <> headerOrder (undefined :: ConfusionMatrix)

instance ToNamedRecord ROCPoint where
  toNamedRecord rp = namedRecord [
      "true_positive_rate" .= rocTruePos rp
    , "false_positive_rate" .= rocFalsePos rp
    , "cutoff" .= rocCutoff rp
    ] <> toNamedRecord (rocConfusion rp)

instance FromNamedRecord ROCPoint where
  parseNamedRecord r = ROCPoint <$>
        r .: "true_positive_rate"
    <*> r .: "false_positive_rate"
    <*> r .: "cutoff"
    <*> parseNamedRecord r

data EffRocState = EffRocState {
    effRocCutoff :: {-# UNPACK #-}!Double
  , effRocTP :: {-# UNPACK #-}!Int
  , effRocFP :: {-# UNPACK #-}!Int
  , effRocTPprev :: {-# UNPACK #-}!Int
  , effRocFPprev :: {-# UNPACK #-}!Int
  , effRocOut :: ![ROCPoint]
  , effRocAUC :: {-# UNPACK #-}!Double
  } deriving (Eq, Show, Read)

emptyEffRocState :: EffRocState
emptyEffRocState = EffRocState inf 0 0 0 0 [] 0
  where inf = 1 / 0

trapezoidArea :: Int -> Int -> Int -> Int -> Double
trapezoidArea x1 x2 y1 y2 = base * height
  where
    base = fromIntegral . abs $ x2 - x1
    height = 0.5 * fromIntegral y2 + 0.5 * fromIntegral y1

addArea :: EffRocState -> EffRocState
addArea s0 = s0 { effRocAUC = effRocAUC s0 + area }
  where
    area = trapezoidArea (effRocFP s0) (effRocFPprev s0) (effRocTP s0) (effRocTPprev s0)

updRocState :: Int -> Int -> EffRocState -> ClassificationScore Bool -> EffRocState
updRocState p n s r
  | effRocCutoff s /= cutoff = nextExample . nextCutoff $ addArea s
  | otherwise = nextExample s
  where
    cutoff = classifiedScore r
    actual = classifiedExample r
    nextCutoff s0 = s0 { effRocCutoff = cutoff, effRocTPprev = effRocTP s0, effRocFPprev = effRocFP s0 }
    nextExample s0 =
      let s1 = if actual then s0 { effRocTP = effRocTP s0 + 1 } else s0 { effRocFP = effRocFP s0 + 1 }
          cm = trueConfMatrix (ClassificationResult p (effRocTP s1)) (ClassificationResult n (n - effRocFP s1))
      in s1 { effRocOut = rocPoint cutoff cm : effRocOut s1 }

-- O(n log n) ROC curve generation and AUC calculation
-- Algorithm 1 and 2 in
-- Fawcett, T., 2006. An Introduction to ROC Analysis. Pattern Recognition Letters 27 (861 - 874)
efficientROC :: Foldable f => Int -> Int -> f (ClassificationScore Bool) -> ROC
efficientROC p n = out . foldl' (updRocState p n) emptyEffRocState
  where
    finAUC effr
      | p == 0 || n == 0 = 0 -- degenerate case, no actual positive or negative examples.
      | otherwise = (effRocAUC effr + trapezoidArea n (effRocFPprev effr) p (effRocTPprev effr)) / fromIntegral (p * n)
    out s = ROC {
        rocAUC = finAUC s
      , rocCurve = V.fromList $ effRocOut s
      }

roc :: V.Vector (ClassificationScore Bool) -> ROC
roc v = efficientROC p n . V.reverse $ sortByScore v
  where
    p = V.length $ V.findIndices classifiedExample v
    n = V.length v - p


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Geometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Geometric distribution. There are two variants of
-- distribution. First is the probability distribution of the number
-- of Bernoulli trials needed to get one success, supported on the set
-- [1,2..] ('GeometricDistribution'). Sometimes it's referred to as
-- the /shifted/ geometric distribution to distinguish from another
-- one.
--
-- Second variant is probability distribution of the number of
-- failures before first success, defined over the set [0,1..]
-- ('GeometricDistribution0').
module Statistics.Distribution.Geometric
    (
      GeometricDistribution
    , GeometricDistribution0
    -- * Constructors
    , geometric
    , geometric0
    -- ** Accessors
    , gdSuccess
    , gdSuccess0
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Applicative ((<$>))
import Control.Monad (liftM)
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_neg_inf)
import qualified Statistics.Distribution as D
import qualified System.Random.MWC.Distributions as MWC

----------------------------------------------------------------
-- Distribution over [1..]

newtype GeometricDistribution = GD {
      gdSuccess :: Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON GeometricDistribution
instance ToJSON GeometricDistribution

instance Binary GeometricDistribution where
    get = GD <$> get
    put (GD x) = put x

instance D.Distribution GeometricDistribution where
    cumulative = cumulative

instance D.DiscreteDistr GeometricDistribution where
    probability (GD s) n
      | n < 1     = 0
      | otherwise = s * (1-s) ** (fromIntegral n - 1)
    logProbability (GD s) n
       | n < 1     = m_neg_inf
       | otherwise = log s + log (1-s) * (fromIntegral n - 1)

instance D.Mean GeometricDistribution where
    mean (GD s) = 1 / s

instance D.Variance GeometricDistribution where
    variance (GD s) = (1 - s) / (s * s)

instance D.MaybeMean GeometricDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance GeometricDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy GeometricDistribution where
  entropy (GD s)
    | s == 0 = m_pos_inf
    | s == 1 = 0
    | otherwise = negate $ (s * log s + (1-s) * log (1-s)) / s

instance D.MaybeEntropy GeometricDistribution where
  maybeEntropy = Just . D.entropy

instance D.DiscreteGen GeometricDistribution where
  genDiscreteVar (GD s) g = MWC.geometric1 s g

instance D.ContGen GeometricDistribution where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

-- | Create geometric distribution.
geometric :: Double                -- ^ Success rate
          -> GeometricDistribution
geometric x
  | x >= 0 && x <= 1 = GD x
  | otherwise        =
    error $ "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x

cumulative :: GeometricDistribution -> Double -> Double
cumulative (GD s) x
  | x < 1        = 0
  | isInfinite x = 1
  | isNaN      x = error "Statistics.Distribution.Geometric.cumulative: NaN input"
  | otherwise    = 1 - (1-s) ^ (floor x :: Int)

----------------------------------------------------------------
-- Distribution over [0..]

newtype GeometricDistribution0 = GD0 {
      gdSuccess0 :: Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON GeometricDistribution0
instance ToJSON GeometricDistribution0

instance Binary GeometricDistribution0 where
    get = GD0 <$> get
    put (GD0 x) = put x

instance D.Distribution GeometricDistribution0 where
    cumulative (GD0 s) x = cumulative (GD s) (x + 1)

instance D.DiscreteDistr GeometricDistribution0 where
    probability    (GD0 s) n = D.probability    (GD s) (n + 1)
    logProbability (GD0 s) n = D.logProbability (GD s) (n + 1)

instance D.Mean GeometricDistribution0 where
    mean (GD0 s) = 1 / s - 1

instance D.Variance GeometricDistribution0 where
    variance (GD0 s) = D.variance (GD s)

instance D.MaybeMean GeometricDistribution0 where
    maybeMean = Just . D.mean

instance D.MaybeVariance GeometricDistribution0 where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy GeometricDistribution0 where
  entropy (GD0 s) = D.entropy (GD s)

instance D.MaybeEntropy GeometricDistribution0 where
  maybeEntropy = Just . D.entropy

instance D.DiscreteGen GeometricDistribution0 where
  genDiscreteVar (GD0 s) g = MWC.geometric0 s g

instance D.ContGen GeometricDistribution0 where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

-- | Create geometric distribution.
geometric0 :: Double                -- ^ Success rate
           -> GeometricDistribution0
geometric0 x
  | x >= 0 && x <= 1 = GD0 x
  | otherwise        =
    error $ "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}

module GrenadeExtras.GradNorm (
  GradNorm(..)
) where

import Grenade (FullyConnected'(FullyConnected'), FullyConnected(FullyConnected), Tanh,
                Logit, Gradient, Relu, Gradients(GNil, (:/>)))
import Numeric.LinearAlgebra.Static ((<.>), toColumns)
import GHC.TypeLits (KnownNat)

class GradNorm x where
  normSquared :: x -> Double

instance (KnownNat i, KnownNat o) => GradNorm (FullyConnected' i o) where
  normSquared (FullyConnected' r l) = r <.> r + (sum . fmap (\c->c<.>c) . toColumns $ l)

instance (KnownNat i, KnownNat o) => GradNorm (FullyConnected i o) where
  normSquared (FullyConnected w _) = normSquared w

instance GradNorm () where
  normSquared _ = 0

instance (GradNorm a, GradNorm b) => GradNorm (a, b) where
  normSquared (a, b) = normSquared a + normSquared b

instance GradNorm Logit where
  normSquared _ = 0

instance GradNorm Tanh where
  normSquared _ = 0

instance GradNorm Relu where
  normSquared _ = 0

instance GradNorm (Gradients '[]) where
  normSquared GNil = 0

instance (GradNorm l, GradNorm (Gradient l), GradNorm (Gradients ls)) => GradNorm (Gradients (l ': ls)) where
  normSquared (grad :/> grest) = normSquared grad + normSquared grest


module Q.Options.Black76
  (
    Black76(..)
  , atmf
  , euOption
  , eucall
  , euput
  , dPlus
  , dMinus
  )
  where

import           Numeric.GSL (derivCentral)
import           Q.Options (DF, Delta (Delta), Forward (..), Gamma (Gamma), OptionType (..),
                            Premium (Premium), Strike (..), TimeScaleable (scale), TotalVar (..),
                            Valuation (Valuation, vPremium), Vega (Vega), Vol (..),
                            YearFrac (YearFrac), discount)
import           Q.Options.ImpliedVol.TimeSlice (TimeSlice (..))
import           Statistics.Distribution (cumulative, density)
import           Statistics.Distribution.Normal (standard)
import Q.Types (LogRelStrike)

data Black76 = Black76 {
    b76F   :: Forward
  , b76DF  :: DF
  , b76T   :: YearFrac
  , b76Vol :: Vol
}

-- | At the money forward strike.
atmf :: Black76 -> Strike
atmf Black76{..} = Strike f
  where (Forward f) = b76F

-- | European option valuation with black 76
euOption :: Black76 -> OptionType -> Strike -> Valuation
euOption Black76{..} cp k = Valuation premium delta vega gamma where
  (Forward f) = b76F
  n           = cumulative standard
  (Vol sigmaSqt) = scale b76T b76Vol
  d1          = dPlus  b76F b76Vol k b76T
  d2          = dMinus b76F b76Vol k b76T
  nd1         = n d1
  nd2         = n d2
  callDelta   = b76DF `discount` nd1
  putDelta    = b76DF `discount` (- (n (-d1)))
  vega        = Vega  $ b76DF `discount` density standard d1 * f * sigmaSqt
  gamma       = Gamma $ b76DF `discount` density standard d1 / (f * sigmaSqt)
  premium  = Premium $ case cp of
    Call -> b76DF `discount` (f * nd1 - nd2 * k')
    Put  -> b76DF `discount` (n (-d2) * k' - n (-d1) * f)
    where (Strike k') = k
  delta | cp == Call = Delta callDelta
        | otherwise = Delta putDelta

-- | see 'euOption'
euput :: Black76 -> Strike -> Valuation
euput b76 =  euOption b76 Put

-- | see 'euOption'
eucall :: Black76 -> Strike -> Valuation
eucall b76 = euOption b76 Call

dPlus :: Forward -> Vol -> Strike -> YearFrac -> Double
dPlus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) + 0.5 * sigma * sigma * t)

dMinus :: Forward -> Vol -> Strike -> YearFrac -> Double
dMinus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) - 0.5 * sigma * sigma * t)

instance TimeSlice Black76 Strike where
  totalVar Black76{..} _ = TotalVar $ vol * vol * t
    where (Vol vol) = b76Vol
          (YearFrac t) = b76T
  dW _ _  = 0
  d2W _ _ = 0

  impliedDensity b76 (Strike k) = let
    dk k' = fst (derivCentral 1e-4 (\k2 -> let (Premium v) = vPremium (eucall b76 (Strike k2)) in v) k')
    in fst (derivCentral 1e-4 dk k)

instance TimeSlice Black76 LogRelStrike  where
  totalVar Black76{..} _ = TotalVar $ vol * vol * t
    where (Vol vol) = b76Vol
          (YearFrac t) = b76T
  dW _ _  = 0
  d2W _ _ = 0

  impliedDensity bs k = impliedDensity f k where
    f :: LogRelStrike  -> TotalVar
    f = totalVar bs


{-# LANGUAGE NamedFieldPuns, TemplateHaskell #-}

module School.Unit.Test.MultiNoulli
( multiNoulliTest ) where

import Conduit ((.|), await, liftIO, yield)
import Data.Default.Class (def)
import Numeric.LinearAlgebra ((><))
import School.TestUtils (addClasses, assertRight, diffCost, randomMatrix,
                         randomNNInts, singleClassOutput, testState, unitCorrect)
import School.Train.AppTrain (AppTrain)
import School.Train.ForwardPass (forwardPass)
import School.Types.FloatEq ((~=))
import School.Types.Slinky (Slinky(..), slinkySingleton)
import School.Unit.CostFunction (CostFunction(..))
import School.Unit.CostParams (CostParams(..))
import School.Unit.LogSoftMax (logSoftMax)
import School.Unit.MultiNoulli
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Unit.UnitParams (UnitParams(..))
import School.Utils.LinearAlgebra (compareDoubleMatrix)
import School.Utils.Tuple (trd3)
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck hiding ((><))
import Test.Tasty.TH
import Test.QuickCheck.Monadic (monadicIO, assert, pre)

eps :: Double
eps = 1e-5

prec :: Double
prec = 5e-2

multi :: CostFunction Double (AppTrain a)
multi = multiNoulli Nothing Nothing

prop_correct_classes :: (Positive Int) -> (Positive Int) -> Property
prop_correct_classes (Positive c) (Positive b) = monadicIO $ do
  classes <- liftIO $ randomNNInts (c - 1) b
  let out = singleClassOutput c classes
  let params = slinkySingleton $ BatchClassTarget classes
  let cost = computeCost multi out params
  assert $ cost ~= Right (-1)

prop_wrong_classes :: (Positive Int) -> (Positive Int) -> Property
prop_wrong_classes (Positive c) (Positive b) = monadicIO $ do
  pre $ c > 1
  classes <- liftIO $ randomNNInts (c - 1) b
  let wrong = map (\k -> mod (k + 1) c) classes
  let out = singleClassOutput c wrong
  let params = slinkySingleton $ BatchClassTarget classes
  let cost = computeCost multi out params
  assert $ cost == Right 0

prop_norm :: (Positive Int) -> (Positive Int) -> Property
prop_norm (Positive c) (Positive b) = monadicIO $ do
  classes <- liftIO $ randomNNInts (c - 1) b
  input <- liftIO $ BatchActivation <$> randomMatrix b c
  let out = apply logSoftMax EmptyParams input
  let params = slinkySingleton $ BatchClassTarget classes
  let cost = computeCost multi out params
  assertRight (>= (-1)) cost

prop_numerical_deriv :: (Positive Int) -> (Positive Int) -> Property
prop_numerical_deriv (Positive c) (Positive b) = monadicIO $ do
  pre $ b < 23 && c < 23
  classes <- liftIO $ slinkySingleton . BatchClassTarget <$> randomNNInts (c - 1) b
  input <- liftIO $ BatchActivation <$> randomMatrix b c
  let deriv = derivCost multi input classes
  let check = (b >< c) [ diffCost multi
                                  input
                                  eps
                                  (j, k)
                                  classes
                               | j <- [0..b-1], k <- [0..c-1] ]
  assertRight (\(BatchGradient g) -> compareDoubleMatrix prec g check)
              deriv

prop_correct_forward_pass :: (Positive Int) -> (Positive Int) -> Property
prop_correct_forward_pass (Positive c) (Positive b) = monadicIO $ do
  classes <- liftIO $ randomNNInts (c - 1) b
  let forward = forwardPass [unitCorrect classes] multi
  activation <- liftIO $ randomMatrix b c
  let input = ([BatchActivation $ addClasses classes activation], SNil)
  let pass = yield input
          .| forward
          .| await
  result <- testState pass def
  assertRight ((maybe False ((~= (-1)) . trd3)) . fst)
              result

multiNoulliTest :: TestTree
multiNoulliTest = $(testGroupGenerator)


module Lib where

import Data.Complex
import qualified Data.Vector as V
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Generic as G
import Codec.Picture

w :: Int
w = 500
h :: Int
h = 500
sz :: Float
sz = fromIntegral st / 2.5
st :: Int
st = 30

colors :: V.Vector PixelRGB8
colors = V.fromList [
    PixelRGB8 22 160 133,
    PixelRGB8 192 57 43,
    PixelRGB8 44 62 80,
    PixelRGB8 142 68 173 ]

roots :: [Complex Float]
roots = [
    10 :+ 6,
    (-2) :+ 4,
    (-3) :+ (-9),
    6 :+ (-3) ]

rootsV :: V.Vector (Complex Float)
rootsV = V.fromList roots

coefs :: [Complex Float]
coefs = coefsFromRoots rootsV
coefs' :: [Complex Float]
coefs' = deriveCoefs coefs

podd :: (Integral a, Num p) => a -> p
podd n
    | odd n = -1
    | otherwise = 1

binom :: Int -> Int -> [S.Vector Int]
binom n k
  | k < 0 = []
  | k > n = []
  | n <= 0 = [G.empty]
  | k == n = [G.enumFromN 1 n]
  | otherwise = binom (n-1) k ++ ((`G.snoc` n) <$> binom (n-1) (k-1))

coefsFromRoots :: (Num a, S.Storable a, G.Vector v a) => v a -> [a]
coefsFromRoots roots =
    let n = G.length roots
     in [ podd (n-k) * sum [ G.product (G.map (\i -> roots G.! (i-1)) p) | p <- binom n (n-k) ] | k <- [0..n] ]

deriveCoefs :: Num b => [b] -> [b]
deriveCoefs coefs = (\(i, c) -> fromIntegral i * c) <$> zip [1..] (tail coefs)

evalPoly :: Num a => [a] -> a -> a
evalPoly [] _ = 0
evalPoly (c0:cs) z = c0 + z * evalPoly cs z

newton :: (RealFloat a, Num t, Eq t) => t -> [Complex a] -> [Complex a] -> Complex a -> Complex a
newton steps coefs coefs' x0 =
    if steps == 0 || magnitude fx < 0.1 then x0
    else newton (steps - 1) coefs coefs' (x0 - fx/fpx)
    where
        fx = evalPoly coefs x0
        fpx = evalPoly coefs' x0

d :: RealFloat a => Complex a -> Complex a -> a
d c1 c2 = magnitude (c2 - c1)

pix :: Int -> Float -> Int -> Int -> PixelRGB8
pix st sz i j =
    let z = realToFrac i / sz :+ realToFrac j / sz
        zr = newton st coefs coefs' z
        (dist, idx) = minimum (zip [d zr r | r <- roots] [0..])
    in colors V.! idx


import Control.Monad.Random
import Learning
import Data.Set
import Numeric.LinearAlgebra

uniformInitializer :: Initializer (Rand StdGen) Double
uniformInitializer 0 set = pure empty
uniformInitializer k set = do
  i <- getRandomR (0, Data.Set.size set - 1)
  let a = elemAt i set
  rest <- uniformInitializer (k - 1) (deleteAt i set)
  return (singleton a `union` rest)

main :: IO ()
main = do
  stuff <- rand 1000000 100
  [lin] <- toRows <$> rand 1 100
  [noise] <- toRows <$> randn 1 1000000
  let res = stuff #> lin
  putStrLn "\nPrinting least squares fit, then real model which was learned with some noise:\n"
  print $ (ols stuff (res + noise) - lin)

--  stuff <- (Data.Set.fromList . toRows) <$> randn 1000 3
--  (centroids, clusters) <- evalRandIO $ kmeans uniformInitializer 3 stuff
--  print centroids


{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module AOC.Common.Numeric (
    fft
  , ifft
  , convolve
  , rconvolve
  , zconvolve
  , FFT.FFTWReal
  ) where

import           Data.Complex
import           GHC.TypeNats
import qualified Data.Array.CArray         as CA
import qualified Data.Array.IArray         as IA
import qualified Data.Ix                   as Ix
import qualified Data.Vector.Generic       as VG
import qualified Data.Vector.Generic.Sized as SVG
import qualified Foreign.Storable          as FS
import qualified Math.FFT                  as FFT
import qualified Math.FFT.Base             as FFT

fft :: (FFT.FFTWReal a, VG.Vector v (Complex a))
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v n (Complex a)
fft = SVG.withVectorUnsafe $
        fromCA
      . FFT.dft
      . toCA

ifft
    :: (FFT.FFTWReal a, VG.Vector v (Complex a))
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v n (Complex a)
ifft = SVG.withVectorUnsafe $
        fromCA
      . FFT.idft
      . toCA

fromCA
    :: (FS.Storable a, VG.Vector v (Complex a))
    => CA.CArray Int (Complex a)
    -> v (Complex a)
fromCA v = VG.generate (Ix.rangeSize (IA.bounds v)) (v IA.!)

toCA
    :: (FS.Storable a, VG.Vector v (Complex a))
    => v (Complex a)
    -> CA.CArray Int (Complex a)
toCA v = IA.listArray (0, VG.length v - 1) (VG.toList v)

-- | FFT-based convolution
convolve
    :: ( VG.Vector v (Complex a)
       , KnownNat n, 1 <= n
       , KnownNat m, 1 <= m
       , FFT.FFTWReal a
       )
    => SVG.Vector v n (Complex a)
    -> SVG.Vector v m (Complex a)
    -> SVG.Vector v (n + m - 1) (Complex a)
convolve x y = ifft $ fft x' * fft y'
  where
    x' = x SVG.++ 0
    y' = y SVG.++ 0

-- | FFT-based real-valued convolution
rconvolve
    :: ( VG.Vector v (Complex a)
       , VG.Vector v a
       , KnownNat n, 1 <= n
       , KnownNat m, 1 <= m
       , FFT.FFTWReal a
       )
    => SVG.Vector v n a
    -> SVG.Vector v m a
    -> SVG.Vector v (n + m - 1) a
rconvolve x y = SVG.map realPart $ convolve (SVG.map (:+ 0) x) (SVG.map (:+ 0) y)

-- | FFT-based integral convolution
zconvolve
    :: ( VG.Vector v (Complex Double)
       , VG.Vector v Double
       , VG.Vector v a
       , KnownNat n, 1 <= n
       , KnownNat m, 1 <= m
       , Integral a
       )
    => SVG.Vector v n a
    -> SVG.Vector v m a
    -> SVG.Vector v (n + m - 1) a
zconvolve x y = SVG.map (round @Double) $
    rconvolve (SVG.map fromIntegral x) (SVG.map fromIntegral y)


{-# LANGUAGE CPP #-}
module Driver (
    E,

    Nat(..),
    Nat2(..),
    Pos(..),
    Pos2(..),
    Index(..),
    Index2(..),
    Assocs(..),
    Assocs2(..),
    BandedAssocs(..),

    module Control.Arrow,
    module Control.Monad,
    module Control.Monad.ST,

    AEq,
    (===),
    (~==),
    module Data.Function,
    module Data.Ix,
    module Data.List,
    module Data.Ord,

    module Debug.Trace,

    module Test.QuickCheck,

    module Text.Printf,

    field,

    module Test.Framework,
    module Test.Framework.Providers.QuickCheck2,
    ) where

import Control.Arrow
import Control.Monad
import Control.Monad.ST

import Data.AEq( AEq )
import qualified Data.AEq as AEq
import Data.Complex
import Data.Ix
import Data.Function
import Data.List
import Data.Ord

import Debug.Trace

import System.IO
import System.Random

import Test.Vector.Dense()
import Test.Matrix.Dense()
import Test.Matrix.Banded()

import Test.QuickCheck hiding ( vector )
import Test.QuickCheck.BLAS( Nat(..), Nat2(..), Pos(..), Pos2(..),
    Index(..), Index2(..), Assocs(..), Assocs2(..), BandedAssocs(..) )
import qualified Test.QuickCheck.BLAS as Test

import Test.Framework
import Test.Framework.Providers.QuickCheck2

import Text.Printf
import Text.Show.Functions

#ifdef COMPLEX
field = "Complex Double"
type E = Complex Double
#else
field = "Double"
type E = Double
#endif

infix 4 ===, ~==

x === y | (AEq.===) x y = True
        | otherwise = trace (printf "expected `%s', but got `%s'" (show y) (show x)) False

x ~== y | (AEq.~==) x y = True
        | otherwise = trace (printf "expected `%s', but got `%s'" (show y) (show x)) False


module Main (main) where

-- package の import
-- CSVの読み込み
import qualified    Text.CSV                                as CSV
import qualified    Data.Text                               as T
import              Data.Text
import              Data.Attoparsec.Text                    as DAT

-- 可視化
import              Graphics.Rendering.Chart.Easy            hiding ( (:<))
import              Graphics.Rendering.Chart.Backend.Cairo
import              Graphics.Rendering.Chart.Axis
import              Graphics.Rendering.Chart.Axis.Int
import              Graphics.Rendering.Chart.Grid

-- 重回帰
import Statistics.Regression

-- その他
import qualified    Control.Monad                           as CM
import qualified    Data.List                               as L
import qualified    Data.Vector.Unboxed                     as VU

-- | Doubleのパーサ
{-# INLINE parseDouble #-}
parseDouble :: Text -> Double
parseDouble tx = case DAT.parseOnly DAT.double tx of
        Right r   -> r
        Left  l   -> error $ "Error on parseDouble : " ++ show l

type Title      = String
type FileName   = String
type Label      = String

-- | 折れ線グラフの作成
plotLine :: (PlotValue a, PlotValue b) => [(Label, [(a, b)])] -> FileName -> Title -> IO ()
plotLine xs file title
    = toFile def (file ++ ".png") $ do
        layout_title .= title
        CM.forM_ xs $ \(title,ys) -> plot $ line title [ys]

plotBar :: [(String,[Double])] -> FileName -> Title -> [Label] -> IO ()
plotBar xs file title labels
    = toFile def (file ++ ".png") $ do
    layout_title .= title
    layout_title_style . font_size .= 10
    layout_x_axis . laxis_generate .= autoIndexAxis (L.map fst xs)
    plot $ fmap plotBars $ bars labels (addIndexes (L.map snd xs))

-- | 散布図の作成
plotGridPoints :: (PlotValue a, PlotValue b)
             => [[[(Label, [(a, b)])]]] -> FileName -> Title -> IO ()
plotGridPoints xs file name
    = let  ys   = aboveN $ (flip L.map) xs
                $ \ cols  -> besideN
                $ (flip L.map) cols -- columns
                $ \ series -> layoutToGrid
                $ execEC
                $ CM.forM_ series
                $ \(t,xs) -> plot $ points t xs
    in CM.void  $ renderableToFile def (file ++ ".png")
                $ fillBackground def
                $ gridToRenderable
                $ title `wideAbove` ys
  where
    title = setPickFn nullPickFn $ label ls HTA_Centre VTA_Centre name
    ls = def { _font_size   = 15 , _font_weight = FontWeightBold }

main = do
    -- CSVファイルの読み込み
    -- Dataフォルダを作成し,そこにデータを入れておきましょう
    xs <- CSV.parseCSVFromFile "Data/multiple_regression_test.csv" >>= \res
       -> case res of
            Right xs -> return $ L.map (L.map T.pack) xs
            Left  xs -> error $ "error on parseCSVFromFile: " ++ show xs

    -- Headerの抽出 Vectorへの変換
    let headers = L.map T.unpack
                $ L.head xs

    -- 値のDoubleへの変換 Vectorへの変換
    let body    = L.transpose
                $ L.map (L.map parseDouble)
                $ L.tail xs

    let year        = body !! 0
        death       = body !! 1
        index       = body !! 2
        hospital    = body !! 3

    -- データをグラフで表示
    plotLine    (L.map (\(x,y) -> (x, L.zip year y))
                [("Number of Death", death)
                ,("Economic Indicators", index)
                ,("Number of Hospitals", hospital)])
                "Data/initial_data"
                "Data Visualization"

    -- 散布図行列を作成する
    plotGridPoints  [[[(xt ++ "_" ++ yt, L.zip xs ys)]
                     | (yt,ys) <- L.zip (L.tail headers) (L.tail body)]
                    | (xt,xs) <- L.zip (L.tail headers) (L.tail body)]
                   "Data/plot_matrix"
                   "Plot Matrix"

    -- heatmap の作成はちょっと難しいのでpass

    -- 重回帰
    let (coeffs, r2) = olsRegress   [ VU.fromList index
                                    , VU.fromList hospital]
                                    ( VU.fromList death)

    putStrLn $ "回帰係数:" ++ show (L.init (VU.toList coeffs))
    putStrLn $ "切片:"     ++ show (L.last (VU.toList coeffs))

    -- 結果のプロット
    plotBar (L.zipWith  (\x y -> (x,[y]))
                        ["Economic Indicators", "Number of Hospitals"]
                        (L.init (VU.toList coeffs)))
            "Data/bar_result"
            "Plot Result"
            ["Coeff"]


{-|
Module      : Qubism.QGate
Description : Types and functions for quantum gates
Copyright   : (c) Keith Pearson, 2018
License     : MIT
Maintainer  : keith@qubitrot.org
-}

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

module Qubism.QGate
  ( QGate
  , (#>)
  , gate
  , ident
  , pauliX
  , pauliY
  , pauliZ
  , hadamard
  , unitary
  , cnot
  , controlled
  , ifBit
  , kronecker
  , onJust
  , onEvery
  , onRange
  ) where

-- For dependent typing
import GHC.TypeLits
import Data.Singletons
import Data.Singletons.TypeLits
import Data.Finite

import           Data.Complex
import           Data.Monoid
import           Control.Monad.Trans.State.Strict
import           Numeric.LinearAlgebra ((><))
import qualified Numeric.LinearAlgebra as LA

import Qubism.Algebra
import Qubism.StateVec
import Qubism.CReg

newtype QGate (n :: Nat) =
  UnsafeMkQGate (LA.Matrix C) -- must be a (2^n) x (2^n) unitary matrix
  deriving (Show)

-- | "Close enough" equality testing
instance Eq (QGate n) where
  (UnsafeMkQGate a) == (UnsafeMkQGate b) =
    LA.norm_2 (a - b) < 0.000001

instance KnownNat n => Semigroup (QGate n) where
  (UnsafeMkQGate a) <> (UnsafeMkQGate b) = UnsafeMkQGate $ a LA.<> b

instance KnownNat n => Monoid (QGate n) where
  mempty = ident

instance KnownNat n => VectorSpace (QGate n) where
  zero = UnsafeMkQGate . (l><l) $ repeat 0 where l = internalLen (sing :: Sing n)
  z                 .: (UnsafeMkQGate b) = UnsafeMkQGate $ LA.scalar z * b
  (UnsafeMkQGate a) +: (UnsafeMkQGate b) = UnsafeMkQGate $ a + b
  neg (UnsafeMkQGate a) = UnsafeMkQGate $ -a

instance KnownNat n => Algebra (QGate n) where
  (*:) = (<>)

internalLen :: (KnownNat n, Num a) => Sing n -> a
internalLen = (2 ^) . fromIntegral . fromSing

-- | Apply a quantum gate to a quantum register. Note that this
-- operator conflicts with the one from Numeric.LinearAlgebra.
infixr 5 #>
(#>) :: QGate n -> StateVec n -> StateVec n
(#>) (UnsafeMkQGate m) (UnsafeMkStateVec v) = UnsafeMkStateVec $ m LA.#> v

-- | Helper function to make state computations nicer
gate :: Monad m => QGate n -> StateT (StateVec n) m ()
gate g = state $ \qr -> ((), g #> qr)

ident :: forall n . KnownNat n => QGate n
ident = UnsafeMkQGate . LA.ident $ internalLen (sing :: Sing n)

-- | Also known as a NOT gate
pauliX :: QGate 1
pauliX = UnsafeMkQGate $
  (2><2) [0 :+ 0, 1 :+ 0,
          1 :+ 0, 0 :+ 0]

pauliY :: QGate 1
pauliY = UnsafeMkQGate $
  (2><2) [0 :+ 0, 0 :+ (-1),
          0 :+ 1, 0 :+ 0]

pauliZ :: QGate 1
pauliZ = UnsafeMkQGate $
  (2><2) [1 :+ 0, 0 :+ 0,
          0 :+ 0, (-1) :+ 0]

hadamard :: QGate 1
hadamard = UnsafeMkQGate $ 1 / sqrt 2 *
  (2><2) [1 :+ 0, 1 :+ 0,
          1 :+ 0, (-1) :+ 0]

-- | An arbitrary element of SU(2),
-- U(theta, phi, lambda)
unitary :: Double -> Double -> Double -> QGate 1
unitary theta phi lambda = UnsafeMkQGate $
  (2><2) [a, b, c, d]
  where a =  cis (phi+lambda/2) * ( cos (theta/2) :+ 0 )
        b = -cis (phi-lambda/2) * ( sin (theta/2) :+ 0 )
        c =  cis (phi-lambda/2) * ( sin (theta/2) :+ 0 )
        d =  cis (phi+lambda/2) * ( cos (theta/2) :+ 0 )

-- | Controlled not between control and target qubits
cnot :: KnownNat n => Finite n -> Finite n -> QGate n
cnot c t = controlled c . onJust t $ pauliX

-- | Transform an arbitrary QGate into a controled one depending on qubit i
controlled :: forall n . KnownNat n => Finite n -> QGate n -> QGate n
controlled finite (UnsafeMkQGate m) =
  UnsafeMkQGate $ (m <> projection) + LA.ident (2^n) - projection
  where
    projection = LA.diag $ (2^n) LA.|> fmap f [0..]
    f j = fromIntegral $ j `quot` 2^(n-i-1) `mod` 2
    n   = fromIntegral $ fromSing (sing :: Sing n)
    i   = fromIntegral $ getFinite finite

-- | If the given Bit is One then apply the gate, otherwise
--  do nothing
ifBit :: KnownNat n => Bit -> QGate n -> QGate n
ifBit b g = if (b == One) then g else ident

-- | The tensor product of a QGate's A and B is a QGate that acts as A on the
-- first n qubits and B on the rest. In the computational basis this is simply
-- the kronecker product of matricies.
kronecker :: QGate n -> QGate m -> QGate (m+n)
kronecker (UnsafeMkQGate a) (UnsafeMkQGate b) =
  UnsafeMkQGate $ LA.kronecker a b

-- | Promote a 1-qubit gate to an n-qubit gate with the original gate acting
-- on qubit i. Other qubits are unaffected.
onJust :: forall n . KnownNat n => Finite n -> QGate 1 -> QGate n
onJust i (UnsafeMkQGate m) = UnsafeMkQGate $  -- it should be possible to do
  pre `LA.kronecker` m `LA.kronecker` post    -- this with just QGate's
  where pre  = LA.ident $ 2^j                 -- kronecker, but I can't get
        post = LA.ident $ 2^(n-j-1)           -- the types to check.
        j    = fromIntegral $ getFinite i
        n    = fromIntegral $ fromSing (sing :: Sing n)

-- | Promote a 1-qubit gate to an n-qubit gate which acts on every qubit
-- identically.
onEvery :: forall n . KnownNat n => QGate 1 -> QGate n
onEvery (UnsafeMkQGate m) = UnsafeMkQGate $ iterate (LA.kronecker m) m !! (n-1)
  where n = fromIntegral $ fromSing (sing :: Sing n)

-- | Promote a 1-qubit gate to an n-qubit gate which acts on a range of
-- qubits, leaving others untouched
onRange :: forall n . KnownNat n => Finite n -> Finite n -> QGate 1 -> QGate n
onRange f l m = mconcat $ map (\i -> onJust i m) [f..l]


module Main where

import Embadd

import qualified Data.Map.Strict as M (fromList, toList, size, lookup)
import Numeric.LinearAlgebra as LA
import System.Environment

main :: IO ()
main = do
  (vectorFile:args) <- getArgs
  case vectorFile of
    "--help" -> putStrLn "Usage: ./emb.hs FILE"
    "NORM" -> computeNorm (head args)
    _ -> do
      vectorCont <- readFile vectorFile
      let vLine (w:ws) = (w, vector $ map read ws :: Vector Double)
          vectorDict = M.fromList . map (vLine . words) . drop 1 . lines $ vectorCont
      putStrLn "Embedding vizualization tool"
      putStrLn "============================"
      putStrLn ""
      putStrLn "W1 + W2 - W3 -> nearest neighbours of v(W1) + v(W2) - v(W3); you can use any expression with words, '+' and '-'"
      putStrLn "PCA W1 W2 ... Wn  -> pca.png will contain {v(W1), ..., v(Wn)} projected to their principal components"
      readEvalPrintLoop vectorDict


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Binomial
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The binomial distribution.  This is the discrete probability
-- distribution of the number of successes in a sequence of /n/
-- independent yes\/no experiments, each of which yields success with
-- probability /p/.

module Statistics.Distribution.Binomial
    (
      BinomialDistribution
    -- * Constructors
    , binomial
    -- * Accessors
    , bdTrials
    , bdProbability
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import qualified Statistics.Distribution as D
import qualified Statistics.Distribution.Poisson.Internal as I
import Numeric.SpecFunctions (choose,incompleteBeta)
import Numeric.MathFunctions.Constants (m_epsilon)
import Data.Binary (put, get)
import Control.Applicative ((<$>), (<*>))

-- | The binomial distribution.
data BinomialDistribution = BD {
      bdTrials      :: {-# UNPACK #-} !Int
    -- ^ Number of trials.
    , bdProbability :: {-# UNPACK #-} !Double
    -- ^ Probability.
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON BinomialDistribution
instance ToJSON BinomialDistribution

instance Binary BinomialDistribution where
    put (BD x y) = put x >> put y
    get = BD <$> get <*> get

instance D.Distribution BinomialDistribution where
    cumulative = cumulative

instance D.DiscreteDistr BinomialDistribution where
    probability = probability

instance D.Mean BinomialDistribution where
    mean = mean

instance D.Variance BinomialDistribution where
    variance = variance

instance D.MaybeMean BinomialDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance BinomialDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy BinomialDistribution where
  entropy (BD n p)
    | n == 0 = 0
    | n <= 100 = directEntropy (BD n p)
    | otherwise = I.poissonEntropy (fromIntegral n * p)

instance D.MaybeEntropy BinomialDistribution where
  maybeEntropy = Just . D.entropy

-- This could be slow for big n
probability :: BinomialDistribution -> Int -> Double
probability (BD n p) k
  | k < 0 || k > n = 0
  | n == 0         = 1
  | otherwise      = choose n k * p^k * (1-p)^(n-k)

-- Summation from different sides required to reduce roundoff errors
cumulative :: BinomialDistribution -> Double -> Double
cumulative (BD n p) x
  | isNaN x      = error "Statistics.Distribution.Binomial.cumulative: NaN input"
  | isInfinite x = if x > 0 then 1 else 0
  | k <  0       = 0
  | k >= n       = 1
  | otherwise    = incompleteBeta (fromIntegral (n-k)) (fromIntegral (k+1)) (1 - p)
  where
    k = floor x

mean :: BinomialDistribution -> Double
mean (BD n p) = fromIntegral n * p

variance :: BinomialDistribution -> Double
variance (BD n p) = fromIntegral n * p * (1 - p)

directEntropy :: BinomialDistribution -> Double
directEntropy d@(BD n _) =
  negate . sum $
  takeWhile (< negate m_epsilon) $
  dropWhile (not . (< negate m_epsilon)) $
  [ let x = probability d k in x * log x | k <- [0..n]]

-- | Construct binomial distribution. Number of trials must be
--   non-negative and probability must be in [0,1] range
binomial :: Int                 -- ^ Number of trials.
         -> Double              -- ^ Probability.
         -> BinomialDistribution
binomial n p
  | n < 0          =
    error $ msg ++ "number of trials must be non-negative. Got " ++ show n
  | p < 0 || p > 1 =
    error $ msg++"probability must be in [0,1] range. Got " ++ show p
  | otherwise      = BD n p
    where msg = "Statistics.Distribution.Binomial.binomial: "


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

import qualified Text.Printf as TP (printf)
import qualified Data.List as L (intercalate)

import Data.Array.Repa (Z(Z), (:.)((:.)), (!), (-^), (+^), (*^), All(All), Any(Any), extent, size)
import qualified Data.Array.Repa as RP (fromFunction, fromListUnboxed, Array, U, D, DIM1, DIM2, DIM4, map)
import qualified Data.Array.Repa as RP (sumAllS, slice, delay, computeS, toList)
import qualified Data.Array.Repa.Repr.Vector as RP (V, fromListVector)
import Data.Complex (Complex((:+)), realPart, imagPart)
import Data.Bits (bit)

pdim1 :: RP.Array RP.U RP.DIM1 (Complex Double) -> String
pdim1 mat = pt ++ "\n" where
    (Z :. n) = extent mat
    pt = concat [ let (xr :+ xi) = mat ! (Z :. i) in TP.printf "%14.4f +I %8.4f" xr xi | i <- [0 .. n - 1] ]

pdim2 :: RP.Array RP.U RP.DIM2 (Complex Double) -> String
pdim2 mat = L.intercalate "\n" [ pt i | i <- [0 .. m - 1]] ++ "\n\n" where
    (Z :. m :. n) = extent mat
    pt i = concat [ let (xr :+ xi) = mat ! (Z :. i :. j) in TP.printf "%14.4f +I %8.4f" xr xi | j <- [0 .. n - 1] ]

identity :: RP.Array RP.U RP.DIM2 (Complex Double)
identity = RP.fromListUnboxed (Z :. 2 :. 2) [1, 0, 0, 1]

sigmaX :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaX = RP.fromListUnboxed (Z :. 2 :. 2) [0, 1, 1, 0]

sigmaY :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaY = RP.fromListUnboxed (Z :. 2 :. 2) [0, 0 :+ (-1), 0 :+ 1, 0]

sigmaZ :: RP.Array RP.U RP.DIM2 (Complex Double)
sigmaZ = RP.fromListUnboxed (Z :. 2 :. 2) [1, 0, 0, -1]

-- | Tensor product.
kron :: RP.Array RP.U RP.DIM2 (Complex Double)
     -> RP.Array RP.U RP.DIM2 (Complex Double)
     -> RP.Array RP.U RP.DIM2 (Complex Double)
kron mata matb = reduceBlock . RP.computeS $ RP.fromFunction (extent mata) iter where
    iter i = RP.computeS $ RP.map (* (mata ! i)) matb :: RP.Array RP.U RP.DIM2 (Complex Double)

-- | Heisenberg Hamiltonian
hamil :: RP.Array RP.U RP.DIM2 (Complex Double)
hamil = RP.computeS $ RP.map (* 0.25) (((kron sigmaX identity) `mmult` (kron identity sigmaX))
    +^ ((kron sigmaY identity) `mmult` (kron identity sigmaY))
    +^ ((kron sigmaZ identity) `mmult` (kron identity sigmaZ)))

-- | S^2 Operator
s2 :: RP.Array RP.U RP.DIM2 (Complex Double)
s2 = RP.computeS $ RP.map (* 0.25) (
    square (RP.computeS $ kron sigmaX identity +^ kron identity sigmaX)
    +^ square (RP.computeS $ kron sigmaY identity +^ kron identity sigmaY)
    +^ square (RP.computeS $ kron sigmaZ identity +^ kron identity sigmaZ))

-- | Square of a matrix
square :: RP.Array RP.U RP.DIM2 (Complex Double) -> RP.Array RP.U RP.DIM2 (Complex Double)
square mat = mmult mat mat

-- | Matrix Matrix Multiplication
mmult :: RP.Array RP.U RP.DIM2 (Complex Double)
      -> RP.Array RP.U RP.DIM2 (Complex Double)
      -> RP.Array RP.U RP.DIM2 (Complex Double)
mmult mata matb = RP.computeS $ RP.fromFunction (Z :. n :. m) iter where
    (Z :. n :. k) = extent mata
    (Z :. kp :. m) = extent matb
    iter :: RP.DIM2 -> Complex Double
    iter (Z :. i :. j) = RP.sumAllS $ RP.slice mata (Any :. i :. All) *^ RP.slice matb (Any :. All :. j)

-- | Matrix Vector Multiplication
mvmult :: RP.Array RP.U RP.DIM2 (Complex Double)
       -> RP.Array RP.U RP.DIM1 (Complex Double)
       -> RP.Array RP.U RP.DIM1 (Complex Double)
mvmult mat vec = RP.computeS $ RP.fromFunction (Z :. n) iter where
    (Z :. n :. m) = extent mat
    iter :: RP.DIM1 -> Complex Double
    iter (Z :. i) = RP.sumAllS $ RP.slice mat (Any :. i :. All) *^ vec

-- | Vector Vector Dot Product
dot :: RP.Array RP.U RP.DIM1 (Complex Double) -> RP.Array RP.U RP.DIM1 (Complex Double) -> Complex Double
dot veca vecb = RP.sumAllS $ veca *^ vecb

-- | Reduce blocked matrix to regular matrix.
reduceBlock :: RP.Array RP.V RP.DIM2 (RP.Array RP.U RP.DIM2 (Complex Double)) -- ^ matrix of blocks
            -> RP.Array RP.U RP.DIM2 (Complex Double)
reduceBlock mat = RP.computeS $ RP.fromFunction (Z :. nl :. nl) iter where
    f :: Int -> RP.Array RP.U RP.DIM2 (Complex Double) -> [(Int, Int)]
    f i block = let (Z :. _ :. n) = extent block in map (i, ) [0..n - 1]
    idx = let l = concat (zipWith f [0..] (RP.toList (RP.slice mat (Any :. (0 :: Int) :. All))))
        in RP.fromListUnboxed (Z :. (length l)) l
    (Z :. nl) = extent idx
    iter (Z :. i :. j) = let (ip, im) = idx ! (Z :. i); (jp, jm) = idx ! (Z :. j)
        in (mat ! (Z :. ip :. jp)) ! (Z :. im :. jm)

-- alpha = 0, beta = 1
-- |babb> = 1011 = 11
stateIndex :: String -> Int
stateIndex x = iter (length x - 1) x where
    iter n (x:xs) = (if n == 0 then 0 else iter (n - 1) xs) + (if x == 'a' then 0 else bit n)

-- | The kth eigenstate represented as n-component vector
eigenState :: Int -> Int -> RP.Array RP.U RP.DIM1 (Complex Double)
eigenState n k = RP.computeS $ RP.fromFunction (Z :. n) (\(Z :. i) -> if k == i then 1 else 0)

state_ab = eigenState (bit 2) $ stateIndex "ab"
state_ba = eigenState (bit 2) $ stateIndex "ba"
state_aa = eigenState (bit 2) $ stateIndex "aa"
state_bb = eigenState (bit 2) $ stateIndex "bb"
singlet = RP.computeS $ RP.map (* (2 ** (-0.5))) (state_ab -^ state_ba)
triplet = RP.computeS $ RP.map (* (2 ** (-0.5))) (state_ab +^ state_ba)

main = do
    putStrLn $ "[Singlet] (Si . Sj) S^2"
    print $ let state = singlet in dot state (hamil `mvmult` state)
    print $ let state = singlet in dot state (s2 `mvmult` state)
    putStrLn $ "[Triplet] (Si . Sj) S^2"
    print $ let state = triplet in dot state (hamil `mvmult` state)
    print $ let state = triplet in dot state (s2 `mvmult` state)
    putStrLn $ "[|aa>] (Si . Sj) S^2"
    print $ let state = state_aa in dot state (hamil `mvmult` state)
    print $ let state = state_aa in dot state (s2 `mvmult` state)


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications, DataKinds, GADTs #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, FunctionalDependencies #-}

module Lib
    ( someFunc
    ) where

import Data.Maybe ( fromMaybe )
import Data.Int (Int32)
import Linear.V1
import Linear.V3
import Linear.V4 hiding (vector)
import Linear.Matrix
import Linear.Metric (dot, norm)
import Graphics.UI.GLUT ( Key(..), SpecialKey(..), MouseButton(..), KeyState(..) )
import Vis
import Vis.Vis ( vis )
import Vis.Camera ( makeCamera, setCamera, cameraMotion, cameraKeyboardMouse, Camera )
import Control.Lens
import Control.Applicative
import Control.Monad (forever)
import Numeric.AD
import Numeric.LinearAlgebra ((#>), det, fromLists, fromList, toList, pinv, inv, vector, matrix, Vector, Matrix, scale, norm_2)
--import Debug.Trace

rotM33XYZ :: (Floating f) => f -> f -> f -> M33 f
rotM33XYZ r p y = rotM33Z y !*! rotM33Y p !*! rotM33X r

rotM33X :: (Floating f) => f -> M33 f
rotM33X a = V3 (V3 1 0 0) (V3 0 ca (-sa))  (V3 0 sa ca) where
    ca = cos a
    sa = sin a

rotM33Y :: (Floating f) => f -> M33 f
rotM33Y a = V3 (V3 ca 0 sa) (V3 0 1 0)  (V3 (-sa) 0 ca) where
    ca = cos a
    sa = sin a

rotM33Z :: (Floating f) => f -> M33 f
rotM33Z a = V3 (V3 ca (-sa) 0)  (V3 sa ca 0) (V3 0 0 1) where
    ca = cos a
    sa = sin a

printV4 :: (Show f) => V4 f -> IO ()
printV4 (V4 v1 v2 v3 v4) = print [v1, v2, v3, v4]

printV3 :: (Show f) => V3 f -> IO ()
printV3 (V3 v1 v2 v3) = print [v1, v2, v3]

printM33 :: (Show f) => M33 f -> IO ()
printM33 (V3 v1 v2 v3) = do
    putStrLn ""
    printV3 v1
    printV3 v2
    printV3 v3
    putStrLn ""

printM44 :: (Show f) => M44 f -> IO ()
printM44 (V4 v1 v2 v3 v4) = do
    putStrLn ""
    printV4 v1
    printV4 v2
    printV4 v3
    printV4 v4
    putStrLn ""

rotM44X :: (Floating f) => f -> M44 f
rotM44X a = identity & _m33 .~ rotM33X a

rotM44Z :: (Floating f) => f -> M44 f
rotM44Z a = identity & _m33 .~ rotM33Z a

transM44X :: (Floating f) => f -> M44 f
transM44X a =  identity & translation .~ V3 a 0 0

transM44Z :: (Floating f) => f -> M44 f
transM44Z a =  identity & translation .~ V3 0 0 a

twistX :: (Floating f) => f -> f -> M44 f
twistX d theta = transM44X d !*! rotM44X theta

twistZ :: (Floating f) => f -> f -> M44 f
twistZ d theta = transM44Z d !*! rotM44Z theta

skew :: (Floating f) => V3 f -> M33 f
skew (V3 x y z) = V3 (V3 0 (-z) y) (V3 z 0 (-x)) (V3 (-y) x 0)

invSkew :: (Floating f) => M33 f -> V3 f
invSkew (V3 (V3 _ _ y) (V3 z _ _) (V3 _ x _)) = V3 x y z

rodrigues :: (Eq f, Floating f) => V3 f -> M33 f
rodrigues r = identity !!* ct  !+! rrt !!* (1 - ct) !+! sr !!* st
    where
    theta = norm r
    ct    = cos theta
    st    = sin theta
    r'    = if theta == 0 then V3 1 0 0 else r / V3 theta theta theta
    sr    = skew r'
    mrt   = V1 r'
    mr    = transpose mrt
    rrt   = mr !*! mrt

invRodrigues :: (Floating f) => M33 f -> V3 f
invRodrigues m = V3 x' y' z' * V3 t t t
    where
    V3 x y z = invSkew (m - transpose m) / 2
    st = norm (V3 x y z)
    t  = asin st
    V3 x' y' z' = V3 x y z / V3 t t t

frame2xyzaxayaz :: (Floating f) => M44 f -> [f]
frame2xyzaxayaz m = [x, y, z, ax, ay, az]
    where
    V3 ax ay az = invRodrigues $ m ^. _m33
    V3 x y z = m ^. translation

class Renderable b f  | b -> f where
    renderFC :: Flavour -> Color -> b -> VisObject f
    render   :: b -> VisObject f
    render = renderFC Wireframe blue

-- M44 を render
instance (Floating f) => Renderable (M44 f) f where
    renderFC _ _ m = convfrm m $ axes (0.05, 10)
        where
        axes (size, aspectRatio) = VisObjects [ Arrow (size, aspectRatio) (V3 1 0 0) red
                                              , Arrow (size, aspectRatio) (V3 0 0 1) blue
                                              ]

class FrameConvertable b f | b -> f where
    convfrm :: M44 f -> b -> b

instance (Floating f) => FrameConvertable (VisObject f) f where
    convfrm m v = Trans t $ RotDcm rot v
        where
        rot = transpose $ m ^. _m33 -- NOTE: vis function use row vector??
        t   = m ^. translation

---- for rendering and collision detection
data LinkShape f = OBB (V3 f) (V3 f) (M33 f) -- center size/2 axis-vector
    deriving Show

instance (Floating f) => FrameConvertable (LinkShape f) f where
    convfrm m (OBB c a r) = OBB c' a r'
        where
        c' = (m !* point c) ^. _xyz
        r' = m^._m33 !*! r

instance (Floating f) => Renderable (LinkShape f) f where
    renderFC f color (OBB c b@(V3 bx by bz) m) = Trans c $ RotDcm rot $ Box (2*bx, 2*by, 2*bz) f color
        where
        rot = transpose $ m ^. _m33 -- NOTE: vis function use row vector??

hasIntersection :: (Ord f, Floating f) => LinkShape f -> LinkShape f -> Bool
hasIntersection (OBB c0 a@(V3 a0 a1 a2)  m0) (OBB c1 b@(V3 b0 b1 b2) m1) = not $ easy || i0 || i1 || i2 || i3 || i4 || i5 || i6 || i7 || i8 || i9 || i10 || i11 || i12 || i13 || i14
    where

    -- not intersection test using bounding sphere
    easy = norm a + norm b < norm vD

    -- not intersection test
    i0  = a0 + (b0*c00' + b1*c01' + b2*c02') < abs sA0D
    i1  = a1 + (b0*c10' + b1*c11' + b2*c12') < abs sA1D
    i2  = a2 + (b0*c20' + b1*c21' + b2*c22') < abs sA2D

    i3  = (a0*c00' + a1*c10' + a2*c20') + b0 < abs (vB0 `dot` vD)
    i4  = (a0*c01' + a1*c11' + a2*c21') + b1 < abs (vB1 `dot` vD)
    i5  = (a0*c02' + a1*c12' + a2*c22') + b2 < abs (vB2 `dot` vD)

    i6  = (a1*c20' + a2*c10') + (b1*c02' + b2*c01') < abs (c10*sA2D - c20*sA1D)
    i7  = (a1*c21' + a2*c11') + (b0*c02' + b2*c00') < abs (c11*sA2D - c21*sA1D)
    i8  = (a1*c22' + a2*c12') + (b0*c01' + b1*c00') < abs (c12*sA2D - c22*sA1D)

    i9  = (a0*c20' + a2*c00') + (b1*c12' + b2*c11') < abs (c20*sA0D - c00*sA2D)
    i10 = (a0*c21' + a2*c01') + (b0*c12' + b2*c10') < abs (c21*sA0D - c01*sA2D)
    i11 = (a0*c22' + a2*c02') + (b0*c11' + b1*c10') < abs (c22*sA0D - c02*sA2D)

    i12 = (a0*c10' + a1*c00') + (b1*c22' + b2*c21') < abs (c00*sA1D - c10*sA0D)
    i13 = (a0*c11' + a1*c01') + (b0*c22' + b2*c20') < abs (c01*sA1D - c11*sA0D)
    i14 = (a0*c12' + a1*c02') + (b0*c21' + b1*c20') < abs (c02*sA1D - c12*sA0D)

    vD = c1 - c0
    V3 vA0 vA1  vA2 = transpose m0
    V3 vB0 vB1  vB2 = transpose m1
    V3 (V3 c00 c01 c02) (V3 c10 c11 c12) (V3 c20 c21 c22)  = transpose m0 !*! m1

    ((c00', c01', c02'), (c10', c11', c12'), (c20', c21', c22'))  = ((abs c00, abs c01, abs c02), (abs c10, abs c11, abs c12), (abs c20, abs c21, abs c22))

    sA0D = vA0 `dot` vD
    sA1D = vA1 `dot` vD
    sA2D = vA2 `dot` vD

-- TODO
hasCollisions :: (Ord f, Floating f) => [LinkShape f] -> [Bool]
hasCollisions (x:xs) = any id (hasIntersection x <$> xs) : hasCollisions xs
hasCollisions [] = []

aabb :: (Ord f, Floating f) => (f, f) -> (f, f) -> (f, f) -> LinkShape f
aabb (x0,x1) (y0,y1) (z0,z1) = OBB c b identity
        where
        c = V3 ((x1+x0)/2) ((y1+y0)/2) ((z1+z0)/2)
        b = V3 ((x1-x0)/2) ((y1-y0)/2) ((z1-z0)/2)

type FrameTrans f = f -> M44 f

data DH f = DH {   __a :: f      -- [m]
               ,   __d :: f      -- [m]
               ,   __alpha :: f  -- [rad]
               } deriving Show

makeLenses ''DH

d1 :: (Floating f) => f
d1 = 0.089159

a2 :: (Floating f) => f
a2 = 0.425

a3 :: (Floating f) => f
a3 = 0.39225

d4 :: (Floating f) => f
d4 = 0.10915

d5 :: (Floating f) => f
d5 = 0.09465

d6 :: (Floating f) => f
d6 = 0.0823

ur5dh :: (Floating f) => [DH f]
ur5dh =
    [ DH 0             d1          (pi/2) -- J1
    , DH (-a2)         0           0      -- J2
    , DH (-a3)         0           0      -- J3
    , DH 0             d4          (pi/2) -- J4
    , DH 0             d5          (-pi/2)-- J5
    , DH 0             d6          0      -- J6(flange)
    ]

tcpOffsetZ :: (Floating f) => f
tcpOffsetZ = 0.15

-- how to draw link (from frame{i} to frame{i+1})
ur5link :: (Ord f, Floating f) => [LinkShape f]
ur5link =
    [ aabb (-r/2, r/2) (-d1, r/2) (-r/2, r/2) -- 1,2
    , aabb (-r/2, a2 + r/2) (-r/2, r/2) (delta+r/2, 3*r/2) -- 2,3
    , aabb (-r/2, a3 + r/2) (-r/2, r/2) (-r/2, r/2) -- 3,4
    , aabb (-r/2, r/2) (r/2 + delta-d4, r/2) (-r/2, r/2) -- 4,5
    , aabb (-r/2, r/2) (-r/2, d5 - r/2 - delta) (-r/2, r/2) -- 5,6
    , aabb (-r/2, r/2) (-r/2, r/2) (delta-d6+r, 0) -- 6,frange
    , aabb (-r/2, r/2) (-r/2, r/2) (delta, tcpOffsetZ)  -- tool
    ]
    where
    r = 0.05
    delta = 0.01

-- https://en.wikipedia.org/wiki/Denavit%E2%80%93Hartenberg_parameters
dr2ft :: (Floating f) => DH f -> FrameTrans f
dr2ft dh theta = twistZ (dh ^. _d) theta !*!  twistX (dh^. _a) (dh^. _alpha)

ur5fts :: (Floating f) => [FrameTrans f]
ur5fts = dr2ft <$> ur5dh

-- calc each frame in base-coordinate for each joint-angle
kinematics :: (Floating f) => [FrameTrans f] -> [f] -> [M44 f]
kinematics (e:es) (a:as) = t : ((t !*! ) <$> kinematics es as)
    where
    t = e a
kinematics [] [] = []
kinematics _ _ = error "invalid args"

joint2tcp :: (Floating f) => [FrameTrans f] -> [f] -> [f]
joint2tcp es js = frame2xyzaxayaz $ last ( kinematics es js) !*! transM44Z tcpOffsetZ

poslabel = ["x", "y", "z", "ax", "ay", "az"]

data World = World { _jPos    :: [Double]
                   , _tcpVel  :: [Double]
                   , _trajct   :: [V3 Double]
                   , _camera :: Camera
                   }

makeLenses ''World

world0 = World [0.1, 0.4, 0.4, 0.1, 0.1, 0.1] [0, 0, 0, 0, 0, 0] [] cameraState0
    where
    defaultCamera = Camera0 { phi0 = 60 , theta0 = 20 , rho0 = 7 }
    cameraState0 = makeCamera $ fromMaybe defaultCamera (optInitialCamera defaultOpts)

motionPolicy :: [Double] -> [Double] ->  [Double] -> [Double]
motionPolicy goal pos vel = toList acc
    where
    alpha = 0.001
    c     = 0.01
    beta  = 0.5
    acc :: Vector Double
    acc = alpha * s (vector goal - vector pos) - beta * vector vel
    s :: Vector Double -> Vector Double
    s v = scale (1 / h (norm_2 v)) v
    h :: Double -> Double
    h z = z + c * log (1 + exp (-2) * c * z)

mySimulateIO :: IO ()
mySimulateIO = vis defaultOpts 0.016 world0 simFun drawFun setCameraFun (Just kmCallback) (Just motionCallback) Nothing
  where
    delta = 0.001

    simFun (w,t) = return $ w  &~ do
                        let tcpp = joint2tcp ur5fts $ w ^. jPos
                        let tcpv = w ^. tcpVel
                        let goal = [0.192428507378705,-0.39791168805212945,-0.43347675593345414,0.7684417340393249,-0.40013839074582813,0.3969016751654795]
                        let tcpa = motionPolicy goal tcpp tcpv
                        let tcpv' = zipWith (+) tcpv  tcpa
                        let j    = fromLists (jacobian (joint2tcp $ ur5fts) (w ^. jPos) :: [[Double]])
                        let jvel = if det j > 0.0001 then toList $ pinv j #> fromList tcpv else [0, 0, 0, 0, 0, 0]
                        let jpos' = zipWith (+) (w ^. jPos) jvel
                        let pos  = joint2tcp ur5fts jpos'
                        jPos .= jpos'
                        tcpVel .= tcpv'
                        trajct %= (:) (V3 (pos !! 0) (pos !! 1) (pos !! 2))

    drawFun (w, t) = do
        let jpos = w ^. jPos
        let frames =  identity : kinematics ur5fts jpos
        let vframes = render <$> frames
        let ts     = tail frames ++ [last frames]
        let links  = zipWith convfrm ts ur5link
        let colors  = (\b -> if b then red else blue) <$> hasCollisions links
        let vlinks  = zipWith (renderFC Wireframe) colors links
        let pos = joint2tcp ur5fts jpos
        let info = [Text2d (show (l, j)) (0, 110  - i * 20) Fixed8By13 red | (l, i,j) <- zip3 poslabel [0..] pos]
        let infoj = [Text2d (show ("J", i, j)) (300, 110  - i * 20) Fixed8By13 red | (i,j) <- zip [0..] jpos]
        let points = Points (w ^. trajct) (Just 1.0) yellow
        print pos
        return (VisObjects $ vlinks ++ vframes ++ info ++ infoj ++ [points], Nothing)

    setCameraFun w = setCamera $ w ^. camera

    kmCallback w k0 k1 _ _ = case (k0,  k1) of
                         (Char 'h', Down)            -> w & tcpVel .~ [ delta, 0, 0, 0, 0, 0]
                         (Char 'l', Down)            -> w & tcpVel .~ [-delta, 0, 0, 0, 0, 0]
                         (Char 'j', Down)            -> w & tcpVel .~ [0,  delta, 0, 0, 0, 0]
                         (Char 'k', Down)            -> w & tcpVel .~ [0, -delta, 0, 0, 0, 0]
                         (Char 'n', Down)            -> w & tcpVel .~ [0, 0,  delta, 0, 0, 0]
                         (Char 'p', Down)            -> w & tcpVel .~ [0, 0, -delta, 0, 0, 0]
                         (SpecialKey KeyLeft,  Down) -> w & tcpVel .~ [0, 0, 0,  delta, 0, 0]
                         (SpecialKey KeyRight, Down) -> w & tcpVel .~ [0, 0, 0, -delta, 0, 0]
                         (SpecialKey KeyDown,  Down) -> w & tcpVel .~ [0, 0, 0, 0,  delta, 0]
                         (SpecialKey KeyUp,    Down) -> w & tcpVel .~ [0, 0, 0, 0, -delta, 0]
                         (Char ',', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0,  delta]
                         (Char '.', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0, -delta]
                         (Char 'a', Down)            -> world0 & camera .~ (w ^. camera)
                         (Char 's', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0, 0]
                         (_, _)           -> w & camera %~ (\x -> cameraKeyboardMouse x k0 k1)

    motionCallback w pos = w & camera %~ (`cameraMotion` pos)

someFunc :: IO ()
someFunc = mySimulateIO

main = mySimulateIO


module Math.IRT.MLE.Truncated
  ( DF (..)
  , MLEResult (..)
  , mleEst
  ) where

import Data.Default.Class

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.MLE.Internal.Generic

data DF = DF { steps         :: !Int
             , thetaEstimate :: !Double
             , lower_bound   :: !Double
             , upper_bound   :: !Double }

instance Default DF where
  def = DF 10 0.0 (-3.5) 3.5

mleEst :: (Distribution d, ContDistr d, DensityDeriv d, LogLikelihood d) => DF -> [Bool] -> [d] -> MLEResult
mleEst (DF n th lb ub) rs params =
    let res = generic_mleEst rs params n th
    in res { theta = min ub $ max lb $ theta res }


{-# LANGUAGE TypeApplications #-}
module ArrayFire.BLASSpec where

import ArrayFire    hiding (not)

import Data.Complex
import Test.Hspec

spec :: Spec
spec =
  describe "BLAS spec" $ do
    it "Should matmul two matrices" $ do
      (matrix @Double (2,2) [[2,2],[2,2]] `matmul` matrix @Double (2,2) [[2,2],[2,2]]) None None
        `shouldBe` matrix @Double (2,2) [[8,8],[8,8]]
    it "Should dot product two vectors" $ do
      dot (vector @Double 2 (repeat 2)) (vector @Double 2 (repeat 2)) None None
        `shouldBe`
           scalar @Double 8
    it "Should produce scalar dot product between two vectors as a Complex number" $ do
      dotAll (vector @Double 2 (repeat 2)) (vector @Double 2 (repeat 2)) None None
        `shouldBe`
           8.0 :+ 0.0
    it "Should take the transpose of a matrix" $ do
      transpose (matrix @Double (2,2) [[1,1],[2,2]]) False
        `shouldBe`
           matrix @Double (2,2) [[1,2],[1,2]]
    it "Should take the transpose of a matrix in place" $ do
      let m = matrix @Double (2,2) [[1,1],[2,2]]
      transposeInPlace m False
      m `shouldBe` matrix @Double (2,2) [[1,2],[1,2]]


{-# LANGUAGE FlexibleContexts,FlexibleInstances,IncoherentInstances #-}
{- Assignment 1 Extra Credit
 - Name: TODO add full name
 - Date: TODO add of completion
 -}
module Assign_2_ExtraCredit where

import Data.Complex

macid = "TODO: put your mac id here"

data GaussianInt a = a :@ a
  deriving Show

class GaussianIntegral g where
  gaussZero :: Integral a => g a
  gaussReal :: Integral a => g a -> a
  gaussImag :: Integral a => g a -> a
  gaussConj :: Integral a => g a -> g a
  gaussAdd :: Integral a => g a -> g a -> g a
  gaussMult :: Integral a => g a -> g a -> g a

{- TODO
 -   implement instances of GaussianIntegral
 -   implmenet instances of Eq, Ord
 -}

gaussNorm :: (Integral a, GaussianIntegral g) => g a -> a
gaussNorm g = error "TODO: implement gaussNorm"

maxGaussNorm :: (Integral a, GaussianIntegral g) => [g a] -> g a
maxGaussNorm gs = error "TODO: implement maxGaussNorm"


module D03
    ( Claim (..), mkClaims, claimsIntersectionSize, isolatedClaim
    ) where

import Data.List
import qualified Data.HashMap.Strict as HM
import Data.Complex
import Text.Regex

data Claim = Claim { identifier :: Int,
                     topLeft :: Complex Int,
                     bottomRight :: Complex Int }
  deriving (Show, Eq)

type World = HM.HashMap (Complex Int) [Claim]

isolatedClaim :: String -> Int
isolatedClaim input = identifier $ head $ claims \\ badClaims
  where
    badClaims = concat $ HM.elems $ HM.filter (\cs -> (length cs) > 1) $ buildWorld claims
    claims = mkClaims input

claimsIntersectionSize :: String -> Int
claimsIntersectionSize = length . filter (> 1) . map length . HM.elems . buildWorld . mkClaims

buildWorld :: [Claim] -> World
buildWorld = foldl addClaimToWorld HM.empty

addClaimToWorld :: World -> Claim -> World
addClaimToWorld world claim = foldl addPoint world [x:+y | x<-[(realPart p1)..(realPart p2)], y<-[(imagPart p1)..(imagPart p2)]]
  where
    p1 = topLeft claim
    p2 = bottomRight claim
    addPoint w p = HM.insertWith (++) p [claim] w

mkClaims :: String -> [Claim]
mkClaims = map mkClaim . lines

mkClaim :: String -> Claim
mkClaim def = Claim (params !! 0) (left :+ top) ((left+width-1) :+ (top+height-1))
  where
    top = params !! 2
    left = params !! 1
    width = params !! 3
    height = params !! 4
    params = extract $ matchRegexAll (mkRegex "#([0-9]+) @ ([0-9]+),([0-9]+): ([0-9]+)x([0-9]+)") def
    extract (Just (_,_,_,subs)) = map read subs


{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.HamMC where

import Math.Probably.MCMC
import qualified Math.Probably.PDF as PDF
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import Control.Applicative

import Numeric.LinearAlgebra
import Numeric.AD
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S

import Debug.Trace
import Data.IORef
import Control.Spoon

data HMCPar = HMCPar { hpXi :: !(Vector Double),
                       hpL :: !Int,
                       hpPi :: !Double,
                       hpEpsilon :: !Double,
                       hpCount :: !Int,
                       hpAccept :: !Int,
                       hpFreezeEps :: Bool }
    deriving Show

hmc1 :: (Vector Double -> (Double,Vector Double))
         -> HMCPar
         -> Sampler HMCPar
hmc1 postGrad (HMCPar current_q l _ epsMean count accept freeze) = do
   let dims = dim current_q
       grad_u = negate . snd . postGrad
       u = negate . fst . postGrad
   current_p <- fmap fromList $ gaussManyUnitD dims
   eps <- uniform (epsMean*0.8) (epsMean*1.2)
   --the initial half momentum step

   let p_half_step = current_p - scale (eps/2) (grad_u current_q)

   let step :: Int -> Vector Double -> Vector Double -> (Vector Double, Vector Double)
       step n p q  -- note that my l is Radford Neal's L+1
        | n == 0 = let qfinal = q + scale eps p -- only q not p update in last loop
                       pfinal = p - scale (eps/2) (grad_u qfinal) -- half momentum step at end
                   in (negate pfinal, -- negate momentum for symmetric proposal
                       qfinal)
        | otherwise =
            let q1 = q + scale eps p
                p1 = p - scale eps (grad_u q1)
            in step (n-1) p1 q1

   let (propose_p, propose_q) = step l p_half_step current_q
       current_U =  u current_q
       current_K =  (current_p `dot` current_p) / 2
       propose_U = u propose_q
       propose_K =  (propose_p `dot` propose_p) / 2
       ratio = exp $ current_U - propose_U + current_K - propose_K
       tr = max 1.0 $ realToFrac count

   u <- unitSample -- 0 to 1
   return $ if  u < ratio
               then HMCPar propose_q l propose_U (if freeze then epsMean else (min 2 $ 1+k_hmc/tr)*epsMean) (count+1) (accept+1) freeze
               else HMCPar current_q l current_U (if freeze then epsMean else (max 0.5 $ 1-k_hmc/tr)**1.8*epsMean) (count+1) (accept) freeze

traceAs s x = trace (s++" = "++show x) x

runHMC postgrad nsam hmc0 = go nsam hmc0 [] where
  go 0 hmcp xs = do io $ putStrLn $ "done with "++show hmcp
                    return (hmcp, xs)
  go n hmcp xs = do hmcnext <- sample $ hmc1 postgrad hmcp
                    if isNaN $ hpPi hmcnext
                       then return (hmcnext, [])
                       else do io $ do putStrLn $ show $ (hpCount hmcnext,
                                                          hpPi hmcnext,
                                                          hpEpsilon hmcnext)
                                       hFlush stdout
                               go (n-1) hmcnext $ hpXi hmcnext : xs

k_hmc = 2


module MachineLearning.MultiSvmClassifierTest
(
  tests
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Approx
import Test.HUnit.Plus

import MachineLearning.DataSets (dataset2)

import qualified Numeric.LinearAlgebra as LA
import qualified MachineLearning as ML
import MachineLearning.Optimization
import MachineLearning.Model
import MachineLearning.MultiSvmClassifier

(x, y) = ML.splitToXY dataset2

model = MultiClass (MultiSvm 1 2)

x1 = ML.addBiasDimension x
onesTheta :: LA.Vector LA.R
onesTheta = LA.konst 1 (2 * LA.cols x1)
zeroTheta :: LA.Vector LA.R
zeroTheta = LA.konst 0 (2 * LA.cols x1)

processX muSigma x = ML.addBiasDimension $ ML.featureNormalization muSigma $ ML.mapFeatures 6 x

muSigma = ML.meanStddev (ML.mapFeatures 6 x)
x2 = processX muSigma x

xPredict = LA.matrix 2 [ -0.5, 0.5
                       , 0.2, -0.2
                       , 1, 1
                       , 1, 0
                       , 0, 0
                       , 0, 1]
xPredict2 = processX muSigma xPredict
yExpected = LA.vector [1, 1, 0, 0, 1, 0]

gradientCheckingEps :: Double
gradientCheckingEps = 1e-3

eps = 0.0001

zeroTheta2 = LA.konst 0 (2 * LA.cols x2)
(thetaGD, _) = minimize (GradientDescent 0.001) model eps 150 (L2 1) x2 y zeroTheta2
(thetaCGFR, _) = minimize (ConjugateGradientFR 0.1 0.1) model eps 30 (L2 0.5) x2 y zeroTheta2
(thetaCGPR, _) = minimize (ConjugateGradientPR 0.1 0.1) model eps 30 (L2 0.5) x2 y zeroTheta2
(thetaBFGS, _) = minimize (BFGS2 0.1 0.1) model eps 30 (L2 0.5) x2 y zeroTheta2

checkGradientTest lambda theta eps = do
  let diffs = take 5 $ map (\e -> checkGradient model lambda x1 y theta e) [1e-3, 1.1e-3 ..]
      diff = minimum diffs
  assertApproxEqual "" eps 0 diff

tests = [  testGroup "gradient checking" [
            testCase "non-zero theta, non-zero lambda" $ checkGradientTest (L2 2) onesTheta 3e-2
              , testCase "zero theta, non-zero lambda" $ checkGradientTest (L2 2) zeroTheta gradientCheckingEps
              , testCase "non-zero theta, zero lambda" $ checkGradientTest (L2 0) onesTheta gradientCheckingEps
              , testCase "zero theta, zero lambda" $ checkGradientTest (L2 0) zeroTheta gradientCheckingEps
              , testCase "non-zero theta, no reg" $ checkGradientTest RegNone onesTheta gradientCheckingEps
              , testCase "zero theta, no reg" $ checkGradientTest RegNone zeroTheta gradientCheckingEps
              ]
        , testGroup "learn" [
            testCase "Gradient Descent" $ assertVector "" 0.01 yExpected (hypothesis model xPredict2 thetaGD)
            , testCase "Conjugate Gradient FR" $ assertVector "" 0.01 yExpected (hypothesis model xPredict2 thetaCGFR)
            , testCase "Conjugate Gradient PR" $ assertVector "" 0.01 yExpected (hypothesis model xPredict2 thetaCGPR)
            , testCase "BFGS" $ assertVector "" 0.01 yExpected (hypothesis model xPredict2 thetaBFGS)
            ]
        ]


{-|
Module      : ODMatrix.ElementaryDecomposition
Description : This module is responsible for the representation of Elementary operations and the process of elementary decomposition of the path between Origin-Destination matrices.

-}
module ODMatrix.ElementaryDecomposition (
    ElementaryMatrix(..)
  , getChildren
  , applyElementary
  , applyPath
  , applicables
  , elementaryValue
  , opposite
  , isNullElementary
  , ODM
  , Bounds
  ) where

  import Numeric.LinearAlgebra
  --import Data.Tree (Tree(..), unfoldTree)

  import ODMatrix (ODM)

  data ElementaryMatrix =
    NullElementary Int |
    Elementary Int (Int, Int) (Int, Int)
    deriving (Show)

  -- | Size of the elementary matrix
  emSize :: ElementaryMatrix -> Int
  emSize (NullElementary s) = s
  emSize (Elementary s _ _) = s

  -- | Bounds of the capacity of the unit (Lower bound, Upper bound).
  type Bounds = (Double, Double)

  -- Size, visited, s
  type ES = (ElementaryMatrix, ODM)

  applyElementaries :: ODM -> [ElementaryMatrix] -> [ODM]
  applyElementaries s es = map (applyElementary s) es

  -- | Apply a elementary operation to the given ODM.
  applyElementary :: ODM -> ElementaryMatrix -> ODM
  applyElementary a e = a + (toMatrix e)

  -- | Apply a sequence of elementary operations to the given ODM.
  applyPath :: ODM -> [ElementaryMatrix] -> ODM
  applyPath = foldl applyElementary

  -- | Given S, such that A + S = B, this function computes the first elementary permutation matrix E such that S' + E = S.
  getChildren :: ODM                  -- ^ Full path permutation matrix
               -> [ElementaryMatrix]  -- ^ All possible elementary matrices
  getChildren odm = es ++ map opposite es
    where al = toAssocList odm
          s = rows odm
          es = [ Elementary s (r1,c1) (r2,c2) |
                  ((r1,c1),_) <- al,
                  ((r2,c2),_) <- al,
                  r1 > r2, c1 < c2 ]

  opposite :: ElementaryMatrix -> ElementaryMatrix
  opposite (Elementary s (r1,c1) (r2,c2)) = Elementary s (r2,c1) (r1,c2)
  opposite (NullElementary s) = NullElementary s

  -- getChildren s = foldl (\acc p -> acc ++ opposites s p pivots) [] pivots
  --   where -- Filtra solo los potenciales pivotes a los que puede agregarse o quitarse un pasajero.
  --         pivots = [ p | p@((i,j),x) <- toAssocList s, i <= j, x /= 0]

  -- reciprocals :: Pivot -> [Pivot] -> [Pivot]
  -- reciprocals ((i,j),x) pivots =
  --   [ p | p@((h,k),y) <- pivots, h < i, j < k, x/y > 0]

  -- opposites :: ODM -> Pivot -> [Pivot] -> [ElementaryMatrix]
  -- opposites s p@(i,j) pivots =
  --   [ Elementary (rows s) (i,j) (h,k)
  --     | ((h,k),_) <- reciprocals p pivots,
  --       atIndex s (h,j) / x < 0,  -- Opposite signs
  --       atIndex s (i,k) / x < 0 ]

  -- | Filter the elementary matrices that cannot be applied to the odm
  applicables :: Bounds             -- ^ Bounds of the values of the target
              -> ODM      -- ^ Target Origin-Destination Matrix
              -> [ElementaryMatrix] -- ^ List of potential elementary operations
              -> [ElementaryMatrix] -- ^ List of applicable elementary operations
  applicables b s = filter (applicable b s)

  -- | Indicates if an elementary operation is applicable for a given origin destination matrix with the given bounds.
  applicable :: Bounds            -- ^ Bounds of the values of the target
             -> ODM               -- ^ Target Origin-Destination Matrix
             -> ElementaryMatrix  -- ^ Elementary operation to apply
             -> Bool
  applicable _ _ (NullElementary _) = True
  applicable (mn,mx) s (Elementary _ (i,j) (h,k)) =
    p (i,k) > mn && p (h,k) < mx &&
    p (i,j) < mx && p (h,j) > mn
    where p = atIndex s

  -- | Compute the value of a elementary matrix in base of a distance functions between cells.
  elementaryValue :: (Int -> Int -> Double) -- ^ Measure function
                  -> ElementaryMatrix       -- ^ E
                  -> Double                 -- ^ Value of E
  elementaryValue _ (NullElementary _) = 0
  elementaryValue m (Elementary _ (r1,_ ) (r2,_)) = s * 2 * m r1 r2
    where s = fromIntegral . signum $ (r1 - r2)

  -- * Convertions

  -- | Transform the internal representation of a Elementary Matrix in a Matrix Double
  toMatrix :: ElementaryMatrix -> ODM
  toMatrix (NullElementary n) = assoc (n,n) 0 []
  toMatrix (Elementary n idx1@(r1,c1) idx2@(r2,c2)) =
    assoc (n,n) 0 [(idx1,1), (idx2,1), ((r1,c2),-1), ((r2,c1),-1)]

  -- | Transform a matrix to an assoc list.
  -- Recovers only the valid, different to 0, elements of the ODM.
  toAssocList :: ODM -> [((Int, Int), Double)]
  toAssocList s = [ c | c@((i,j),v) <- zip idxs elems, i <= j, v /= 0]
    where n = rows s
          idxs = [ (i,j) | i <- [0..(n-1)], j <- [0..(n-1)] ]
          elems = toList . flatten $ s

  -- | Return True if the elementary matrix is a null elementary matrix.
  isNullElementary :: ElementaryMatrix -> Bool
  isNullElementary (NullElementary s) = True
  isNullElementary _ = False


{-# LANGUAGE ScopedTypeVariables #-}
module Main
    ( main
    ) where

-------------------------------------------------------------------------------
import           Data.Foldable             as FT
import           Data.Ratio
import           Data.Sequence             as Seq
import           Test.Tasty
import           Test.Tasty.HUnit
import           Test.Tasty.QuickCheck
-------------------------------------------------------------------------------
import           Statistics.RollingAverage
-------------------------------------------------------------------------------

main :: IO ()
main = defaultMain tests

tests :: TestTree
tests = testGroup "rolling-average"
  [
    testProperty "window size of 0 always produces 0" $ \(samples :: [Int]) ->
      ravg (avgList 0 samples) === (0.0 :: Double)
  , testProperty "never grows samples larger than the limit" $ \(Positive lim) (samples :: [Int]) ->
      Seq.length (ravgSamples (avgList lim samples)) <= lim
  , testCase "can calculate rolling average as a Ratio" $ do
      ravg (avgList 3 ([0, 1, 2, 3, 4, 5] :: [Int])) @?= (((3 + 4 + 5) % 3) :: Rational)
  ]

-------------------------------------------------------------------------------
avgList :: Num a => Int -> [a] -> RollingAvg a
avgList lim samples = FT.foldl' ravgAdd (mkRavg lim) samples


-- Modulo destinado a todas as funcoes relativas ao treino da Rede.

module Training
(train) where

import Execution
import InputOutput
import Types
import Data.List.Split
import System.Random
import System.Random.Shuffle
import Numeric.LinearAlgebra.HMatrix
import Numeric.LinearAlgebra.Data

type Image = [Double]
type Sample = (Int, Image)

-- Funcao de treino que executa uma
-- dada quantidade de epocas de treino
train :: Int -> IO String
train epochAmount = do
                    trainingSet <- getTraining
                    testSet <- getTest
                    network <- readIn
                    manageTrainingEpoch epochAmount trainingSet testSet network

--
manageTrainingEpoch :: Int -> [Sample] -> [Sample] -> Data -> IO String
manageTrainingEpoch 0 _ _ _ = return ""
manageTrainingEpoch epochAmount trainingSet testSet network = do
                                                newNetwork <- trainingEpoch trainingSet network
                                                let
                                                    correctCnt = testEpoch testSet newNetwork
                                                    totalAmount = length testSet
                                                do  return $ printEpoch epochAmount correctCnt totalAmount
                                                    writeIn newNetwork
                                                    manageTrainingEpoch (epochAmount - 1) trainingSet testSet newNetwork

trainingEpoch :: [Sample] -> Data -> IO Data
trainingEpoch trainingSet network = do
                                standardGenerator <- getStdGen
                                let trainingSize = length trainingSet
                                    shuffledTrainingSet = shuffle' trainingSet trainingSize standardGenerator
                                    minibatchAmount = 20 -- adequar quantidade
                                    minibatchSize = trainingSize `div` 20 -- adequar quantidade
                                    minibatches = chunksOf minibatchSize shuffledTrainingSet
                                    newNetwork = manageMinibatch minibatchAmount 0 network minibatches
                                return newNetwork

manageMinibatch :: Int -> Int -> Data -> [[Sample]] -> Data
manageMinibatch amount counter network minibatches = if amount /= (counter + 1)
                                                        then
                                                          let changes = minibatchEvaluation (minibatches !! counter) amount network
                                                          in plus (plus network changes) (manageMinibatch amount (counter + 1) network minibatches)
                                                        else
                                                          generateBasedOf network

minibatchEvaluation :: [Sample] -> Int -> Data -> Data
minibatchEvaluation minibatch amount network = let sumedDesiredChanges = manageSample minibatch amount network
                                                   averageDesiredChanges = divide sumedDesiredChanges (fromIntegral amount)
                                               in averageDesiredChanges

manageSample :: [Sample] -> Int -> Data -> Data
manageSample minibatch counter networkModel = if counter > 0
                                    then
                                        let representedInt = fst $ minibatch !! counter
                                            image = snd $ minibatch !! counter
                                            network = feedforward image networkModel
                                            expectedOutput = buildExpectedOutput representedInt
                                            desiredChanges = backpropagation network image expectedOutput
                                            sumChanges = generateBasedOf networkModel
                                        in plus (plus sumChanges desiredChanges) (manageSample minibatch counter networkModel)
                                    else
                                        generateBasedOf networkModel

buildExpectedOutput :: Int -> [Double]
buildExpectedOutput representedValue = let indexes = [0.0 .. 9.0]
                                       in [if x == (fromIntegral representedValue) then x else 0.0 | x <- indexes]

-- Recebe as informacoes da rede neural, o resultado
-- esperado e retorna um Data com as modificacoes necessarias
-- na rede
-- N = network, E = expected list, I = image values
backpropagation :: Data -> Image -> [Double] -> Data
backpropagation n i e
    | isEmpty n = error "Backpropagation error: Data is empty"
    | length e /= 10 = error "Backpropagation error: expectedOutput list is invalid"
    | otherwise = let outputError = computeOutputError (aOutput n) (fromList e) (zetaOutput n)
                      hiddenError = computeHiddenError (wOutput n) outputError (zetaHidden n)
                      outputDesiredChanges = computeOutputDesiredChanges outputError (aHidden n)
                      hiddenDesiredChanges = computeHiddenDesiredChanges hiddenError i
                      wH = fst hiddenDesiredChanges
                      bH = snd hiddenDesiredChanges
                      wO = fst outputDesiredChanges
                      bO = snd outputDesiredChanges
                      newNetwork = (Data wH bH (aHidden n) (zetaHidden n) wO bO (aOutput n) (zetaOutput n))
                  in newNetwork

-- Gera o vetor de erro da camada output, recebe
-- as ativacoes do output atual, as ativacoes esperadas
-- e a lista zeta do output
computeOutputError :: Vector Double -> Vector Double -> Vector Double -> Vector Double
computeOutputError a e z = fromList $ hadamardV (toList (add a (scale (-1) e))) (sigV' $ toList z)

computeHiddenError :: Matrix R -> Vector Double -> Vector Double -> Vector Double
computeHiddenError ow oe hz = fromList $ hadamardV (toList $ (tr' ow) #> oe) (sigV' (toList hz))

computeOutputDesiredChanges :: Vector Double -> Vector Double -> (Matrix R, Vector Double)
computeOutputDesiredChanges oe ah = let owDesired = oe `outer` ah
                                        ohDesired = oe
                                    in (owDesired, ohDesired)

computeHiddenDesiredChanges :: Vector Double -> Image -> (Matrix R, Vector Double)
computeHiddenDesiredChanges he image = let hwDesired = he `outer` (fromList image)
                                           hbDesired = he
                                       in (hwDesired, hbDesired)

testEpoch :: [Sample] -> Data -> Int
testEpoch testSet network = manageEpoch testSet network (length testSet)

manageEpoch :: [Sample] -> Data -> Int -> Int
manageEpoch testSet network counter = if counter > 0
                                        then
                                            let newNetwork = feedforward (snd $ testSet !! counter) network
                                            in if (toList $ aOutput newNetwork) == buildExpectedOutput (fst $ testSet !! counter)
                                                then
                                                    1 + manageEpoch testSet network (counter - 1)
                                                else
                                                    0 + manageEpoch testSet network (counter - 1)
                                        else
                                            0


{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies  #-}

-- | Provides spherical harmonic models of scalar-valued functions.
module Math.SphericalHarmonics
(
  SphericalHarmonicModel
, sphericalHarmonicModel
, scaledSphericalHarmonicModel
, evaluateModel
, evaluateModelCartesian
, evaluateModelGradient
, evaluateModelGradientCartesian
, evaluateModelGradientInLocalTangentPlane
)
where

import Data.Complex
import Data.VectorSpace hiding (magnitude)
import Math.SphericalHarmonics.AssociatedLegendre
import Numeric.AD

-- | Represents a spherical harmonic model of a scalar-valued function.
data SphericalHarmonicModel a = SphericalHarmonicModel [[(a, a)]]
  deriving (Functor)

-- | Creates a spherical harmonic model.
-- Result in an error if the length of the list is not a triangular number.
sphericalHarmonicModel :: (Fractional a) => [[(a, a)]] -- ^ A list of g and h coefficients for the model
                       -> SphericalHarmonicModel a -- ^ The spherical harmonic model
sphericalHarmonicModel cs | valid = SphericalHarmonicModel cs
                          | otherwise = error "The number of coefficients is not a triangular number."
  where
    valid = and $ zipWith (==) (fmap length cs) [1..length cs]

-- | Creates a spherical harmonic model, scaling coefficients for the supplied reference radius.
-- Result in an error if the length of the list is not a triangular number.
scaledSphericalHarmonicModel :: (Fractional a) => a -- ^ The reference radius
                       -> [[(a, a)]] -- ^ A list of g and h coefficients for the model
                       -> SphericalHarmonicModel a -- ^ The spherical harmonic model
scaledSphericalHarmonicModel r cs = sphericalHarmonicModel cs'
  where
    cs' = normalizeReferenceRadius r cs

instance(Fractional a, Eq a) => AdditiveGroup (SphericalHarmonicModel a) where
  zeroV = SphericalHarmonicModel [[(0,0)]]
  negateV = fmap negate
  (SphericalHarmonicModel m1) ^+^ (SphericalHarmonicModel m2) = SphericalHarmonicModel (combineCoefficients m1 m2)
    where
      combineCoefficients [] cs = cs
      combineCoefficients cs [] = cs
      combineCoefficients (c1:cs1) (c2:cs2) = zipWith addPairs c1 c2 : combineCoefficients cs1 cs2
      addPairs (g1, h1) (g2, h2) = (g1 + g2, h1 + h2)

instance (Fractional a, Eq a) => VectorSpace (SphericalHarmonicModel a) where
  type Scalar (SphericalHarmonicModel a) = a
  x *^ m = fmap (* x) m

normalizeReferenceRadius :: (Fractional a) => a -> [[(a, a)]] -> [[(a, a)]]
normalizeReferenceRadius r = zipWith (fmap . mapWholePair . transform) [0 :: Int ..]
  where
    transform n = (* (r ^ (2 + n)))

-- | Computes the scalar value of the spherical harmonic model at a specified spherical position.
evaluateModel :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
              -> a -- ^ Spherical radius
              -> a -- ^ Spherical colatitude (radian)
              -> a -- ^ Spherical longitude (radian)
              -> a -- ^ Model value
evaluateModel m r colat lon = evaluateModel' m r (cos colat) (cis lon)

-- | Computes the scalar value of the spherical harmonic model at a specified Cartesian position.
evaluateModelCartesian :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                       -> a -- ^ X position
                       -> a -- ^ Y position
                       -> a -- ^ Z position
                       -> a -- ^ Model value
evaluateModelCartesian m x y z = evaluateModel' m r cosColat cisLon
  where
    r = sqrt $ (x*x) + (y*y) + (z*z)
    cosColat = z / r
    cisLon = normalize $ mkPolar x y

evaluateModel' :: (RealFloat a, Ord a) => SphericalHarmonicModel a
               -> a -- r
               -> a -- cosColat
               -> Complex a -- cisLon
               -> a
evaluateModel' (SphericalHarmonicModel cs) r cosColat cisLon = sum $ zipWith (*) (iterate (/ r) (recip r)) (zipWith evaluateDegree [0..] cs)
  where
    sines = 1 : iterate (* cisLon) cisLon
    evaluateDegree n cs' = sum $ zipWith3 evaluateOrder (fmap (schmidtSemiNormalizedAssociatedLegendreFunction n) [0..n]) cs' sines
    evaluateOrder p (g, h) cisMLon = ((g * realPart cisMLon) + (h * imagPart cisMLon)) * (p (cosColat))

-- | Computes the gradient of the scalar value of the spherical harmonic model, in spherical coordinates, at a specified location.
evaluateModelGradient :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                      -> a -- ^ Spherical radius
                      -> a -- ^ Spherical colatitude (radian)
                      -> a -- ^ Spherical longitude (radian)
                      -> (a, a, a) -- ^ Radial, colatitudinal, and longitudinal components of gradient
evaluateModelGradient model r colat lon = makeTuple . fmap negate $ modelGrad [r, colat, lon]
  where
    modelGrad = grad (\[r', c', l'] -> evaluateModel (fmap auto model) r' c' l')

-- | Computes the gradient of the scalar value of the spherical harmonic model at a specified location, in Cartesian coordinates.
-- The result is expressed in right-handed coordinates centered at the origin of the sphere, with the positive Z-axis piercing the
-- north pole and the positive x-axis piercing the reference meridian.
evaluateModelGradientCartesian :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                               -> a -- ^ X position
                               -> a -- ^ Y position
                               -> a -- ^ Z position
                               -> (a, a, a) -- X, Y, and Z components of gradient
evaluateModelGradientCartesian model x y z = makeTuple . fmap negate $ modelGrad [x, y, z]
  where
    modelGrad = grad (\[x', y', z'] -> evaluateModelCartesian (fmap auto model) x' y' z')

-- | Computes the gradient of the scalar value of the spherical harmonic model at a specified location, in Cartesian coordinates.
-- The result is expressed in a reference frame locally tangent to the sphere at the specified location.
evaluateModelGradientInLocalTangentPlane :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                                         -> a -- ^ Spherical radius
                                         -> a -- ^ Spherical colatitude (radian)
                                         -> a -- ^ Spherical longitude (radian)
                                         -> (a, a, a) -- ^ East, North, and up components of gradient
evaluateModelGradientInLocalTangentPlane model r colat lon = (e, n, u)
  where
    (r', colat', lon') = evaluateModelGradient model r colat lon
    e = lon' / (r * sin colat)
    n = -colat' / r -- negated because the colatitude increase southward
    u = r'

normalize :: (RealFloat a) => Complex a -> Complex a
normalize r@(x :+ y) | isInfinite m' = 0
                     | otherwise = (x * m') :+ (y * m')
  where
    m' = recip . magnitude $ r

mapWholePair :: (a -> b) -> (a, a) -> (b, b)
mapWholePair f (a, b) = (f a, f b)

makeTuple :: [a] -> (a, a, a)
makeTuple [x, y, z] = (x, y, z)


module Main where

import Data.Monoid

import Core
import Render
import Codec.Picture
import Codec.Picture.Types
import Numeric.LinearAlgebra(vector)

settings :: RenderSettings
settings = RenderSettings {
    background = PixelRGBF 1 1 1,
    width = 800,
    height = 800,
    path = "out.png",
    topLeft = vector [1,1,0],
    topRight = vector [1,-1,0],
    bottomRight = vector [-1, -1,0],
    origin = vector [0,0,-1],
    antialiasing = 10}

scene
    = plane (diffuse (scaleColor 0.5)) (Ray (vec3 (-2) 0 0) (vec3 (-2) 0 0))
    <> sphere 0.9 (vec3 0 0 2) (diffuse (scaleColor 0.5))
    <> plane (skylike white blue) (Ray (vec3 10 0 0) (vec3 1 0 0))

main :: IO ()
main = do
    img <- render settings scene
    savePngImage (path settings) img