{- |
Raw, unnormalized versions of the transforms in @fftw@.

Note that the forwards and backwards transforms of this module are not actually
inverses.  For example, @run idft (run dft v) /= v@ in general.

For more information on the individual transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.
-}
module Numeric.FFT.Vector.Unnormalized(
                    -- * Creating and executing 'Plan's
                    run,
                    plan,
                    execute,
                    -- * Complex-to-complex transforms
                    dft,
                    idft,
                    -- * Real-to-complex transforms
                    dftR2C,
                    dftC2R,
                    -- * Real-to-real transforms
                    -- $dct_size
                    -- ** Discrete cosine transforms
                    dct1,
                    dct2,
                    dct3,
                    dct4,
                    -- ** Discrete sine transforms
                    dst1,
                    dst2,
                    dst3,
                    dst4,
                    ) where

import Numeric.FFT.Vector.Base
import Foreign
import Foreign.C
import Data.Complex

#include <fftw3.h>

-- | Whether the complex fft is forwards or backwards.
type CDirection = CInt

-- | The type of the cosine or sine transform.
type CKind = (#type fftw_r2r_kind)

foreign import ccall unsafe fftw_plan_dft_1d
    :: CInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> CDirection
        -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_dft_r2c_1d
    :: CInt -> Ptr Double -> Ptr (Complex Double) -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_dft_c2r_1d
    :: CInt -> Ptr (Complex Double) -> Ptr Double -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_r2r_1d
    :: CInt -> Ptr Double -> Ptr Double -> CKind -> CFlags -> IO (Ptr CPlan)

dft1D :: CDirection -> Transform (Complex Double) (Complex Double)
dft1D d = Transform {
            inputSize = id,
            outputSize = id,
            creationSizeFromInput = id,
            makePlan = \n a b -> withPlanner . fftw_plan_dft_1d n a b d,
            normalization = const id
            }

-- | A forward discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = sum_(j=0)^(n-1) x_j e^(-2pi i j k/n)@
dft :: Transform (Complex Double) (Complex Double)
dft = dft1D (#const FFTW_FORWARD)

-- | A backward discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = sum_(j=0)^(n-1) x_j e^(2pi i j k/n)@
idft :: Transform (Complex Double) (Complex Double)
idft = dft1D (#const FFTW_BACKWARD)

-- | A forward discrete Fourier transform with real data.  If the input size is @n@,
-- the output size will be @n \`div\` 2 + 1@.
dftR2C :: Transform Double (Complex Double)
dftR2C = Transform {
            inputSize = id,
            outputSize = \n -> n `div` 2 + 1,
            creationSizeFromInput = id,
            makePlan = \n a b -> withPlanner . fftw_plan_dft_r2c_1d n a b,
            normalization = const id
        }

-- | A backward discrete Fourier transform which produces real data.
--
-- This 'Transform' behaves differently than the others:
--
--  - Calling @plan dftC2R n@ creates a 'Plan' whose /output/ size is @n@, and whose
--    /input/ size is @n \`div\` 2 + 1@.
--
--  - If @length v == n@, then @length (run dftC2R v) == 2*(n-1)@.
dftC2R :: Transform (Complex Double) Double
dftC2R = Transform {
            inputSize = \n -> n `div` 2 + 1,
            outputSize = id,
            creationSizeFromInput = \n -> 2 * (n-1),
            makePlan = \n a b -> withPlanner . fftw_plan_dft_c2r_1d n a b,
            normalization = const id
        }

r2rTransform :: CKind -> Transform Double Double
r2rTransform kind = Transform {
                    inputSize = id,
                    outputSize = id,
                    creationSizeFromInput = id,
                    makePlan = \n a b -> withPlanner . fftw_plan_r2r_1d n a b kind,
                    normalization = const id
                }

-- $dct_size
-- The real-even (DCT) and real-odd (DST) transforms.  The input and output sizes
-- are the same (@n@).

-- | A type-1 discrete cosine transform.
--
-- @y_k = x_0 + (-1)^k x_(n-1) + 2 sum_(j=1)^(n-2) x_j cos(pi j k\/(n-1))@
dct1 :: Transform Double Double
dct1 = r2rTransform (#const  FFTW_REDFT00)

-- | A type-2 discrete cosine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)k\/n)@
dct2 :: Transform Double Double
dct2 = r2rTransform (#const  FFTW_REDFT10)

-- | A type-3 discrete cosine transform.
--
-- @y_k = x_0 + 2 sum_(j=1)^(n-1) x_j cos(pi j(k+1\/2)\/n)@
dct3 :: Transform Double Double
dct3 = r2rTransform (#const  FFTW_REDFT01)

-- | A type-4 discrete cosine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)(k+1\/2)\/n)@
dct4 :: Transform Double Double
dct4 = r2rTransform (#const  FFTW_REDFT11)

-- | A type-1 discrete sine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j sin(pi(j+1)(k+1)\/(n+1))@
dst1 :: Transform Double Double
dst1 = r2rTransform (#const  FFTW_RODFT00)

-- | A type-2 discrete sine transform.
--
-- @y_k = 2 sum_(j=0)^(n-1) x_j sin(pi(j+1\/2)(k+1)\/n)@
dst2 :: Transform Double Double
dst2 = r2rTransform (#const  FFTW_RODFT10)

-- | A type-3 discrete sine transform.
--
-- @y_k = (-1)^k x_(n-1) + 2 sum_(j=0)^(n-2) x_j sin(pi(j+1)(k+1\/2)/n)@
dst3 :: Transform Double Double
dst3 = r2rTransform (#const  FFTW_RODFT01)

-- | A type-4 discrete sine transform.
--
-- @y_k = sum_(j=0)^(n-1) x_j sin(pi(j+1\/2)(k+1\/2)\/n)@
dst4 :: Transform Double Double
dst4 = r2rTransform (#const FFTW_RODFT11)


{-# LANGUAGE GADTs
           , DataKinds
           , TypeApplications
           , ScopedTypeVariables
           , ExplicitForAll
           , AllowAmbiguousTypes
           , KindSignatures
           , TypeOperators
           , RankNTypes
           , GeneralizedNewtypeDeriving
           , StrictData #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

module Lib where

import Numeric.LinearAlgebra hiding (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA
import qualified Data.Vector.Storable as VS
import GHC.TypeLits
import Data.Proxy
import System.Random
import Control.Monad
import Numeric.AD

newtype Vector (n :: Nat) = Vector { _vec :: VS.Vector Float }
  deriving (Show, Floating)

instance KnownNat n => Num (Vector n) where
  (Vector v1) + (Vector v2) = Vector $ v1 + v2
  (Vector v1) - (Vector v2) = Vector $ v1 - v2
  (Vector v1) * (Vector v2) = Vector $ v1 * v2
  abs                       = vectorMap abs
  fromInteger               = vecFromList . (:[]) . fromInteger
  negate                    = vectorMap negate
  signum                    = vectorMap signum

instance KnownNat n => Fractional (Vector n) where
  (Vector v1) / (Vector v2) = Vector $ v1 / v2
  fromRational              = vecFromList . (:[]) . fromRational

natToInt :: forall n. KnownNat n => Int
natToInt = fromIntegral . natVal $ Proxy @n

vecFromList :: forall n. KnownNat n => [Float] -> Vector n
vecFromList = Vector . VS.fromList . take (natToInt @n) . cycle

vecLength :: forall n. KnownNat n => Vector n -> Int
vecLength _ = natToInt @n

vecTail :: forall n. KnownNat n => Vector (n+1) -> Vector n
vecTail (Vector v) = Vector $ VS.tail v

newtype Matrix (n :: Nat) (m :: Nat) = Matrix { _mat :: LA.Matrix Float }
  deriving Show

matrixWidth :: forall n m. (KnownNat n, KnownNat m) => Matrix n m -> Int
matrixWidth _ = natToInt @n

matrixHeight :: forall n m. (KnownNat n, KnownNat m) => Matrix n m -> Int
matrixHeight _ = natToInt @m

matFromList :: forall n m. (KnownNat n, KnownNat m) => [Float] -> Matrix n m
matFromList = Matrix . (h><w) . cycle
  where w = natToInt @n
        h = natToInt @m

randomMatrix :: forall n m. (KnownNat n, KnownNat m) => IO (Matrix n m)
randomMatrix = matFromList <$> replicateM (w*h) (randomRIO (0.001, 0.1))
    where w = natToInt @n
          h = natToInt @m

data Network :: Nat -> Nat -> * where
  Layer :: forall m n. (KnownNat n, KnownNat m) =>
                       Matrix (n+1) m ->
                       (forall a. Floating a => a -> a) ->
                       Network n m
  Composition :: (KnownNat n, KnownNat m, KnownNat k) =>
                 Network m k ->
                 Network n m ->
                 Network n k

conductSignal :: forall n m. (KnownNat n, KnownNat m) =>
                               Network n m ->
                               Vector n ->
                               Vector m
conductSignal (Layer mat f) = vectorMap f . mvMul mat . vectorCons 1
conductSignal (Composition network2 network1) = conductSignal network2 . conductSignal network1

vectorMap :: (Float -> Float) -> Vector n -> Vector n
vectorMap f (Vector v) = Vector $ VS.map f v

mvMul :: Matrix n m -> Vector n -> Vector m
mvMul (Matrix m) (Vector v) = Vector $ m #> v

vectorCons :: KnownNat n => Float -> Vector n -> Vector (n+1)
vectorCons n (Vector v) = Vector $ VS.cons n v

randomLayer :: forall n m. (KnownNat n, KnownNat m) => (forall a. Floating a => a -> a) -> IO (Network n m)
randomLayer f = (`Layer` f) <$> randomMatrix

scaleMatrix :: (KnownNat n, KnownNat m) => Float -> Matrix m n -> Matrix m n
scaleMatrix s (Matrix m) = Matrix $ LA.scale s m

transposeMat :: (KnownNat n, KnownNat m) => Matrix n m -> Matrix m n
transposeMat (Matrix mat) = Matrix $ LA.tr mat

outerProduct :: (KnownNat n, KnownNat m) => Vector n -> Vector m -> Matrix n m
outerProduct (Vector v1) (Vector v2) = Matrix $ v2 `outer` v1

vvMul :: (KnownNat n) => Vector n -> Vector n -> Vector n
vvMul (Vector v1) (Vector v2) = Vector $ v1 * v2

matrixSub :: (KnownNat m, KnownNat n) => Matrix n m -> Matrix n m -> Matrix n m
matrixSub (Matrix m1) (Matrix m2) = Matrix $ m1 - m2

useTrainingExample :: forall m n. (KnownNat n, KnownNat m) =>
                        Vector m ->
                        Network n m ->
                        Vector n ->
                        (Network n m, Vector n)
useTrainingExample diffLoss (Layer weights activation) input = (trainedLayer, inputGradient)
  where
    trainedLayer = Layer newWeights activation
    newWeights :: Matrix (n+1) m
    newWeights = matrixSub weights $ scaleMatrix 0.01 weightGradient
    consedInput = vectorCons 1 input
    preActivationOutput = mvMul weights consedInput
    output = vectorMap activation preActivationOutput
    diffActivation = vectorMap (diff activation) preActivationOutput
    weightGradient = outerProduct consedInput (diffLoss * diffActivation)
    inputGradient  = vecTail $ mvMul (transposeMat weights) (diffLoss * diffActivation)
useTrainingExample diffLoss (Composition network2 network1) input =
    (Composition trainedNet2 trainedNet1, inputGradient)
  where (trainedNet2,outputGrad1)   = useTrainingExample diffLoss network2 input2
        input2                      = conductSignal network1 input
        (trainedNet1,inputGradient) = useTrainingExample outputGrad1 network1 input

showExample :: (KnownNat n, KnownNat m) =>
                (Vector m -> Vector m -> Vector m) ->
                Network n m ->
                (Vector n, Vector m) ->
                Network n m
showExample lossGrad network (input,expected) = fst $ useTrainingExample diffLoss network input
  where diffLoss = lossGrad expected output
        output   = conductSignal network input

sigmoid input = 1/(1+exp(-input))
                -- (1/) . (1+) . exp . negate

epsilon :: Floating a => a
epsilon = 0.0000001

logLoss :: Floating a => [a] -> a
logLoss [expected, output] = negate $ expected * log (output + epsilon) + (1 - expected) * log (1 - output + epsilon)

gradLogLoss expected output = grad logLoss [expected, output] !! 1

-- "mnist/train-set"

getTrainSet :: IO [(Vector 784, Vector 10)]
getTrainSet =
  map (\(a,b) -> (vecFromList a, toOneHot b)) .
  map (read @([Float],Int)) .
  lines <$> readFile  "mnist/train-set"
  where
    toOneHot n = vecFromList $ map (\x -> if x == n then 1 else 0) [0..]

eval network (input,expected) = do
  putStrLn "EXPECTED:"
  print expected
  putStrLn "PREDICTED:"
  print $ conductSignal network input

trainTest k n = do
  trainSet <- take n <$> getTrainSet
  untrainedLayer1 <- randomLayer sigmoid
  untrainedLayer2 <- randomLayer @50 @10 sigmoid
  let untrainedNetwork = Composition untrainedLayer2 untrainedLayer1
      trainedNetwork = foldl (showExample gradLogLoss) untrainedNetwork . take (k*n) $ cycle trainSet
      evalSet = take 10 trainSet
  mapM_ (eval trainedNetwork) evalSet

{-

  diff (loss . activation . mvMul(weights, input))

   => diff loss (activation . mvMul (weights,input)) * diff (activation . mvMul (weights,input))

   => diff loss (output) * diff activation (mvMul (weights,input)) * diff mvMul(weights,input)
-}


module STCR2Z1T0 where

import           Control.Monad             as M
import           Control.Monad.Parallel    as MP
import           Data.Array.Repa           as R
import           Data.Binary               (decodeFile)
import           Data.Complex
import           Data.List                 as L
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:initDistStr:histFilePath:alphaStr:pinwheelFlagStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      init = read initStr :: (Double, Double, Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      initDist = read initDistStr :: [R2S1RPPoint]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numThread = read numThreadStr :: Int
      sourceDist = L.take 1 initDist
      sinkDist = L.drop 1 initDist
      folderPath = "output/test/STCR2Z1T0"
  flag <- doesFileExist histFilePath
  -- arrR2Z1T0 <-
  --   if pinwheelFlag
  --     then computeR2Z1T0Array numPoint numPoint alpha thetaFreqs theta0Freqs
  --     else if flag
  --            then getNormalizedHistogramArr <$> decodeFile histFilePath
  --            else solveMonteCarloR2Z1T0
  --                   numThread
  --                   numTrail
  --                   maxTrail
  --                   numPoint
  --                   numPoint
  --                   sigma
  --                   tao
  --                   1
  --                   len
  --                   theta0Freqs
  --                   thetaFreqs
  --                   histFilePath
  --                   (emptyHistogram
  --                      [ numPoint
  --                      , numPoint
  --                      , L.length theta0Freqs
  --                      , L.length thetaFreqs
  --                      ]
  --                      0)
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z1T0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          sigma
          tao
          0
          theta0Freqs
          thetaFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length theta0Freqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z1T0 <-
    computeUnboxedP $
    computeR2Z1T0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      thetaFreqs
      theta0Freqs
  let arr3d =
        rotate3D . R.slice arrR2Z1T0 $
        (Z :. All :. (L.length theta0Freqs - 1) :. All :. All)
  createDirectoryIfMissing True folderPath
  MP.mapM_
    (\i ->
       plotImageRepaComplex
         (folderPath </> "GreensR2Z1T0_" L.++ show (i + 1) L.++ ".png") .
       ImageRepa 8 .
       computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr3d $
       (Z :. All :. All :. i))
    [0 .. (L.length thetaFreqs) - 1]
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  sourceDistArr <-
    computeInitialDistributionR2T0 plan numPoint numPoint theta0Freqs sourceDist
  sinkDistArr <-
    computeInitialDistributionR2T0 plan numPoint numPoint theta0Freqs sinkDist
  arrR2Z1T0F <- (computeP . makeFilter2D $ arrR2Z1T0) >>= dftR2Z1T0 plan
  arrR2Z1T0TRF <-
    (computeP . makeFilter2D . timeReverseR2Z1T0 thetaFreqs theta0Freqs $
     arrR2Z1T0) >>=
    dftR2Z1T0 plan
  -- Source field
  sourceArr <- convolveR2T0 plan arrR2Z1T0F sourceDistArr
  -- let sourceR2Z1 = R.sumS . rotateR2Z1T0Array $ sourceArr
  --     sourceField =
  --       computeS .
  --       R.extend (Z :. (1 :: Int) :. All :. All) .
  --       R.sumS . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
  --       sourceR2Z1
  sourceR2Z1 <- R.sumP . rotateR2Z1T0Array $ sourceArr
  sourceField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
    sourceR2Z1
  plotImageRepaComplex (folderPath </> "Source.png") . ImageRepa 8 $ sourceField
  -- Sink field
  sinkArr' <- convolveR2T0 plan arrR2Z1T0F sinkDistArr
  let   sinkArr  = computeSinkFromSourceR2Z1T0 thetaFreqs theta0Freqs sinkArr'
  -- let sinkR2Z1 = R.sumS . rotateR2Z1T0Array $ sinkArr
  --     sinkField =
  --       computeS .
  --       R.extend (Z :. (1 :: Int) :. All :. All) .
  --       R.sumS . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
  --       sinkR2Z1
  sinkR2Z1 <- R.sumP . rotateR2Z1T0Array $ sinkArr
  sinkField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
    sinkR2Z1
  plotImageRepaComplex (folderPath </> "Sink.png") . ImageRepa 8 $ sinkField
  -- Completion Filed
  completionFiled <-
    convolveR2Z1 plan thetaFreqs sourceR2Z1 sinkR2Z1
  let completionFiledR2 =
        R.sumS . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
        completionFiled
  -- completionFiledR2 <-
  --   R.sumP . R.map magnitude . rotate3D . r2z1Tor2s1 numOrientation thetaFreqs $
  --   completionFiled
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.map magnitude $
    completionFiledR2


module Main where

import Network.Layer
import Network.Network
import Network.Neuron
import Network.Trainer

import Numeric.LinearAlgebra
import System.Random

main :: IO ()
main = do

  g <- newStdGen
  let l   = LayerDefinition sigmoidNeuron 2 connectFully
      l'  = LayerDefinition sigmoidNeuron 2 connectFully
      l'' = LayerDefinition sigmoidNeuron 1 connectFully

  let n = createNetwork normals g [l, l', l'']

  let t = BackpropTrainer 3 quadraticCost quadraticCost'

  let dat = [(fromList [0, 1], fromList [1]),
             (fromList [1, 1], fromList [0]),
             (fromList [1, 0], fromList [1]),
             (fromList [0, 0], fromList [0])]

  let n' = trainNTimes n t online dat 10000

  putStrLn "==> XOR predictions: "
  putStr "(0,0): "
  print $ predict (fromList [0, 0]) n'
  putStr "(1,0): "
  print $ predict (fromList [1, 0]) n'
  putStr "(0,1): "
  print $ predict (fromList [0, 1]) n'
  putStr "(1,1): "
  print $ predict (fromList [1, 1]) n'

  putStrLn "\n==> Trained network: "
  print $ map layerToShowable $ layers n'

{-
  saveNetwork "xor.ann" n'

  putStrLn "==> Network saved and reloaded: "
  n'' <- loadNetwork "xor.ann" [l, l', l'']

  print $ predict (fromList [0, 0]) n''
  print $ predict (fromList [1, 0]) n''
  print $ predict (fromList [0, 1]) n''
  print $ predict (fromList [1, 1]) n''
-}


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.Color.YUV (
  YUV(..)
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C
import Data.Vector.Image.Color.RGB

data YUV a = YUV {
   yuv_y :: a
 , yuv_u :: a
 , yuv_b :: a
} deriving (Show,Eq,Ord)

instance Functor YUV where
  fmap func (YUV r g b) = YUV (func r) (func g) (func b)

rgb2yuv :: (Num a,Integral a) => RGB a -> YUV a
rgb2yuv (RGB r g b) = YUV ((r+2*g+b)`div`4) (b-g) (r-g)

yuv2rgb :: (Num a,Integral a) => YUV a -> RGB a
yuv2rgb (YUV y u v) = RGB (v+g) g (u+g)
  where
    g = y - ((u + v)`div`4)

instance (Num a,Integral a) => ColorSpace YUV a where
  toRGB = yuv2rgb
  fromRGB = rgb2yuv
  toList (YUV r g b) = [r,g,b]
  fromList [r,g,b] = (YUV r g b)
  fromList val = error $ "List length must be 3 but " L.++  show (L.length val)
  dim _ = 3


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE PatternSynonyms     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# LANGUAGE ViewPatterns        #-}
module Numeric.Quaternion.Internal
    ( Quaternion (..), Quater(Quater)
    ) where

import Numeric.DataFrame.Type  (KnownBackend)
import Numeric.Matrix.Internal (Matrix)
import Numeric.PrimBytes       (PrimBytes)
import Numeric.Vector.Internal (Vector)
import Text.Read

-- | @(x,y,z,w)@ of a quaternion
pattern Quater :: Quaternion t => t -> t -> t -> t -> Quater t
pattern Quater a b c d <- (unpackQ# -> (# a, b, c, d #))
  where
    Quater = packQ
{-# COMPLETE Quater #-}

-- | Quaternion operations
class ( Floating (Quater t), Floating t, Ord t, PrimBytes t
      , KnownBackend t '[3], KnownBackend t '[4]
      , KnownBackend t '[3,3], KnownBackend t '[4, 4]
      ) => Quaternion t where
    -- | Quaternion data type. The ordering of coordinates is @(x,y,z,w)@,
    --   where @w@ is the argument, and @x y z@ are the components of a 3D vector
    data Quater t
    -- | Set the quaternion in format @(x,y,z,w)@
    packQ :: t -> t -> t -> t -> Quater t
    -- | Get the values of the quaternion in format @(x,y,z,w)@
    unpackQ# :: Quater t -> (# t, t, t, t #)
    -- | Set the quaternion from 3D axis vector and argument
    fromVecNum :: Vector t 3 -> t -> Quater t
    -- | Set the quaternion from 4D vector in format @(x,y,z,w)@
    fromVec4 :: Vector t 4 -> Quater t
    -- | Transform the quaternion to 4D vector in format @(x,y,z,w)@
    toVec4 :: Quater t -> Vector t 4
    -- | Get scalar square of the quaternion.
    --
    --   >>> realToFrac (square q) == q * conjugate q
    square :: Quater t -> t
    -- | Imaginary part of the quaternion (orientation vector)
    im :: Quater t -> Quater t
    -- | Real part of the quaternion
    re :: Quater t -> Quater t
    -- | Imaginary part of the quaternion as a 3D vector
    imVec :: Quater t -> Vector t 3
    -- | Real part of the quaternion as a scalar
    taker :: Quater t -> t
    -- | i-th component
    takei :: Quater t -> t
    -- | j-th component
    takej :: Quater t -> t
    -- | k-th component
    takek :: Quater t -> t
    -- | Conjugate quaternion (negate imaginary part)
    conjugate :: Quater t -> Quater t
    -- | Rotates and scales vector in 3D using quaternion.
    --   Let \( q = c (\cos \frac{\alpha}{2}, v \sin \frac{\alpha}{2}) \)
    --     , \( c > 0 \), \( {|v|}^2 = 1 \);
    --   then the rotation angle is \( \alpha \), and the axis of rotation is \(v\).
    --   Scaling is proportional to \( c^2 \).
    --
    --   >>> rotScale q x == q * x * (conjugate q)
    rotScale :: Quater t -> Vector t 3 -> Vector t 3
    -- | Creates a quaternion @q@ from two vectors @a@ and @b@,
    --   such that @rotScale q a == b@.
    getRotScale :: Vector t 3 -> Vector t 3 -> Quater t
    -- | Creates a rotation versor from an axis vector and an angle in radians.
    --   Result is always a unit quaternion (versor).
    --   If the argument vector is zero, then result is a real unit quaternion.
    axisRotation :: Vector t 3 -> t -> Quater t
    -- | Quaternion rotation angle \( \alpha \)
    --   (where \( q = c (\cos \frac{\alpha}{2},  v \sin \frac{\alpha}{2}) \)
    --     , \( c > 0 \), \( {|v|}^2 = 1 \)).
    --
    --   >>> q /= 0 ==> axisRotation (imVec q) (qArg q) == signum q
    qArg :: Quater t -> t
    -- | Create a quaternion from a rotation matrix.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, there result of this
    --   function may be ambiguious. Assume the sign of the result to be chosen arbitrarily.
    fromMatrix33 :: Matrix t 3 3 -> Quater t
    -- | Create a quaternion from a homogenious coordinates trasform matrix.
    --   Ignores matrix translation transform.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, there result of this
    --   function may be ambiguious. Assume the sign of the result to be chosen arbitrarily.
    fromMatrix44 :: Matrix t 4 4 -> Quater t
    -- | Create a rotation matrix from a quaternion.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, so the following property holds:
    --
    --   >>> toMatrix33 q == toMatrix33 (-q)
    toMatrix33 :: Quater t -> Matrix t 3 3
    -- | Create a homogenious coordinates trasform matrix from a quaternion.
    --   Translation of the output matrix is zero.
    --   Note, that rotations of \(q\) and \(-q\) are equivalent, so the following property holds:
    --
    --   >>> toMatrix44 q == toMatrix44 (-q)
    toMatrix44 :: Quater t -> Matrix t 4 4

instance (Show t, Quaternion t, Ord t, Num t) => Show (Quater t) where
    showsPrec p (Quater x y z w)
        = case finS of
            SEmpty -> showChar '0'
            Simple -> finF
            SParen -> showParen (p > 6) finF
      where
        (finS, finF) = go SEmpty
          [(w, Nothing), (x, Just 'i'), (y, Just 'j'), (z, Just 'k')]
        go :: ShowState -> [(t, Maybe Char)] -> (ShowState, ShowS)
        go s ((v,l):xs)
          | (s0, f0) <- showComponent s v l
          , (s', f') <- go s0 xs
            = (s', f0 . f')
        go s [] = (s, id)
        showLabel Nothing  = id
        showLabel (Just c) = showChar c
        showComponent :: ShowState -> t -> Maybe Char -> (ShowState, ShowS)
        showComponent sState val mLabel = case (sState, compare val 0) of
          (_     , EQ) -> ( sState, id )
          (SEmpty, GT) -> ( Simple, shows val . showLabel mLabel )
          (SEmpty, LT) -> ( SParen, shows val . showLabel mLabel )
          (_     , GT) -> ( SParen
                          , showString " + " . shows val . showLabel mLabel )
          (_     , LT) -> ( SParen
                          , showString " - " . shows (negate val) . showLabel mLabel )

data ShowState = SEmpty | Simple | SParen
    deriving Eq

instance (Read t, Quaternion t, Num t) => Read (Quater t) where
    readPrec     = parens $ readPrec >>= go id 0 0 0 0
      where
        go :: (t -> t) -> t -> t -> t -> t -> t -> ReadPrec (Quater t)
        go f x y z w new =
          let def = pure (Quater x y z (f new))
              withLabel EOF         = def
              withLabel (Ident "i")
                = (lexP >>= proceed (f new) y z w) <++ pure (Quater (f new) y z w)
              withLabel (Ident "j")
                = (lexP >>= proceed x (f new) z w) <++ pure (Quater x (f new) z w)
              withLabel (Ident "k")
                = (lexP >>= proceed x y (f new) w) <++ pure (Quater x y (f new) w)
              withLabel l           = proceed x y z (f new) l
          in (lexP >>= withLabel) <++ def
        proceed :: t -> t -> t -> t -> Lexeme -> ReadPrec (Quater t)
        proceed x y z w (Symbol "+") = readPrec >>= go id x y z w
        proceed x y z w (Symbol "-") = readPrec >>= go negate x y z w
        proceed x y z w EOF          = pure (Quater x y z w)
        proceed _ _ _ _ _            = pfail

    readListPrec = readListPrecDefault
    readList     = readListDefault


{-|
Module: MachineLearning.Random
Description: Random generation utility functions.
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Randon generation uitility functions.
-}

module MachineLearning.Random
(
  sample
  , sampleM
  , getRandomRListM
  , getRandomRVectorM
  , getRandomRMatrixM
)

where

import Control.Monad (when)
import System.Random (RandomGen, Random)
import qualified Control.Monad.ST as ST
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV
import Numeric.LinearAlgebra ((><))
import qualified Data.Vector.Mutable as MV
import qualified Control.Monad.Random as RndM
import MachineLearning.Types (R, Vector, Matrix)

-- | Samples `n` (given as a second parameter) values from `list` (given as a third parameter).
sample :: RandomGen g => g -> Int -> V.Vector a -> (V.Vector a, g)
sample gen n xs = RndM.runRand (sampleM n xs) gen

-- | Samples `n` (given as a second parameter) values from `list` (given as a third parameter) inside RandomMonad.
sampleM :: RandomGen g => Int -> V.Vector a -> RndM.Rand g (V.Vector a)
sampleM n xs = do  -- Random Monad starts
  let rangeList = V.fromList $ zip (repeat 0) [n..(length xs)-1]
  rnds <- randomsInRangesM rangeList
  let (pre, post) = V.splitAt n xs
  let ys = ST.runST $ do  -- ST Monad starts
        mv <- V.thaw pre
        V.zipWithM_ (\val r -> when (r < n) $ MV.write mv (mod r n) val) post rnds
        V.unsafeFreeze mv
  return ys

-- | Returns a list of random values distributed in a closed interval `range`
getRandomRListM :: (RandomGen g, Random a) =>
                   Int             -- ^ list's lengths
                   -> (a, a)       -- ^ range
                   -> RndM.Rand g [a]          -- ^ list of random values inside RandomMonad
getRandomRListM size range = mapM (\_ -> RndM.getRandomR range) [1..size]

-- | Returns a vector of random values distributed in a closed interval `range`
getRandomRVectorM :: RandomGen g =>
                    Int                     -- ^ vector's length
                    -> (R, R)               -- ^ range
                    -> RndM.Rand g Vector   -- vector of randon values inside RandomMonad
getRandomRVectorM size range = SV.fromList <$> getRandomRListM size range

-- | Returns a matrix of random values distributed in a closed interval `range`
getRandomRMatrixM :: RandomGen g =>
                    Int                     -- ^ number of rows
                    -> Int                  -- ^ number of columns
                    -> (R, R)               -- ^ range
                    -> RndM.Rand g Matrix   -- vector of randon values inside RandomMonad
getRandomRMatrixM r c range = (r><c) <$> getRandomRListM (r*c) range

-- | Takes a list of ranges `(lo, hi)`,
-- returns a list of random values uniformly distributed in the list of closed intervals [(lo, hi)].
randomsInRangesM :: (RndM.RandomGen g, RndM.Random a) => V.Vector (a, a) -> RndM.Rand g (V.Vector a)
randomsInRangesM rangeList = mapM RndM.getRandomR rangeList


{-# LANGUAGE FlexibleContexts #-}
module Main where

import           System.Exit
import           System.Random
import           System.TimeIt
import           Control.Monad.State
import           Numeric.LinearAlgebra
import           Control.Monad.Trans.Maybe      ( runMaybeT )
import           Control.Monad.Writer           ( runWriter )
import           Criterion.Main
import           Criterion.Types

import           Neural.Network
import           Neural.Activation              ( sigmoid
                                                , sigmoid'
                                                )
import           Neural.Training
import           Data.MNIST                     ( loadData )

shallowShape :: [Int]
shallowShape = [784, 30, 10]

deepShape :: [Int]
deepShape = [784, 30, 30, 30, 10]

-- HYPER PARAMETERS
-- learing rate
μ :: Double
μ = 0.5

-- Mini batch size
mbs :: Int
mbs = 100

-- Amount of training epochs
epochs :: Int
epochs = 50

main :: IO ()
main = defaultMain
  [ env setupEnv $ \ ~(trainEx, testEx) ->
      let trainShallow = timeIt (randomIO >>= trainNetwork trainEx testEx shallowShape)
          trainDeep    = timeIt (randomIO >>= trainNetwork trainEx testEx deepShape)
      in  bgroup
            "main"
            [ bench "shallow" $ nfIO trainShallow
            , bench "deep"    $ nfIO trainDeep
            ]
  ]

trainNetwork
  :: [TrainingExample Double]
  -> [TestingExample Double]
  -> [Int]
  -> Seed
  -> IO (Network Double)
trainNetwork trainEx testEx shape seed = do
  let gen             = mkStdGen seed
  let trainFunc = train sigmoid sigmoid' μ trainEx (Just testEx) epochs mbs
  let (randNet, gen') = runState (randomNetwork shape) gen
  let (net, logs) = runWriter $ evalStateT (trainFunc randNet) gen'
  putStrLn "Training network..."

  mapM_ putStrLn logs
  return net

setupEnv
    :: IO ([TrainingExample Double], [TestingExample Double])
setupEnv = do
  mnistData <- runMaybeT loadData
  case mnistData of
    Nothing -> putStrLn "Failed to load MNIST data." >> exitFailure
    Just (trainEx, testEx) -> return (trainEx, testEx)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

{-|
Module      : Numeric.LinearAlgebra.MatrixFree
Description : Represent matrices as a pair of functions for forward and adjoint.
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearAlgebra.MatrixFree
    ( LinearMap(..)
    , (##>)
    , (*#)
    , (+#)
    , (<#>)
    , avgRepeatedFree
    , exactDimsFree
    , eyeFree
    , removeAvg
    , fromL
    , innerDimsFree
    , toL
    , konst
    , diag
    , forward
    , adjoint
    )
where

import           Data.Proxy                     ( Proxy(..) )
import           Data.Type.Equality             ( (:~:)(Refl) )
import           GHC.TypeLits
import           Numeric.LinearAlgebra          ( Transposable(..) )
import           Numeric.LinearAlgebra.Static hiding (konst, diag)
import qualified Numeric.LinearAlgebra.Static  as LS
import qualified Data.Vector.Storable          as V

-- | A linear operator/map between two vector spaces
data LinearMap (m :: Nat) (n :: Nat) where
  LinearMap :: (KnownNat m, KnownNat n) =>
               (R n -> R m) -- ^ Forward Function
            -> (R m -> R n) -- ^ Adjoint FunctionLS
            -> LinearMap m n

-- | Extract the foward function from a `LinearMap`
forward :: LinearMap m n -> (R n -> R m)
forward (LinearMap f _) = f

-- | Extract the adjoint function from a `LinearMap`
adjoint :: LinearMap m n -> (R m -> R n)
adjoint (LinearMap _ a) = a

instance Transposable (LinearMap n m) (LinearMap m n) where
   -- We do not have different functions for conjugate transposes in the
   -- LinearMap constructor (yet?)
    tr  (LinearMap f a) = LinearMap a f
    tr' = tr

---------------------------------------------------------------------
--        Conversions between matrices and matrix free forms       --
---------------------------------------------------------------------

-- | Convert a static matrix into the matrix free data type. This merely
-- wraps the matrix in a function (eta abstraction); it does not determine
-- the function associated with a given matrix.
fromL :: (KnownNat m, KnownNat n) =>
         L m n -- ^ Static matrix to convert to matrix free data type
      -> LinearMap m n
fromL a = LinearMap (a #>) (tr a #>)

-- | Convert a matrix free linear map into a static matrix.
--
-- This methods is very inefficient as it traverses several components
-- linearly several times.
toL :: (KnownNat m, KnownNat n) => LinearMap m n -> L m n
toL (LinearMap f _) =
    matrix $ concatMap (V.toList . extract . f) (toColumns eye)

---------------------------------------------------------------------------
--                           Example LinearMaps                          --
---------------------------------------------------------------------------

-- | The identity matrix in matrix free form.
eyeFree :: (KnownNat n) => LinearMap n n
eyeFree = LinearMap id id

-- | Take the average of a vector.
avg :: forall n. (KnownNat n) => R n -> Double
avg v = v <.> 1 / n' where n' = fromIntegral . natVal $ (Proxy :: Proxy n)

-- | mean(v) * 1, where 1 is the ones vector. of length v.
avgRepeatedFree :: (KnownNat n) => LinearMap n n
avgRepeatedFree = LinearMap f f where f v = LS.konst (avg v)

-- | Remove the average value of a vector from each element of the vector.
removeAvg :: (KnownNat n) => LinearMap n n
removeAvg = eyeFree +# ((-1) *# avgRepeatedFree)

-- | The constant matrix, which sums the vector and sets every element of
-- the resulting vector to that sum.
konst :: (KnownNat m, KnownNat n) => Double -> LinearMap m n
konst k = LinearMap sumAndRepeat sumAndRepeat
  where
    sumAndRepeat :: (KnownNat p, KnownNat q) => R p -> R q
    sumAndRepeat v = LS.konst (v <.> LS.konst k)

-- | Convert a vector into a diagonal matrix.
diag :: (KnownNat n) => R n -> LinearMap n n
diag r = LinearMap f f
  where
    f v = zipWithVector (*) r v

-- TODO: Define fft in matrix free form using FFT and IFFT

---------------------------------------------------------------------
--                Operators for composing LinearMaps               --
---------------------------------------------------------------------

-- | Apply matrix free LinearMap to a vector
(##>) :: LinearMap n m -> R m -> R n
(##>) (LinearMap f _) = f

infixr 8 ##>

-- | Matrix multiply of two `LinearMap`s
(<#>) :: LinearMap m n
      -> LinearMap n p
      -> LinearMap m p
(<#>) (LinearMap f_a a_a) (LinearMap f_b a_b) =
    LinearMap (f_a . f_b) (a_b . a_a)

infixr 8 <#>

-- | Multiply a LinearMap matrix by a constant real number
(*#) :: (KnownNat m, KnownNat n) => Double -> LinearMap m n -> LinearMap m n
(*#) s (LinearMap f a) = LinearMap (\v -> LS.konst s * f v) (\v -> LS.konst s * a v)

infixr 8 *#

-- | Add two `LinearMap`s together if they have the same shape.
(+#) :: (KnownNat m, KnownNat n) =>
        LinearMap m n
     -> LinearMap m n
     -> LinearMap m n
(+#) (LinearMap f_a a_a) (LinearMap f_b a_b) =
    LinearMap (\v -> f_b v + f_a v) (\v -> a_b v + a_a v)

infixr 7 +#

-----------------------------------------------------------------------------
--                    Runtime Proofs on LinearMap shape                    --
-----------------------------------------------------------------------------

-- | Useful for constraining two dependently typed LinearMap matrices to
-- match each other in dimensions when they are unknown at compile-time.
exactDimsFree :: forall m n j k.
                 (KnownNat n, KnownNat m, KnownNat j, KnownNat k) =>
                 LinearMap m n
              -> Maybe (LinearMap j k)
exactDimsFree (LinearMap f b) = do
    Refl <- sameNat (Proxy :: Proxy m) (Proxy :: Proxy j)
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy k)
    return $ LinearMap f b

-- | Useful for constraining that two dependently typed LinearMap matrices
-- match in their inner dimensions (enabling matrix multiplication) each
-- other in dimensions when they are unknown at compile-time.
innerDimsFree :: forall m n p q. (KnownNat n, KnownNat p) =>
                 LinearMap m n   -- ^ The left LinearMap in a matrix multiply
              -> LinearMap p q   -- ^ The right LinearMap in a matrix multiply
              -> Maybe (n :~: p) -- ^ Proof that the inner dimensions are equal
innerDimsFree (LinearMap _ _) (LinearMap _ _) = do
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy p)
    return (Refl)


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Linear algebra types and operations
--

module Numeric.LinearAlgebra (
    -- * Vector types
    Vector,
    RVector,
    STVector,
    IOVector,

    -- * Matrix types
    Matrix,
    RMatrix,
    STMatrix,
    IOMatrix,

    -- * Packed matrix types
    Packed,
    RPacked,
    STPacked,
    IOPacked,

    module Numeric.LinearAlgebra.Types,
    ) where

import Numeric.LinearAlgebra.Types
import Numeric.LinearAlgebra.Vector( Vector, RVector, STVector, IOVector )
import Numeric.LinearAlgebra.Matrix( Matrix, RMatrix, STMatrix, IOMatrix )
import Numeric.LinearAlgebra.Packed( Packed, RPacked, STPacked, IOPacked )


module MachineLearning.PCATest
(
  tests
)

where

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

import qualified Numeric.LinearAlgebra as LA

import MachineLearning.PCA

getDimReducerSmokeTest =
  let nFeatures = 4
      nExamples = 7
      m = LA.matrix nFeatures [1 .. fromIntegral $ nFeatures*nExamples]
      m10 = m * 10
      (reduceDims, retainedVariance, mReduced) = getDimReducer m 2
      m10Reduced = reduceDims m10
  in do
    assertEqual "dimension equality (getDimReducer)" (LA.cols mReduced) 2
    assertEqual "dimension equality (reduceDims)" (LA.cols m10Reduced) 2
    assertApproxEqual "retained variance" 1e-10 1 retainedVariance

getDimReducer_rvSmokeTest rv=
  let nFeatures = 4
      nExamples = 7
      m = LA.matrix nFeatures [1 .. fromIntegral $ nFeatures*nExamples]
      m10 = m * 10
      (reduceDims, k, mReduced) = getDimReducer_rv m rv
      m10Reduced = reduceDims m10
  in do
    assertEqual "dimension equality (getDimReducer_rv)" (LA.cols mReduced) k
    assertEqual "dimension equality (reduceDims_rv)" (LA.cols m10Reduced) k
    assertBool "reduced number of dimensions" $ k <= nFeatures

tests = [ testGroup "smoke test" [
            testCase "getDimReducer" getDimReducerSmokeTest
            , testCase "getDimReducer_rv, rv = 0" $ getDimReducer_rvSmokeTest 0
            , testCase "getDimReducer_rv, rv = 0.5" $ getDimReducer_rvSmokeTest 0.5
            , testCase "getDimReducer_rv, rv = 1" $ getDimReducer_rvSmokeTest 1
            , testCase "getDimReducer_rv, rv = 2" $ getDimReducer_rvSmokeTest 2
            ]
        ]


module D03Spec (spec) where

import Test.Hspec
import D03
import Data.Complex

spec :: Spec
spec = do

  describe "utilities" $ do
    it "parses into claims" $ mkClaims "#1 @ 146,196: 19x14\n#2 @ 641,817: 27x28" `shouldBe` [ Claim 1 (146:+196) (164:+209), Claim 2 (641:+817) (667:+844) ]

  describe "intersections size" $ do
    it "simple example" $ claimsIntersectionSize "#1 @ 1,3: 4x4\n#2 @ 3,1: 4x4\n#3 @ 5,5: 2x2" `shouldBe` 4
    it "other example" $ claimsIntersectionSize "#1 @ 1,2: 3x3\n#2 @ 2,3: 3x4\n#3 @ 1,5: 3x3" `shouldBe` 8
    it "overlap example" $ claimsIntersectionSize "#1 @ 1,2: 3x3\n#2 @ 2,3: 3x4\n#3 @ 1,5: 3x3\n#4 @ 2,5: 2x3" `shouldBe` 10
    it "other overlap example" $ claimsIntersectionSize "#1 @ 1,2: 3x3\n#2 @ 2,3: 3x4\n#3 @ 1,5: 3x3\n#4 @ 2,5: 2x3\n#5 @ 2,2: 3x3" `shouldBe` 14
    it "submits answer" $ do input <- readFile "test/D03.txt"; claimsIntersectionSize input `shouldBe` 106501

  describe "isolated claim" $ do
    it "simple example" $ isolatedClaim "#1 @ 1,3: 4x4\n#2 @ 3,1: 4x4\n#3 @ 5,5: 2x2" `shouldBe` 3
    it "submits answer" $ do input <- readFile "test/D03.txt"; isolatedClaim input `shouldBe` 632


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UnboxedTuples         #-}
{-# LANGUAGE UndecidableInstances  #-}
module Numeric.Matrix.Internal
  ( MatrixTranspose (..)
  , SquareMatrix (..)
  , MatrixDeterminant (..)
  , MatrixInverse (..)
  , Matrix
  , HomTransform4 (..)
  , Mat22f, Mat23f, Mat24f
  , Mat32f, Mat33f, Mat34f
  , Mat42f, Mat43f, Mat44f
  , Mat22d, Mat23d, Mat24d
  , Mat32d, Mat33d, Mat34d
  , Mat42d, Mat43d, Mat44d
  , mat22, mat33, mat44
  , (%*)
  ) where

import GHC.Base
import Numeric.DataFrame.Contraction        ((%*))
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Scalar.Internal
import Numeric.Vector.Internal

-- | Alias for DataFrames of rank 2
type Matrix (t :: l) (n :: k) (m :: k) = DataFrame t '[n,m]

class MatrixTranspose t (n :: k) (m :: k) where
    -- | Transpose Mat
    transpose :: Matrix t n m -> Matrix t m n

class SquareMatrix t (n :: Nat) where
    -- | Mat with 1 on diagonal and 0 elsewhere
    eye :: Matrix t n n
    -- | Put the same value on the Mat diagonal, 0 otherwise
    diag :: Scalar t -> Matrix t n n
    -- | Sum of diagonal elements
    trace :: Matrix t n n -> Scalar t

class MatrixDeterminant t (n :: Nat) where
    -- | Determinant of  Mat
    det :: Matrix t n n -> Scalar t

class MatrixInverse t (n :: Nat) where
    -- | Matrix inverse
    inverse :: Matrix t n n -> Matrix t n n

-- | Operations on 4x4 transformation matrices and vectors in homogeneous coordinates.
--   All angles are specified in radians.
--
--   Note: since version 2 of @easytensor@, DataFrames and matrices are row-major.
--         A good SIMD implementation may drastically improve performance
--         of 4D vector-matrix products of the form @v %* m@, but not so much
--         for products of the form @m %* v@ (due to memory layout).
--         Thus, all operations here assume the former form to benefit more from
--         SIMD in future.
class HomTransform4 t where
    -- | Create a translation matrix from a vector.  The 4th coordinate is ignored.
    --
    --   If @p ! 3 == 1@ and @v ! 3 == 0@, then
    --
    --   > p %* translate4 v == p + v
    --
    translate4  :: Vector t 4 -> Matrix t 4 4
    -- | Create a translation matrix from a vector.
    --
    --   If @p ! 3 == 1@, then
    --
    --   > p %* translate3 v == p + toHomVector v
    --
    translate3  :: Vector t 3 -> Matrix t 4 4
    -- | Rotation matrix for a rotation around the X axis, angle is given in radians.
    --   e.g. @p %* rotateX (pi/2)@ rotates point @p@ around @Ox@ by 90 degrees.
    rotateX     :: t -> Matrix t 4 4
    -- | Rotation matrix for a rotation around the Y axis, angle is given in radians.
    --   e.g. @p %* rotateY (pi/2)@ rotates point @p@ around @Oy@ by 90 degrees.
    rotateY     :: t -> Matrix t 4 4
    -- | Rotation matrix for a rotation around the Z axis, angle is given in radians.
    --   e.g. @p %* rotateZ (pi/2)@ rotates point @p@ around @Oz@ by 90 degrees.
    rotateZ     :: t -> Matrix t 4 4
    -- | Rotation matrix for a rotation around an arbitrary normalized vector
    --   e.g. @p %* rotate (pi/2) v@ rotates point @p@ around @v@ by 90 degrees.
    rotate      :: Vector t 3 -> t -> Matrix t 4 4
    -- | Rotation matrix from the Euler angles roll (axis @Z@), yaw (axis @Y'@), and pitch (axis @X''@).
    --   This order is known as Tait-Bryan angles (@Z-Y'-X''@ intrinsic rotations), or nautical angles, or Cardan angles.
    --
    --   > rotateEuler pitch yaw roll == rotateZ roll %* rotateY yaw %* rotateX pitch
    --
    --   https://en.wikipedia.org/wiki/Euler_angles#Conventions_2
    rotateEuler :: t -- ^ pitch (axis @X''@)
                -> t -- ^ yaw (axis @Y'@)
                -> t -- ^ roll (axis @Z@)
                -> Matrix t 4 4
    -- | Create a transform matrix using up direction, camera position and a point to look at.
    --   Just the same as GluLookAt.
    lookAt      :: Vector t 3 -- ^ The up direction, not necessary unit length or perpendicular to the view vector
                -> Vector t 3 -- ^ The viewers position
                -> Vector t 3 -- ^ The point to look at
                -> Matrix t 4 4
    -- | A perspective symmetric projection matrix. Right-handed coordinate system. (@x@ - right, @y@ - top)
    --   http://en.wikibooks.org/wiki/GLSL_Programming/Vertex_Transformations
    perspective :: t -- ^ Near plane clipping distance (always positive)
                -> t -- ^ Far plane clipping distance (always positive)
                -> t -- ^ Field of view of the y axis, in radians
                -> t -- ^ Aspect ratio, i.e. screen's width\/height
                -> Matrix t 4 4
    -- | An orthogonal symmetric projection matrix. Right-handed coordinate system. (@x@ - right, @y@ - top)
    --   http://en.wikibooks.org/wiki/GLSL_Programming/Vertex_Transformations
    orthogonal  :: t -- ^ Near plane clipping distance
                -> t -- ^ Far plane clipping distance
                -> t -- ^ width
                -> t -- ^ height
                -> Matrix t 4 4
    -- | Add one more dimension and set it to 1.
    toHomPoint  :: Vector t 3 -> Vector t 4
    -- | Add one more dimension and set it to 0.
    toHomVector :: Vector t 3 -> Vector t 4
    -- | Transform a homogenous vector or point into a normal 3D vector.
    --   If the last coordinate is not zero, divide the rest by it.
    fromHom     :: Vector t 4 -> Vector t 3

-- Type abbreviations

type Mat22f = Matrix Float 2 2
type Mat32f = Matrix Float 3 2
type Mat42f = Matrix Float 4 2
type Mat23f = Matrix Float 2 3
type Mat33f = Matrix Float 3 3
type Mat43f = Matrix Float 4 3
type Mat24f = Matrix Float 2 4
type Mat34f = Matrix Float 3 4
type Mat44f = Matrix Float 4 4

type Mat22d = Matrix Double 2 2
type Mat32d = Matrix Double 3 2
type Mat42d = Matrix Double 4 2
type Mat23d = Matrix Double 2 3
type Mat33d = Matrix Double 3 3
type Mat43d = Matrix Double 4 3
type Mat24d = Matrix Double 2 4
type Mat34d = Matrix Double 3 4
type Mat44d = Matrix Double 4 4

-- | Compose a 2x2D matrix
mat22 :: PrimBytes (t :: Type)
      => Vector t 2 -> Vector t 2 -> Matrix t 2 2
mat22 = packDF @_ @2 @'[2]
{-# INLINE mat22 #-}

-- | Compose a 3x3D matrix
mat33 :: PrimBytes (t :: Type)
      => Vector t 3 -> Vector t 3 -> Vector t 3 -> Matrix t 3 3
mat33 = packDF @_ @3 @'[3]
{-# INLINE mat33 #-}

-- | Compose a 4x4D matrix
mat44 :: PrimBytes (t :: Type)
      => Vector t 4
      -> Vector t 4
      -> Vector t 4
      -> Vector t 4
      -> Matrix t 4 4
mat44 = packDF @_ @4 @'[4]
{-# INLINE mat44 #-}

instance ( KnownDim n, KnownDim m
         , PrimArray t (Matrix t n m)
         , PrimArray t (Matrix t m n)
         ) => MatrixTranspose t (n :: Nat) (m :: Nat) where
    transpose df = case uniqueOrCumulDims df of
      Left a -> broadcast a
      Right _
         | wm <- dimVal' @m
         , wn <- dimVal' @n
         , m <- case wm of W# w -> word2Int# w
         , n <- case wn of W# w -> word2Int# w
         -> let f ( I# i,  I# j )
                  | isTrue# (i ==# n) = f ( 0, I# (j +# 1#) )
                  | otherwise         = (# ( I# (i +# 1#), I# j )
                                         , ix# (i *# m +# j) df
                                         #)
            in case gen# (CumulDims [wm*wn, wn, 1]) f (0,0) of (# _, r #) -> r

instance MatrixTranspose (t :: Type) (xn :: XNat) (xm :: XNat) where
    transpose (XFrame (df :: DataFrame t ns))
      | ((D :: Dim n) :* (D :: Dim m) :* U) <- dims @ns
      , Dict <- inferPrimElem df
      = XFrame (transpose df :: Matrix t m n)
    transpose _ = error "MatrixTranspose/transpose: impossible argument"

instance (KnownDim n, PrimArray t (Matrix t n n), Num t)
      => SquareMatrix t n where
    eye
      | n <- dimVal' @n
      = let f 0 = (# n    , 1 #)
            f k = (# k - 1, 0 #)
        in case gen# (CumulDims [n*n, n, 1]) f 0 of
            (# _, r #) -> r
    diag se
      | n <- dimVal' @n
      , e <- unScalar se
      = let f 0 = (# n    , e #)
            f k = (# k - 1, 0 #)
        in case gen# (CumulDims [n*n, n, 1]) f 0 of
            (# _, r #) -> r
    trace df
      | I# n <- fromIntegral $ dimVal' @n
      , n1 <- n +# 1#
      = let f 0# = ix# 0# df
            f k  = ix# k  df + f (k -# n1)
        in scalar $ f (n *# n -# 1#)


-- TODO: Avoid recomputing norms and reg const

{-# LANGUAGE BangPatterns, FlexibleContexts, NamedFieldPuns #-}

module School.Unit.LogSoftMax
( logSoftMax ) where

import Numeric.LinearAlgebra (Container, Matrix, R, Vector, add, atIndex,
                              build, cmap, maxElement, size, sumElements)
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitParams (UnitParams(..))
import School.Unit.UnitGradient (UnitGradient(..))
import School.Utils.LinearAlgebra (mapRows, sumCols)

logSoftMax :: Unit R
logSoftMax = Unit
  { apply = lsmApply
  , deriv = lsmDeriv
  }

getNorm :: (Container Vector a, Floating a)
        => a
        -> Vector a
        -> a
getNorm reg = log
      . sumElements
      . cmap (exp . flip (-) reg)

handleRow :: (Container Vector a, Floating a)
          => Vector a
          -> Vector a
handleRow v =
  let !reg = maxElement v
      !norm = getNorm reg v
  in cmap (flip (-) (norm + reg)) v

lsmApply :: UnitParams R
         -> UnitActivation R
         -> UnitActivation R
lsmApply EmptyParams (BatchActivation inMatrix) =
  BatchActivation $ mapRows handleRow inMatrix
lsmApply _ (ApplyFail msg) = ApplyFail msg
lsmApply _ _ = ApplyFail "Failure when applying log softmax unit"

failure :: String
        -> ( UnitGradient R
           , UnitParams R
           )
failure msg = (GradientFail msg, EmptyParams)

builder :: (Container Vector a, RealFrac a)
        => Matrix a
        -> Vector a
        -> Vector a
        -> (a -> a -> a)
builder exps norms gradCols dj dk =
  let j = round dj :: Int
      k = round dk :: Int
  in (-1) * atIndex exps (j, k)
   * atIndex gradCols j
   / atIndex norms j

lsmDeriv :: UnitParams R
         -> UnitGradient R
         -> UnitActivation R
         -> ( UnitGradient R
            , UnitParams R
            )
lsmDeriv EmptyParams
         (BatchGradient inGrad)
         (BatchActivation input) =
  (outGrad, EmptyParams) where
    !reg = maxElement input
    !exps = cmap (exp . flip (-) reg) input
    !norms = sumCols exps
    !gradCols = sumCols inGrad
    delta = build (size input) (builder exps norms gradCols)
    outGrad = BatchGradient $ add inGrad delta
lsmDeriv _ (GradientFail msg) _ = failure msg
lsmDeriv _ _ (ApplyFail msg) = failure msg
lsmDeriv _ _ _ = failure "Failure when differentiating log softmax unit"


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}

module Models where

import Foreign.Storable (Storable)
import Numeric.LinearAlgebra ((<.>), (<>), (#>), inv, sumElements, Container)
import Numeric.LinearAlgebra.Data (vector, matrix, Matrix, (><), Vector, tr'
                                  ,cmap, ident, size, fromRows)

class Trainable m a b where
  train :: m a b -> [(a,b)] -> m a b

class Predictive m where
  predict :: m a b -> a -> b

type LBWeights = Matrix Double
type LBFuncs a = [a -> Double]

data LBReg a b = LBReg { lbFuncs   :: LBFuncs a
                       , lbWeights :: LBWeights
                       , lbError   :: Double
                       , lbReg     :: Double
                       , lbPredict :: LBWeights
                                   -> LBFuncs a
                                   -> a
                                   -> b
                       }

instance Trainable LBReg [a] (Vector Double) where
  train = trainLBReg

instance Predictive LBReg where
  predict m = lbPredict m (lbWeights m) (lbFuncs m)

updateWeights x w e = x { lbWeights = w, lbError = e }

trainLBReg :: LBReg [a] (Vector Double) -> [([a], Vector Double)] -> LBReg [a] (Vector Double)
trainLBReg m d = updateWeights m w e where
  t    = fromRows (fmap snd d)
  fs   = lbFuncs m
  dps  = fmap fst d
  p    = lbPredict m
  reg  = lbReg m
  iden = ident $ 1 + (length . head) dps :: Matrix Double
  dm   = tr' $ (length fs >< length dps) $ fs <*> dps
  w    = (inv (cmap (reg*) iden + (tr' dm <> dm)) <>  tr' dm) <> t
  e    = 2--sosError t pr where pr = vector $ fmap (p w fs) dps

sosError :: Vector Double -> Vector Double -> Double
sosError t pr = 0--sumElements $ cmap ((/2) . (**2)) $ t - pr

lbEval :: LBWeights -> LBFuncs a -> a -> Vector Double
lbEval w fs d = tr' w #> vector (zipWith ($) fs $ replicate (length fs) d)

mkLBId :: Int -> LBFuncs [Double]
mkLBId n = const 1 : [(!!x) | x <- [0 .. n-1]]

mkSLBReg :: Double -> [([Double], Vector Double)] -> LBReg [Double] (Vector Double)
mkSLBReg reg ds = train LBReg { lbFuncs   = mkLBId n
                              , lbWeights = (n><k) [1 .. ]
                              , lbPredict = lbEval
                              , lbError   = fromIntegral (maxBound :: Int)
                              , lbReg     = reg
                              } ds where n = length $ fst $ head ds
                                         k = size $ snd $ head ds