{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
{-# LANGUAGE StrictData       #-}
module FourierPinwheel.Harmonics where

import           Data.Array.IArray   as IA
import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List
import           Data.List           as L
import           Data.Vector.Generic as VG
import           Data.Vector.Unboxed as VU
import           Filter.Utils
import           Math.Gamma
import           Utils.Parallel

data FPData vector = FPData
  { getFPDataHarmonics       :: [vector]
  , getFPDataHarmonicsOffset :: [vector]
  , getFPDataCoef            :: [vector]
  , getFPDataCoefHollow      :: [vector]
  }

createHarmonics ::
     ( RealFloat e
     , Unbox e
     , NFData e
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     , Gamma (Complex e)
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array U DIM4 (Complex e)
  -> IO (FPData (vector (Complex e)))
createHarmonics numR2Freqs phiFreq rhoFreq thetaFreq rFreq alpha periodR2 periodEnv coefficients = do
  let harmonicVecs =
        parMap
          rdeepseq
          (\radialFreq ->
             createVector
               numR2Freqs
               (2 * phiFreq + 1)
               1
               (pinwheel periodR2 radialConst alpha radialFreq))
          [-rhoFreq .. rhoFreq]
      harmonicVecsOffset =
        parMap
          rdeepseq
          (\radialFreq ->
             createVector
               numR2Freqs
               (2 * thetaFreq + 1)
               (-1)
               (logpolarHarmonics periodR2 radialConst radialFreq)) .
        L.reverse $
        [-rFreq .. rFreq]
  phaseShiftedCoefArr <-
    computeUnboxedP .
    R.zipWith (*) coefficients . fromFunction (extent coefficients) $ \(Z :. r :. theta :. rho :. phi) ->
      phaseShift
        (phi - phiFreq - (theta - thetaFreq))
        (rho - rhoFreq - (r - rFreq))
        radialConst
        alpha
  let phaseShiftedCoefVecs =
        parMap
          rdeepseq
          (\r ->
             VU.convert . toUnboxed . computeS . R.slice phaseShiftedCoefArr $
             (Z :. r :. All :. All :. All))
          [0 .. 2 * rFreq]
  coefHollowArr <- computeUnboxedP $ sumArr *^ coefficients
  let coefHollowVecs =
        parMap
          rdeepseq
          (\r ->
             VU.convert . toUnboxed . computeS . R.slice coefHollowArr $
             (Z :. r :. All :. All :. All))
          [0 .. 2 * rFreq]
  return $
    FPData harmonicVecs harmonicVecsOffset phaseShiftedCoefVecs coefHollowVecs
  where
    radialConst = 2 * pi / log periodEnv
    !sumArr1 =
      computePinwheelArray
        numR2Freqs
        phiFreq
        rhoFreq
        thetaFreq
        rFreq
        alpha
        periodR2
        periodEnv
    sumArr =
      fromListUnboxed (extent coefficients) .
      parMap
        rdeepseq
        (\(r, theta, rho, phi) ->
           let angularFreq = phi - theta
               radialFreq = rho - r
            in sumArr1 IA.! (radialFreq, angularFreq)) $
      [ (r, theta, rho, phi)
      | r <- [-rFreq .. rFreq]
      , theta <- [-thetaFreq .. thetaFreq]
      , rho <- [-rhoFreq .. rhoFreq]
      , phi <- [-phiFreq .. phiFreq]
      ]

{-# INLINE createVector #-}
createVector ::
     (RealFloat e, Unbox e, VG.Vector vector (Complex e))
  => Int
  -> Int
  -> Int
  -> (Int -> e -> e -> Complex e)
  -> vector (Complex e)
createVector numR2Freqs numAngularFreqs sign f =
  let centerR2 = div numR2Freqs 2
      centerAngular = div numAngularFreqs 2
   in VG.convert .
      toUnboxed .
      computeS .
      makeFilter2D .
      fromFunction (Z :. numAngularFreqs :. numR2Freqs :. numR2Freqs) $ \(Z :. k :. i :. j) ->
        let x = fromIntegral (i - centerR2)
            y = fromIntegral (j - centerR2)
            rho = sqrt $ x ^ 2 + y ^ 2
            phi = atan2 y x
         in if rho == 0
              then 0
              else f (sign * (k - centerAngular)) phi rho

{-# INLINE pinwheel #-}
pinwheel :: (RealFloat e) => e -> e -> e -> Int -> Int -> e -> e -> Complex e
pinwheel periodR2 radialConst alpha radialFreq angularFreq phi rho =
  pi * cis (fromIntegral angularFreq * phi) *
  ((periodR2 / (pi * rho) :+ 0) **
   ((2 + alpha) :+ (radialConst * fromIntegral radialFreq)))

{-# INLINE logpolarHarmonics #-}
logpolarHarmonics ::
     (RealFloat e) => e -> e -> Int -> Int -> e -> e -> Complex e
logpolarHarmonics periodR2 radialConst radialFreq angularFreq phi rho =
  cis (fromIntegral angularFreq * phi) *
  ((periodR2 / (pi * rho) :+ 0) **
   (0 :+ (radialConst * fromIntegral radialFreq)))

{-# INLINE phaseShift #-}
phaseShift ::
     (RealFloat e, Gamma (Complex e)) => Int -> Int -> e -> e -> Complex e
phaseShift angularFreq radialFreq radialConst alpha =
  ((0 :+ (-1)) ^ abs angularFreq) *
  (gamma $
   ((2 + fromIntegral (abs angularFreq) + alpha) :+
    (radialConst * fromIntegral radialFreq)) /
   2) /
  (gamma $
   ((fromIntegral (abs angularFreq) - alpha) :+
    (radialConst * fromIntegral (-radialFreq))) /
   2)

{-# INLINE computeIndicesFromRadius #-}
computeIndicesFromRadius :: Double -> [(Int,Int)]
computeIndicesFromRadius radius =
  let rho = round radius
      r2 = rho ^ 2
   in L.filter
        (\(x, y) ->
           let r = fromIntegral $ x ^ 2 + y ^ 2
            in r < r2)
        [(x, y) | x <- [-rho .. rho], y <- [-rho .. rho]]

{-# INLINE computePinwheelArray #-}
computePinwheelArray ::
     ( Num e
     , RealFloat e
     , Gamma (Complex e)
     , Unbox e
     , NFData e
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IA.Array (Int, Int) (Complex e)
computePinwheelArray numR2Freqs phiFreq rhoFreq thetaFreq rFreq alpha periodR2 periodEnv =
  let maxAngularFreq = phiFreq + thetaFreq
      maxRadialFreq = rhoFreq + rFreq
      radialConst = 2 * pi / log periodEnv
      centerR2 = div numR2Freqs 2
      idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [-maxRadialFreq .. maxRadialFreq]
        , angularFreq <- [-maxAngularFreq .. maxAngularFreq]
        ]
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             let arr =
                   fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
                     let x = fromIntegral $ i - centerR2
                         y = fromIntegral $ j - centerR2
                         rho = sqrt $ x ^ 2 + y ^ 2
                         phi = atan2 y x
                      in if rho == 0
                           then 0
                           else pinwheel
                                  periodR2
                                  radialConst
                                  alpha
                                  radialFreq
                                  angularFreq
                                  phi
                                  rho *
                                phaseShift
                                  angularFreq
                                  radialFreq
                                  radialConst
                                  alpha
              in sumAllS arr / (fromIntegral numR2Freqs^2))
          idxs
   in listArray
        ((-maxRadialFreq, -maxAngularFreq), (maxRadialFreq, maxAngularFreq))
        pinwheels


{-# LANGUAGE ViewPatterns #-}
-- |
-- Tests for quantile
module Tests.Quantile (tests) where

import Control.Exception
import qualified Data.Vector.Unboxed as U
import Test.Framework
import Test.Framework.Providers.HUnit
import Test.Framework.Providers.QuickCheck2
import Test.HUnit (Assertion,assertEqual,assertFailure)
import Test.QuickCheck hiding (sample)
import Numeric.MathFunctions.Comparison (ulpDelta,ulpDistance)
import Statistics.Quantile

tests :: Test
tests = testGroup "Quantiles"
  [ testCase "R alg. 4" $ compareWithR cadpw (0.00, 0.50, 2.50, 8.25, 10.00)
  , testCase "R alg. 5" $ compareWithR hazen (0.00, 1.00, 5.00, 9.00, 10.00)
  , testCase "R alg. 6" $ compareWithR spss  (0.00, 0.75, 5.00, 9.25, 10.00)
  , testCase "R alg. 7" $ compareWithR s     (0.000, 1.375, 5.000, 8.625,10.00)
  , testCase "R alg. 8" $ compareWithR medianUnbiased
      (0.0, 0.9166666666666667, 5.000000000000003, 9.083333333333334, 10.0)
  , testCase "R alg. 9" $ compareWithR normalUnbiased
      (0.0000, 0.9375, 5.0000, 9.0625, 10.0000)
  , testProperty "alg 7." propWeigtedAverage
    -- Test failures
  , testCase "weightedAvg should throw errors" $ do
      let xs  = U.fromList [1,2,3]
          xs0 = U.fromList []
      shouldError "Empty sample" $ weightedAvg 1 4 xs0
      shouldError "N=0"  $ weightedAvg 1 0 xs
      shouldError "N=1"  $ weightedAvg 1 1 xs
      shouldError "k<0"  $ weightedAvg (-1) 4 xs
      shouldError "k>N"  $ weightedAvg 5    4 xs
  , testCase "quantile should throw errors" $ do
      let xs  = U.fromList [1,2,3]
          xs0 = U.fromList []
      shouldError "Empty xs" $ quantile s 1 4 xs0
      shouldError "N=0"  $ quantile s 1 0 xs
      shouldError "N=1"  $ quantile s 1 1 xs
      shouldError "k<0"  $ quantile s (-1) 4 xs
      shouldError "k>N"  $ quantile s 5    4 xs
    --
  , testProperty "quantiles    are OK" propQuantiles
  , testProperty "quantilesVec are OK" propQuantilesVec
  ]

sample :: U.Vector Double
sample = U.fromList [0, 1, 2.5, 7.5, 9, 10]

-- Compare quantiles implementation with reference R implementation
compareWithR :: ContParam -> (Double,Double,Double,Double,Double) -> Assertion
compareWithR p (q0,q1,q2,q3,q4) = do
  assertEqual "Q 0" q0 $ quantile p 0 4 sample
  assertEqual "Q 1" q1 $ quantile p 1 4 sample
  assertEqual "Q 2" q2 $ quantile p 2 4 sample
  assertEqual "Q 3" q3 $ quantile p 3 4 sample
  assertEqual "Q 4" q4 $ quantile p 4 4 sample

propWeigtedAverage :: Positive Int -> Positive Int -> Property
propWeigtedAverage (Positive k) (Positive q) =
  (q >= 2 && k <= q) ==> let q1 = weightedAvg k q sample
                             q2 = quantile s k q sample
                         in counterexample ("weightedAvg   = " ++ show q1)
                          $ counterexample ("quantile      = " ++ show q2)
                          $ counterexample ("delta in ulps = " ++ show (ulpDelta q1 q2))
                          $ ulpDistance q1 q2 <= 16

propQuantiles :: Positive Int -> Int -> Int -> NonEmptyList Double -> Property
propQuantiles (Positive n)
              ((`mod` n) -> k1)
              ((`mod` n) -> k2)
              (NonEmpty xs)
  =   n >= 2
  ==> [x1,x2] == quantiles s [k1,k2] n rndXs
  where
    rndXs = U.fromList xs
    x1 = quantile s k1 n rndXs
    x2 = quantile s k2 n rndXs

propQuantilesVec :: Positive Int -> Int -> Int -> NonEmptyList Double -> Property
propQuantilesVec (Positive n)
                 ((`mod` n) -> k1)
                 ((`mod` n) -> k2)
                 (NonEmpty xs)
  =   n >= 2
  ==> U.fromList [x1,x2] == quantilesVec s (U.fromList [k1,k2]) n rndXs
  where
    rndXs = U.fromList xs
    x1 = quantile s k1 n rndXs
    x2 = quantile s k2 n rndXs

shouldError :: String -> a -> Assertion
shouldError nm x = do
  r <- try (evaluate x)
  case r of
    Left  (ErrorCall{}) -> return ()
    Right _             -> assertFailure ("Should call error: " ++ nm)


{-# LANGUAGE TemplateHaskell #-}
module AI.NeuralNetwork
    ( Network
    , layerSizes
    , biases
    , weights
    , sigmoid
    , newNetwork
    , feedforward
    , toFile
    , fromFile ) where

import           Control.Lens
import           Control.Monad
import           Control.Monad.State
import           Data.Random.Normal
import           Numeric.LinearAlgebra
import           System.IO
import           System.Random

data Network = Network { _layerSizes :: [Int]
                       , _biases     :: [Vector Double]
                       , _weights    :: [Matrix Double] }
                       deriving (Eq, Show)

makeLenses ''Network

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

newNetwork :: Monad m => [Int] -> StateT StdGen m Network
newNetwork sizes = do
    let matrixSizes = zip (init sizes) (tail sizes)

    biasesL  <- mapM sequence
        [[state normal | _ <- [1..y]] | y <- tail sizes]
    weightsL <- mapM sequence $
        [[state normal | _ <- [1..x], x <- [1..y]] | (x, y) <- matrixSizes]

    let randBiases  = map fromList biasesL
    let randWeights = map (\((x, y), l) -> (y><x) l) (zip matrixSizes weightsL)

    return Network { _layerSizes = sizes
                   , _biases = randBiases
                   , _weights = randWeights }

feedforward :: Vector Double -> Network -> Vector Double
feedforward input nn = foldl (\a (b, w) -> cmap sigmoid ((w #> a) + b)) input bw
    where bw = zip (nn^.biases) (nn^.weights)

toFile :: FilePath -> Network -> IO ()
toFile filepath nn =
    let headerStr  = show (nn^.layerSizes) ++ "\n"
        biasesStr  = unlines $ show . toList  <$> (nn^.biases)
        weightsStr = unlines $ show . toLists <$> (nn^.weights)
        nnStr      = headerStr ++ biasesStr ++ weightsStr
     in writeFile filepath nnStr

fromFile :: FilePath -> IO Network
fromFile filepath = withFile filepath ReadMode (\handle -> do
    layerSizes' <- (read :: String -> [Int]) <$> hGetLine handle
    biases'     <- forM (tail layerSizes') $ \_ -> do
        line <- hGetLine handle
        return $ fromList $ (read :: String -> [Double]) line
    weights'    <- forM (tail layerSizes') $ \_ -> do
        line <- hGetLine handle
        return $ fromLists $ (read :: String -> [[Double]]) line
    return $ Network { _layerSizes = layerSizes'
                     , _biases     = biases'
                     , _weights    = weights' })


module FFTSerialSpec where

import qualified Clash.Prelude as Clash
import Clash.Prelude (Signal, Vec(..), BitVector, Index, Signed, Unsigned, SFixed, Bit, SNat(..),
                      simulate, simulate_lazy, listToVecTH, KnownNat, pack, unpack, (++#), mealy, mux, bundle, unbundle,
                      HiddenClockResetEnable, System)
import Test.Hspec
import Test.QuickCheck

import qualified Numeric.FFT as FFT
import qualified Data.Complex as C
import Clash.DSP.Complex (Complex, fromComplex, toComplex)
import Clash.DSP.FFT.Serial
import Clash.DSP.FFT.Example
import Clash.DSP.FFT.Twiddle (twiddleFactors)

--serial FFT
spec = describe "Serial FFTs" $ do
    specify "Decimation in time equals known good implementation"      $ property prop_fftSerialDIT
    specify "Decimation in frequency equals known good implementation" $ property prop_fftSerialDIF

approxEqual :: Double -> Double -> Bool
approxEqual x y = abs (x - y) < 0.0001

approxEqualComplex (a C.:+ b) (c C.:+ d) = approxEqual a c && approxEqual b d

twiddles4 :: Vec 4 (Complex Double)
twiddles4 = $(listToVecTH (twiddleFactors 4))

ditInputReorder :: Vec 8 a -> Vec 4 (a, a)
ditInputReorder (a :> b :> c :> d :> e :> f :> g :> h :> Nil) = (a, e) :> (c, g) :> (b, f) :> (d, h) :> Nil

ditOutputReorder :: [(a, a)] -> [a]
ditOutputReorder ((a, b) : (c, d) : (e, f) : (g, h) : _) = a : c : e : g : b : d : f : h : []

difInputReorder :: Vec 8 a -> Vec 4 (a, a)
difInputReorder (a :> b :> c :> d :> e :> f :> g :> h :> Nil) = (a, e) :> (b, f) :> (c, g) :> (d, h) :> Nil

difOutputReorder :: [(a, a)] -> [a]
difOutputReorder ((a, b) : (c, d) : (e, f) : (g, h) : _) = a : e : c : g : b : f : d : h : []

prop_fftSerialDIT :: Vec 8 (C.Complex Double) -> Bool
prop_fftSerialDIT vec = and $ zipWith approxEqualComplex (map toComplex result) (FFT.fft (Clash.toList vec))
    where
    result = ditOutputReorder $ drop 8 $ simulate_lazy @System (fftSerialDIT twiddles4 (pure True)) $ (Clash.toList (ditInputReorder (Clash.map fromComplex vec))) ++ repeat (0, 0)

prop_fftSerialDIF :: Vec 8 (C.Complex Double) -> Bool
prop_fftSerialDIF vec = and $ zipWith approxEqualComplex (map toComplex result) (FFT.fft (Clash.toList vec))
    where
    result = difOutputReorder $ drop 8 $ simulate_lazy @System (fftSerialDIF twiddles4 (pure True)) $ (Clash.toList (difInputReorder (Clash.map fromComplex vec))) ++ repeat (0, 0)


module EquivarianceDiagram where

import           Data.Array.Repa     as R
import           Data.Complex
import           Data.List           as L
import           Data.Vector.Unboxed as VU
import           DFT.Plan
import           Filter.Gaussian
import           Image.IO
import           System.Directory
import           System.Environment
import           System.FilePath

main = do
  (inputPath:sigmaStr:_) <- getArgs
  (ImageRepa _ img) <- readImageRepa inputPath True
  let (Z :. channels :. cols :. rows) = extent img
      sigma = read sigmaStr :: Double
      folderPath = "output/test/EquivarianceDiagram"
  createDirectoryIfMissing True folderPath
  lock <- getFFTWLock
  (plan1, imgF) <-
    dft1dGPlan lock emptyPlan [channels, cols, rows] [1, 2] .
    VU.convert . VU.map (:+ 0) . toUnboxed $
    img
  (plan2, _) <- idft1dGPlan lock plan1 [channels, cols, rows] [1, 2] imgF
  (plan, gaussianFilterF) <-
    gaussian2DFilter plan2 (Gaussian2DParams sigma rows cols)
  convolvedImg <- convolveGaussian2D plan gaussianFilterF . R.map (:+ 0) $ img
  plotImageRepa (folderPath </> (takeBaseName inputPath) L.++ "_out.png") .
    ImageRepa 8 . computeS . R.map magnitude $
    convolvedImg


-- A simple example testing arithmetic functions.
{-# LANGUAGE FlexibleInstances #-}
import           Control.Monad
import           Data.Complex
import           Data.Either
import           Data.Matrix
import           Data.Ord
import           Data.Vector                (Vector)
import           Numeric.Natural
import           QuickSpec                  hiding (A)
import           QuickSpec.Internal         hiding (A)
import qualified QuickSpec.Internal.Haskell as Haskell
import           Test.QuickCheck
import           Test.QuickCheck.Instances
import           Twee.Pretty

dim = 5

type A = Rational
type T = Matrix A

instance Ord a => Ord (Matrix a) where
  compare = comparing toLists

instance (RealFrac a, Ord a) => Ord (Complex a) where
  compare = comparing (\(x :+ y) -> (truncate (x*1000), truncate (y*1000)))

--instance Ord Cyclotomic where
--  compare = comparing show
--
--instance Arbitrary Cyclotomic where
--  arbitrary = sized arb
--    where
--      arb n =
--        frequency $
--          [(5, fromInteger  <$> arbitrary),
--           (5, fromRational <$> arbRat),
--           (5, gaussianRat <$> arbRat <*> arbRat),
--           (5, polarRat <$> arbRat <*> arbRat),
--           (5, e <$> fmap (succ . abs) arbitrary),
--           (3, sqrtInteger <$> arbitrary),
--           (3, sqrtRat <$> arbRat),
--           (2, sinRev <$> arbRat),
--           (2, cosRev <$> arbRat),
--           (2, return i)] ++
--          concat
--            [ [(5, conj <$> arb (n-1)),
--               (3, real <$> arb (n-1)),
--               (3, imag <$> arb (n-1))]
--            | n > 0 ] ++
--          concat
--            [ let arb2 = arb (n `div` 2) in
--              [(n*5, (+) <$> arb2 <*> arb2),
--               (n*5, (*) <$> arb2 <*> arb2)]
--            | n > 0 ]
--      arbRat = sized $ \n -> do
--        x <- choose (-(n `min` 3), n `min` 3)
--        d <- fmap (succ . abs) (choose (0, n `min` 5))
--        n <- choose (0, d-1)
--        return (fromIntegral x + fromIntegral n / fromIntegral d)

instance Arbitrary T where
  arbitrary =
    makeHermitian <$> frequency [
--(1, return $ zero dim dim),
--      (1, return $ identity dim),
--      (1, return $ negate $ identity dim),
--      (1, genDiagonal),
      (5, fromList dim dim <$> infiniteListOf genA)]

genA :: Gen A
genA = fromInteger <$> arbitrary

genDiagonal :: Gen T
genDiagonal = diagonalList dim 0 <$> infiniteListOf genA

isDiagonal :: T -> Bool
isDiagonal m = m == diagonal 0 (getDiag m)

makeUpperTriangular :: T -> T
makeUpperTriangular = mapPos (\(i, j) x -> if i > j then 0 else x)

isUpperTriangular :: T -> Bool
isUpperTriangular m = makeUpperTriangular m == m

genUpperTriangular :: Gen T
genUpperTriangular = makeUpperTriangular <$> arbitrary

makeLowerTriangular :: T -> T
makeLowerTriangular = mapPos (\(i, j) x -> if i > j then 0 else x)

isLowerTriangular :: T -> Bool
isLowerTriangular m = makeLowerTriangular m == m

genLowerTriangular :: Gen T
genLowerTriangular = makeLowerTriangular <$> arbitrary

makeHermitian :: T -> T
makeHermitian m = m + {-fmap conjugate-} transpose m

isHermitian :: T -> Bool
isHermitian m = m == {-fmap conjugate-} transpose m

genHermitian :: Gen T
genHermitian = makeHermitian <$> arbitrary

instance Floating Rational where
  sqrt = error "sqrt"
instance RealFloat Rational

newtype RowV = RowV {unRowV :: T} deriving (Eq, Ord)
newtype ColV = ColV {unColV :: T} deriving (Eq, Ord)

instance Arbitrary RowV where
  arbitrary = RowV <$> oneof [return $ zero 1 dim, fromList 1 dim <$> infiniteList]
instance Arbitrary ColV where
  arbitrary = (ColV <$> transpose) . unRowV <$> arbitrary

times :: T -> T -> T
times x y = scaleMatrix (1/2) (x*y+y*x)

pow :: Natural -> T -> T
pow 1 t = t
pow n t | n > 0 = times t (pow (n-1) t)

main = quickSpec [
  series [
    background
    [withMaxTermSize 7,
     -- Good pruning settings for Hard Maths (TM)
     --withPruningTermSize 9,
     arith (Proxy :: Proxy A),
     arith (Proxy :: Proxy Natural) `without` ["0"],
     con "*" ((*) :: Natural -> Natural -> Natural),
     monoType (Proxy :: Proxy T),
     con "0" (zero dim dim :: T),
     con "I" (identity dim :: T),
     con "+" ((+) :: T -> T -> T),
     vars ["A", "B", "C"] (Proxy :: Proxy T),
     vars ["n", "m", "l"] (Proxy :: Proxy Natural),
     instFun (arbitrary `suchThat` (\n -> n > 0 && n < 4) :: Gen Natural)],
     --con "*" ((*) :: T -> T -> T),
     --con "-" (negate :: T -> T)],
    --[(con "&&" (&&))],
    --[predicateGen "hermitian" isHermitian (\() -> liftM2 (,) genHermitian (return ()))],
    con "∘" times,
    con "^" (flip pow)]]
{-    [--con "perm" (permMatrix dim :: Int -> Int -> T),
     con "transpose" (transpose :: T -> T),
     --con "inverse" (either (const Nothing) Just . (inverse :: T -> Either String T)),
     con "trace" (trace :: T -> A),
     con "diagProd" (diagProd :: T -> A),
     con "det" (detLU :: T -> A)],
     {-predicateGen "isDiagonal" isDiagonal (\() -> liftM2 (,) genDiagonal (return ())),
     predicateGen "isUpperTriangular" isUpperTriangular (\() -> liftM2 (,) genUpperTriangular (return ())),
     predicateGen "isLowerTriangular" isLowerTriangular (\() -> liftM2 (,) genLowerTriangular (return ()))],-}
    [monoType (Proxy :: Proxy RowV),
     con "0R" (RowV (zero 1 dim)),
     con "+R" (\(RowV x) (RowV y) -> RowV (x+y)),
     con "*R" (\(RowV x) y -> RowV (x*y)),
     con "-R" (\(RowV x) -> RowV (negate x))],
    [con "transposeR" (\(RowV x) -> ColV (transpose x)),
     con "detR" (detLU . unRowV)],
    [monoType (Proxy :: Proxy ColV),
     con "0C" (ColV (zero dim 1)),
     con "+C" (\(ColV x) (ColV y) -> ColV (x+y)),
     con "*C" (\(ColV x) y -> ColV (x*y)),
     con "-C" (\(ColV x) -> ColV (negate x))],
    [con "transposeC" (\(ColV x) -> RowV (transpose x)),
     con "detC" (detLU . unColV)]]]-}


module Quantum
  ( QRegister,
    QGate,
    C,
    qbitsNum,
    dimensionsNum,
    qbit,
    qregister,
    qbitL,
    qregisterL,
    makeGateC,
    makeGate,
    (|>),
    (|+>),
    measure,
  )
where

import Data.Complex
import Data.List (sort)
import Matrix
import System.Random (randomRIO)
import Utils

type C = Complex Double

data QRegister = QRegister
  { qbitsNum :: Int,
    states :: Matrix C
  }

instance Show QRegister where
  show (QRegister n xm) = show xm

data QGate = QGate
  { dimensionsNum :: Int,
    operator :: Matrix C
  }

instance Show QGate where
  show (QGate n xm) = show xm

class Entanglable a where
  entangle :: a -> a -> a

instance Entanglable QRegister where
  entangle (QRegister n1 xm) (QRegister n2 ym) = QRegister n zm
    where
      n = n1 + n2
      zm = kron xm ym

instance Entanglable QGate where
  entangle (QGate n1 xm) (QGate n2 ym) = QGate n zm
    where
      n = n1 + n2
      zm = kron xm ym

qbitL :: Int -> [C]
qbitL 0 = [1 :+ 0, 0 :+ 0]
qbitL 1 = [0 :+ 0, 1 :+ 0]
qbitL _ = error "Expecting |0> or |1>"

qbit :: Int -> QRegister
qbit = QRegister 1 . vectorV . qbitL

applyGate :: QGate -> QRegister -> QRegister
applyGate (QGate n1 xm) (QRegister n2 ym)
  | n1 /= n2 = error "Dimensions of gate and number of qbits in register disagree"
  | otherwise = QRegister n2 (mult xm ym)

qregister :: [Int] -> QRegister
qregister = foldl1 entangle . map qbit

qregisterL :: [Int] -> [C]
qregisterL = concat . toList . states . qregister

makeGateC :: Matrix C -> QGate
makeGateC xm
  | r /= c = error "Gate matrix must be square"
  | not . powerOfTwo $ r = error "Gate matrix dimensions must be power of two"
  | otherwise = QGate n xm
  where
    r = rowsNum xm
    c = colsNum xm
    n = log2 r
    powerOfTwo n = odd n || (elem n . take (log2 n + 1) . map (2 ^) $ [1 ..])
    log2 = fromIntegral . truncate . logBase 2.0 . fromIntegral

makeGate :: Matrix Double -> QGate
makeGate = makeGateC . fmap (:+ 0.0)

infixl 7 |>

(|>) :: QRegister -> QGate -> QRegister
(|>) = flip applyGate

infixl 8 |+>

(|+>) :: Entanglable a => a -> a -> a
(|+>) = entangle

measure :: QRegister -> IO Int
measure (QRegister n xm) = do
  rnd <- randomRIO (0, 1) :: IO Double
  let state = length . takeWhile (< rnd) $ probs
  return (state - 1)
  where
    probs = sort . (\xs -> 0.0 : xs ++ [1]) . scanl1 (+) . map ((^ 2) . magnitude) . concat . toList $ xm


module Space where

import Data.Set
import qualified Data.Set as S

import Data.Vector
import qualified Data.Vector as V

import Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra as M

-- INITIALISATION ---------------------------------------------------------------------------------

-- The data Space (short for simplicial complex) consist of a set of simplices with
-- the following property: all the faces of all the simplices in the set must also
-- belong to the set. A simplex is implemented as a Set for efficiency.
-- For convenience, the set of simplices is split into a list according to the
-- dimension. The defining condition is unchecked as, in practice, the helper
-- function allSimplices (below) is always used to iniatialise the data. The output
-- of the helper function is always valid.

data Space t = Space { subsimplices :: [Set (Set t)] } deriving (Show)

-- This helper function takes a set of simplices as input and returns a valid
-- simplicial complex. The input consists of all the highest dimensional simplices
-- which make up the space. The function takes all their boundaries and sorts
-- them by dimension. The lowest dimension is -1 and it is the dimension of the
-- empty set. It is considered a simplex.
allSimplices :: (Ord t) => [[t]] -> [Set (Set t)]
allSimplices [] = error "This space has no faces. For empty space input '[[]]'."
allSimplices faces =

    let associatedSet = S.fromList (Prelude.map (S.fromList) faces)
        powerFaces = S.map (S.powerSet) associatedSet
        allMixedFaces = S.foldl (S.union) S.empty powerFaces
        lengthFunctions = [ \xs -> (S.size xs == n) | n <- [0,1..] ]

    in  sieveAndForget lengthFunctions allMixedFaces

sieveAndForget :: [t -> Bool] -> Set t -> [Set t]
sieveAndForget (f:fs) xs
    | (S.size xs == 0) = []
    | otherwise        = p1 : (sieveAndForget fs p2)
        where (p1,p2) = S.partition (f) xs

-- HOMOLOGY ---------------------------------------------------------------------------------------

-- This function only computes the free rank of the homology groups. It does not give
-- any information about the torsion subgroups. These require knowledge of the
-- Smith decomposition of a matrix, an algorithm which has never been developed in
-- Haskell (over a PID). To compute the free rank, it is enough to know the rank
-- of the boundary operators (through the package Numeric.LinearALgebra)

-- This function returns the homology numbers of a space from 0 up to its dimension.
-- All the other homology numbers must be zero.
homology :: (Ord t) => Space t -> [Int]
homology (Space { subsimplices = xs }) =
    let operators = chainComplex (Space { subsimplices = xs })
        dimensionDomains = Prelude.map (M.cols) operators
        dimensionImages = Prelude.map (M.rank) (Prelude.tail operators)
        dimensionAugmentedImages = 0:dimensionImages
        dimensionShiftedImages = dimensionImages Prelude.++ [0]
        dimensionKernels = Prelude.zipWith (-) dimensionDomains dimensionAugmentedImages

    in  Prelude.zipWith (-) dimensionKernels dimensionShiftedImages

-- Very small variation on homology. It is used to understand if a space is connected.
reducedHomology :: (Ord t) => Space t -> [Int]
reducedHomology space = reduced
    where reduced = (x-1):xs
          (x:xs) = homology space

-- Computes all boundary operators in a chain complex
chainComplex :: (Ord t) => Space t -> [M.Matrix R]
chainComplex (Space { subsimplices = xs }) =

    let listOfSimplices = Prelude.map (S.toList) xs
        domainListOfSimplices = Prelude.tail listOfSimplices
        imageListOfSimplices = Prelude.init listOfSimplices
        domainVectorOfSimplices = Prelude.map (V.fromList) domainListOfSimplices
        imageVectorOfSimplices = Prelude.map (V.fromList) imageListOfSimplices
        domainImage = Prelude.zip domainVectorOfSimplices imageVectorOfSimplices

    in  Prelude.map (boundaryOperator) domainImage

-- Computes boundary operator for all simplices of a fixed dimension.
boundaryOperator :: (Ord t) => (V.Vector (Set t) , V.Vector (Set t))-> M.Matrix R
boundaryOperator (domain , image) = convertToMatrix vectorMatrix
    where vectorMatrix = V.map (f) domain
          f = \x -> (piecewiseBoundaryOperator x image)

-- Computes boundary operator of a single simplex.
piecewiseBoundaryOperator :: (Ord t) => Set t -> V.Vector (Set t) -> V.Vector Int
piecewiseBoundaryOperator element image =
    let f imageFace
              | (what == False) = 0
              | odd indicator   = 1
              | otherwise       = -1

              where (s1 , what , s2) = S.splitMember imageFace facesOfElement
                    -- Shortcut to compute the sign of a face
                    indicator = S.size element + S.size s1
                    facesOfElement = S.map (g) element
                    g = \y -> (S.delete y element)

    in  V.map (f) image

-- MATRIX CONVERSION ------------------------------------------------------------------------------

-- These algorithms are necessary to convert from V.Vector to M.Matrix.
-- Unfortunately these two data lack compatibility. Furthermore
-- M.Vector cannot store Sets, hence they cannot be used from the beginning.
-- Ideally, the Smith Normal Form should be called directly on the V.Vector
-- so to avoid using M.Matrix all together.

convertToMatrix :: V.Vector (V.Vector Int) -> M.Matrix R
convertToMatrix vectorMatrix = M.matrix s doubleFlattenedListMatrix
    where s = V.length vectorMatrix
          doubleFlattenedListMatrix = Prelude.map (fromIntegral) flattenedListMatrix
          flattenedListMatrix = Prelude.foldl (Prelude.++) [] listMatrix
          listMatrix = Prelude.map (V.toList) listVectorMatrix
          listVectorMatrix = V.toList transposedVectorMatrix
          transposedVectorMatrix = transpose vectorMatrix

transpose :: V.Vector (V.Vector Int) -> V.Vector (V.Vector Int)
transpose matrix = V.foldr (g) emptyVector matrix
    where emptyVector = V.replicate s V.empty
          s = V.length ((V.!) matrix 0)
          g = \x y -> V.zipWith (V.cons) x y

-- BASIC FUNCTIONS --------------------------------------------------------------------------------

-- These functions are basic as they do not rely on computing homology.

-- The function returns the dimension of the highest dimensional simplices.
-- 2 is subtracted to account for the simplices of dimension 0 and -1.
dimension :: (Ord t) => Space t -> Int
dimension (Space { subsimplices = xs }) = Prelude.length xs - 2

-- This function returns the number of simplices of a given dimension.
numberOfSimplices :: (Ord t) => Space t -> Int -> Int
numberOfSimplices (Space { subsimplices = xs }) dim
    | dim < -1 || dim > maxDim = 0
    | otherwise                   = S.size simplicesDim

    where maxDim = dimension (Space { subsimplices = xs })
          simplicesDim = xs !! index
          index = dim + 1

eulerCharacteristic :: (Ord t) => Space t -> Int
eulerCharacteristic (Space { subsimplices = xs }) =
    alternatingSum (Prelude.tail lengths)

        where alternatingSum :: (Num t) => [t] -> t
              alternatingSum [] = error "Not defined."
              alternatingSum [x] = x
              alternatingSum (x:xs) = x - alternatingSum xs

              lengths = (Prelude.map (S.size) xs)

-- ADVANCED FUNCTIONS -----------------------------------------------------------------------------

-- These functions all rely on homology.

isConnected :: (Ord t) => Space t -> Bool
isConnected space = (zeroHomology == 1)
    where zeroHomology = Prelude.head (homology space)

-- This function can tell if a space is not homotopic to a topological manifold.
-- Note that, for the reasons above, it only checks Poincaré duality up to torsion.
rationalPoincareDuality :: (Ord t) => Space t -> Bool
rationalPoincareDuality space = (homNumbers == revHomNumbers)
    where revHomNumbers = Prelude.reverse homNumbers
          homNumbers = homology space

-- If this function returns False then the answer is definitely No. If it returns
-- True, then the algorthm is inconcludent. The function first checks the Euler
-- characteristic as this is a less expensice invariant to compute, albeit coarser.
-- If it is the same, homology ranks are computed.
couldItBeContractible :: (Ord t) => Space t -> Bool
couldItBeContractible space = (euler == 1) && (isZero reduced)
    where euler = eulerCharacteristic space
          reduced = reducedHomology space

-- Simple helper function to check if all the entries in a list are zero
isZero :: [Int] -> Bool
isZero [] = error "Empty list."
isZero [0] = True
isZero (x:xs) = (x == 0) && (isZero xs)

-- Again, if the function returns True, no new information is gained. The
-- Euler characteristic is computed first. This algorithm only checks
-- if the homology numbers are the same, however, it is necessary to add
-- enough zeros to the shorter list to equate the lengths.
couldTheyBeHomotopic :: (Ord t) => Space t -> Space t -> Bool
couldTheyBeHomotopic space1 space2 = (euler1 == euler2) && (hom1 == hom2)
    where euler1 = eulerCharacteristic space1
          euler2 = eulerCharacteristic space2
          hom1 = tempHom1 Prelude.++ zeroList1
          hom2 = tempHom2 Prelude.++ zeroList2
          tempHom1 = homology space1
          tempHom2 = homology space2
          zeroList1 = [ 0 | x <- [1..s1]]
          zeroList2 = [ 0 | x <- [1..s2]]
          s1 = maxLength - length1
          s2 = maxLength - length2
          length1 = Prelude.length tempHom1
          length2 = Prelude.length tempHom2
          maxLength = Prelude.maximum [length1, length2]


{-# LANGUAGE FlexibleInstances #-}

module Main where

import Control.Monad
import Graphics.Rendering.Chart
import Data.Colour
import Data.Colour.Names
import Data.Default.Class
import Graphics.Rendering.Chart.Backend.Cairo
import Control.Lens
import Data.Complex

setLinesBlue :: PlotLines a b -> PlotLines a b
setLinesBlue = plot_lines_style  . line_color .~ opaque blue

chart :: Renderable ()
chart = toRenderable layout
  where
    sumf = plot_lines_values .~ [[ (x,(s x)) | x <- [1.3,(1.31)..10]]]
              $ plot_lines_style  . line_color .~ opaque black
              $ plot_lines_title .~ "S(a)"
              $ def

    tln = plot_lines_values .~ [[ (x,(taylor1 2.0 x)) | x <- [1.3,(1.4)..3.5]]]
              $ plot_lines_style  . line_color .~ opaque red
              $ plot_lines_title .~ "First taylor"
              $ def

    tln2 = plot_lines_values .~ [[ (x,(taylor2 2.0 x)) | x <- [1.3,(1.4)..4.0]]]
              $ plot_lines_style  . line_color .~ opaque orange
              $ plot_lines_title .~ "Second taylor"
              $ def

    tln3 = plot_lines_values .~ [[ (x,(taylor3 2.0 x)) | x <- [1.3,(1.4)..3.2]]]
              $ plot_lines_style  . line_color .~ opaque blue
              $ plot_lines_title .~ "Third taylor"
              $ def

    layout = layout_title .~ "Amplitude Modulation"
           $ layout_plots .~ [toPlot sumf, toPlot tln, toPlot tln2, toPlot tln3]
           $ def

renderChart :: IO (PickFn ())
renderChart = renderableToFile def "example1_big.png" chart

approxS :: Int -> Double -> Double
approxS i a = sum $ map (f.fromIntegral) [1..i]
  where f n = n**(-a)

class Abs a where
  absoluteValue :: a -> Double

instance Abs Double where
  absoluteValue = abs

instance Abs (Complex Double) where
  absoluteValue = magnitude

untilConverged :: (Abs a, Floating a, Num a) => Double -> [a] -> a
untilConverged epsilon (x1:xs) = go x1 xs
  where go last (y:ys) =
         if absoluteValue (last - y) < epsilon
           then y
           else go y ys

sumSeries :: (Abs a, Floating a, Num a) => (Int -> a) -> a
sumSeries seriesGen = untilConverged 0.000001 $ tail $ scanl1 (+) $ map seriesGen [1..]

s :: Double -> Double
s a = sumSeries $ \n-> (fromIntegral n)**(-a)

sDerivative :: (Abs a, Floating a, Num a) => Int -> a -> a
sDerivative d a = sign * (sumSeries $ \nInt ->
    let n = fromIntegral nInt in
    (fromIntegral nInt)**(-a) * (realToFrac $ log n)^d)
  where sign = if even d then 1 else (-1)

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

taylorCoeffients :: (Fractional a, Num a) => (Int -> a -> a) -> a -> [a]
taylorCoeffients f a0 = map g [0..]
  where g n = (f n a0) / fromIntegral (fact n)

taylor1 :: Double -> Double -> Double
taylor1 a0 a = s a0 + (a - a0) * sDerivative 1 a0

taylor2 :: Double -> Double -> Double
taylor2 a0 a = s a0 + (a - a0) * sDerivative 1 a0 + 0.5 * (a - a0)^2 * sDerivative 2 a0

taylor3 :: Double -> Double -> Double
taylor3 a0 a = s a0 + (a - a0) * sDerivative 1 a0 + 0.5 * (a - a0)^2 * sDerivative 2 a0 + (1.0/6.0) * (a - a0)^3 * sDerivative 3 a0

main :: IO ()
main = void $ renderChart


module EKF
  ( KF
  , kf_x
  , kf_p
  , kf_t

  , makeFilter
  , predict
  , update
  ) where

import Prelude hiding ((<>))

import App (Sensor (..))

import Numeric.LinearAlgebra

data KF = KF
  { kf_x :: Vector Double
  , kf_p :: Matrix Double
  , kf_t :: Word
  } deriving Show

makeFilter (Laser px py t) =
  KF { kf_x = 4 |> [px,
                    py,
                    0,
                    0]
     , kf_p = (4><4) [1, 0,    0,    0,
                      0, 1,    0,    0,
                      0, 0, 1000,    0,
                      0, 0,    0, 1000]
     , kf_t = t }

makeFilter (Radar rho phi rho' t) =
  KF { kf_x = 4 |> [rho  * cos phi,
                    rho  * sin phi,
                    rho' * cos phi,
                    rho' * sin phi]
     , kf_p =  (4><4) [1, 0,    0,    0,
                       0, 1,    0,    0,
                       0, 0, 1000,    0,
                       0, 0,    0, 1000]
     , kf_t = t }

predict :: Word -> KF -> KF
predict t kf =
  KF { kf_x = x
     , kf_p = p
     , kf_t = t }
  where
    -- Time inside the filter.
    filterTime =
      kf_t kf

    -- Delta-time. Essentially a timestep.
    dt =
      (fromIntegral (t - filterTime)) / 1000000.0
    dt2 = dt ** 2
    dt3 = dt ** 3
    dt4 = dt ** 4

    -- Acceleration noise tuning parameters
    ax = 3
    ay = 3

    noiseAx = ax ** 2
    noiseAy = ay ** 2

    -- State transition matrix
    f = (4><4) [1, 0, dt,  0,
                0, 1,  0, dt,
                0, 0,  1,  0,
                0, 0,  0,  1]

    -- Noise covariance matrix
    q = (4><4) [-- Row 1
                noiseAx * 0.25 * dt4, 0, noiseAx * 0.5 * dt3, 0,
                -- Row 2
                0, noiseAy * 0.25 * dt4, 0, noiseAy * 0.5 * dt3,
                --  Row 3
                noiseAx *0.5 * dt3, 0, noiseAx * dt2, 0,
                -- Row 4
                0, noiseAy * 0.5 * dt3, 0, noiseAy * dt2]

    -- Noise vector
    v  = 4 |> [ax * 0.5 * dt2,
               ay * 0.5 * dt2,
               ax * dt,
               ay * dt]

    -- Estimated step based on constant velocity (CV) motion model
    x = f #> kf_x kf + v
    p = f <> kf_p kf <> (tr f) + q

update (Laser px py t) kf =
  KF { kf_x = x'
     , kf_p = p'
     , kf_t = t }
  where
    kf' = predict t kf

    -- Measurment transition matricies
    h = (2><4) [1, 0, 0, 0,
                0, 1, 0, 0]
    r = (2><2) [0.0225, 0.0000,
                0.0000, 0.0225]
    y = 2 |> [px, py] - (h #> kf_x kf')

    -- Kalman Filter Equations
    i  = ident 4
    ht = tr h
    s  = h <> kf_p kf' <> ht + r
    si = inv s
    k  = kf_p kf' <> ht <> si

    -- Estimation based on kalman filter gain
    x' = kf_x kf' + (k #> y)
    p' = (i - k <> h) <> kf_p kf'

update (Radar rho phi rho' t) kf =
  KF { kf_x = x'
     , kf_p = p'
     , kf_t = t }
  where
    kf' = predict t kf

    -- Calculate the Jacobian matrix values
    px = kf_x kf' ! 0
    py = kf_x kf' ! 1
    vx = kf_x kf' ! 2
    vy = kf_x kf' ! 3

    c1 = px ** 2 + py ** 2
    c2 = sqrt c1
    c3 = c1 * c2

    zRho  = c2;
    zPhi  = (atan2 py px)
    zRho' = (px*vx + py*vy) / zRho

    -- Measurment transition matricies
    h = (3><4) [-- Row 1
                px / c2, py / c2, 0, 0,
                -- Row 2
                -py / c1, px / c1, 0, 0,
                -- Row 3
                py*(vx * py - vy * px)/c3,
                px*(vy * px - vx * py)/c3,
                px/c2,
                py/c2]
    r = (3><3) [0.09, 0.0000, 0.00,
                0.00, 0.0009, 0.00,
                0.00, 0.0000, 0.09]
    z = 3 |> [rho, phi, rho'] - 3 |> [zRho, zPhi, zRho']
    y = 3 |> [z ! 0, atan2 (sin (z ! 1)) (cos (z ! 1)), z ! 2]

    -- Kalman filter equations
    i  = ident 4
    ht = tr h
    s  = h <> kf_p kf' <> ht + r
    si = inv s
    k  = kf_p kf' <> ht <> si

    -- Estimation based on kalman filter gain
    x' = kf_x kf' + (k #> y)
    p' = (i - k <> h) <> kf_p kf'


import Criterion.Main
import Data.Default.Class
import qualified Data.Vector.Unboxed as U
import Text.Printf

import Numeric.SpecFunctions
import Numeric.Polynomial
import Numeric.RootFinding

-- Uniformly sample logGamma performance between 10^-6 to 10^6
benchmarkLogGamma logG =
  [ bench (printf "%.3g" x) $ nf logG x
  | x <- [ m * 10**n | n <- [ -8 .. 8 ]
                     , m <- [ 10**(i / tics) | i <- [0 .. tics-1] ]
         ]
  ]
  where tics = 3
{-# INLINE benchmarkLogGamma #-}

-- Power of polynomial to be evaluated (In other words length of coefficients vector)
coef_size :: [Int]
coef_size = [ 1,2,3,4,5,6,7,8,9
            , 10,    30
            , 100,   300
            , 1000,  3000
            , 10000, 30000
            ]
{-# INLINE coef_size #-}

-- Precalculated coefficients
coef_list :: [U.Vector Double]
coef_list = [ U.replicate n 1.2 | n <- coef_size]
{-# NOINLINE coef_list #-}

main :: IO ()
main = defaultMain
  [ bgroup "logGamma" $
    benchmarkLogGamma logGamma
  , bgroup "logGammaL" $
    benchmarkLogGamma logGammaL
  , bgroup "incompleteGamma" $
      [ bench (show p) $ nf (incompleteGamma p) p
      | p <- [ 0.1
             , 1,   3
             , 10,  30
             , 100, 300
             , 999, 1000
             ]
      ]
  , bgroup "factorial"
    [ bench (show n) $ nf factorial n
    | n <- [ 0, 1, 3, 6, 9, 11, 15
           , 20, 30, 40, 50, 60, 70, 80, 90, 100
           ]
    ]
  , bgroup "incompleteBeta"
    [ bench (show (p,q,x)) $ nf (incompleteBeta p q) x
    | (p,q,x) <- [ (10,      10,      0.5)
                 , (101,     101,     0.5)
                 , (1010,    1010,    0.5)
                 , (10100,   10100,   0.5)
                 , (100100,  100100,  0.5)
                 , (1001000, 1001000, 0.5)
                 , (10010000,10010000,0.5)
                 ]
    ]
  , bgroup "log1p"
      [ bench (show x) $ nf log1p x
      | x <- [ -0.9
             , -0.5
             , -0.1
             ,  0.1
             ,  0.5
             ,  1
             ,  10
             ,  100
             ] :: [Double]
      ]
  , bgroup "sinc" $
        bench "sin" (nf sin (0.55 :: Double))
      : [ bench (show x) $ nf sinc x
        | x <- [0, 1e-6, 1e-3,  0.5]
        ]
  , bgroup "poly"
      $  [ bench ("vector_"++show (U.length coefs)) $ nf (\x -> evaluatePolynomial x coefs) (1 :: Double)
         | coefs <- coef_list ]
      ++ [ bench ("unpacked_"++show n) $ nf (\x -> evaluatePolynomialL x (map fromIntegral [1..n])) (1 :: Double)
         | n <- coef_size ]
  , bgroup "RootFinding"
    [ bench "ridders sin" $ nf (ridders       def (0,pi/2))     (\x -> sin x - 0.525)
    , bench "newton sin"  $ nf (newtonRaphson def (0,1.2,pi/2)) (\x -> (sin x - 0.525,cos x))
    ]
  ]


module DistanceMatrixRank where

import Data.List as L
import Data.Vector as V
import Data.Matrix as M

import Statistics.Sample

data DistanceMatrix = DistanceMatrix
  {
    distanceMatrix :: Matrix Double,
    origins :: [String],
    destinations :: [String]
  }
  deriving (Show)

rankDestinations :: DistanceMatrix -> [(String, Double)]
rankDestinations m = L.zip (destinations m) (L.map stdDev $ columns (distanceMatrix m))

columns :: Matrix a -> [Vector a]
columns m = L.map (flip M.getCol m) [1..cols]
  where
    cols = M.ncols m


module Main where

import           Criterion
import           Criterion.Main.Options
import           Criterion.Types
import           Data.Complex
import           Data.Vector
import qualified Numeric.FFT            as FFT

tstvec :: Int -> Vector (Complex Double)
tstvec sz = generate sz (\i -> let ii = fromIntegral i
                               in sin (2*pi*ii/1024) + sin (2*pi*ii/511))

main :: IO ()
main = do
  p <- FFT.plan 256
  benchmarkWith (defaultConfig { resamples = 1000 })
    (nf (FFT.fftWith p) (tstvec 256))


module Echo where

import Data.Complex
import Linear.Metric
import Linear.V3

echo :: (Floating a) => (a -> Complex a) -> V3 a -> V3 a -> a -> Complex a
echo signal carrier obj t = signal (t - 2 * d / c)
  where
    d = distance carrier obj
    c = 299792458


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators    #-}
module FokkerPlanck.Interpolation where

import           Data.Array.Repa as R
import           Data.Complex
import           Data.List       as L
import           Filter.Utils
import           Types

{-# INLINE radialCubicInterpolation #-}
radialCubicInterpolation ::
     (R.Source r1 Double, R.Source r2 (Complex Double), Shape sh)
  => R.Array r1 (sh :. Int) Double
  -> Double
  -> R.Array r2 (sh :. Int :. Int) (Complex Double)
  -> R.Array D (sh :. Int :. Int) (Complex Double)
radialCubicInterpolation radialArr scaleFactor inputArr =
  let (_ :. n) = extent radialArr
      (_ :. nx :. ny) = extent inputArr
   in if scaleFactor == 0
        then error "radialCubicInterpolation: scaleFactor = 0"
        else R.traverse2 inputArr radialArr const $ \fInput fRadial idx@(sh :. i :. j) ->
               let r =
                     (sqrt . fromIntegral $
                      (i - center nx) ^ 2 + (j - center ny) ^ 2) /
                     scaleFactor
                in if r < 0 || r > (fromIntegral $ n - 1) -- out of boundary. When r = 0, log r makes no sense, and therefore the information at r = 0 should not be used to interpoalte 0 < r <= 1.
                     then 0
                     else if r <= 1 -- linear interpolation
                            then fInput idx *
                                 ((fRadial (sh :. 1) +
                                   (r - 1) *
                                   (fRadial (sh :. 2) - fRadial (sh :. 1))) :+
                                  0)
                            else if r >= (fromIntegral $ n - 2) -- linear interpolation
                                   then fInput idx *
                                        ((fRadial (sh :. n - 2) +
                                          (r - (fromIntegral $ n - 2)) *
                                          (fRadial (sh :. n - 1) -
                                           fRadial (sh :. n - 2))) :+
                                         0)
                                   else let x0 = floor r -- cubic interpolation
                                            p0 = fRadial (sh :. x0 - 1)
                                            p1 = fRadial (sh :. x0)
                                            p2 = fRadial (sh :. x0 + 1)
                                            p3 = fRadial (sh :. x0 + 2)
                                            x = r - fromIntegral x0
                                         in fInput idx *
                                            (((-0.5 * p0 + 1.5 * p1 - 1.5 * p2 +
                                               0.5 * p3) *
                                              x ^ 3 +
                                              (p0 - 2.5 * p1 + 2 * p2 - 0.5 * p3) *
                                              x ^ 2 +
                                              (-0.5 * p0 + 0.5 * p2) * x +
                                              p1) :+
                                             0)

{-# INLINE cubicInterpolation #-}
cubicInterpolation ::
     (R.Source s Double)
  => R.Array s DIM5 Double
  -> Double
  -> R.Array D DIM4 Double
cubicInterpolation arr r =
  let (Z :. _ :. _ :. _ :. _ :. n) = extent arr
   in R.traverse arr (\(sh :. _) -> sh) $ \f sh@(Z :. tf :. sf :. t0f :. s0f) ->
        if r < 0 || r > (fromIntegral $ n - 1) -- out of boundary. When r = 0, log r makes no sense, and therefore the information at r = 0 should not be used to interpoalte 0 < r <= 1.
          then 0
          else if r <= 1 -- linear interpolation
                 then (f (sh :. 1) + (r - 1) * (f (sh :. 2) - f (sh :. 1)))
                 else if r >= (fromIntegral $ n - 2) -- linear interpolation
                        then (f (sh :. n - 2) +
                              (r - (fromIntegral $ n - 2)) *
                              (f (sh :. n - 1) - f (sh :. n - 2)))
                        else let x0 = floor r -- cubic interpolation
                                 p0 = f (sh :. x0 - 1)
                                 p1 = f (sh :. x0)
                                 p2 = f (sh :. x0 + 1)
                                 p3 = f (sh :. x0 + 2)
                                 x = r - fromIntegral x0
                              in ((-0.5 * p0 + 1.5 * p1 - 1.5 * p2 + 0.5 * p3) *
                                  x ^ 3 +
                                  (p0 - 2.5 * p1 + 2 * p2 - 0.5 * p3) * x ^ 2 +
                                  (-0.5 * p0 + 0.5 * p2) * x +
                                  p1)


module Examples where

import Control.Lens hiding (para)
import Prelude hiding (id, (.))
import System.Random
import Control.Monad

import Numeric.LinearAlgebra.Array
import Numeric.LinearAlgebra.Array.Util

import CategoricDefinitions
import Autodiff.GAD
import Autodiff.D
import Ops
import Para
import TrainUtils
import TensorUtils
import OnesLike

{-
f (p, q) a = (p + q*a)
-}
linRegFn :: (Additive a, Num a) => DType ((a, a), a) a
linRegFn = addC . (id `x` mulC) . assocL

-- returns (inp, out) pair we want to learn
sampleData :: IO (Double, Double)
sampleData = do
    x <- randomIO :: IO Double
    return (x, 3 + 7*x)

run :: IO ()
run = do
    initialParams <- randomIO :: IO (Double, Double)
    let initialLearner = Learner initialParams (Para linRegFn) sgd
        sampler = zip3 [0..] (repeat sampleData) (repeat sqDiff)

    finalLearner <- foldM trainStepWithCost initialLearner (take 10000 sampler)
    putStrLn $ "Starting network parameters:\n" ++ show initialParams
    putStrLn $ "Final network parameters:\n" ++ show (finalLearner ^. p)

    return ()

sampleDataTensor :: IO (Tensor, Tensor)
sampleDataTensor = do
    d <- randomTensor [5, 3] "bf"
    return (d, 3 + 7 * d)

-- multiply two arrays you get and sum the "b" axis
l :: _ => ParaDType (Tensor, Tensor) Tensor Tensor
l = Para $ sumAxes "b" . linRegFn

run1 :: IO _
run1 = do
    p1 <- randomTensor [3, 2] "fo"
    p2 <- randomTensor [2] "o"
    let initialLearner = Learner (p1, p2) l sgd
        sampler = zip3 [0..] (repeat sampleDataTensor) (repeat (sumAxes "b" . sqDiff))

    finalLearner <- foldM trainStepWithCost initialLearner (take 1 sampler)
    putStrLn $ "Starting network parameters:\n" ++ arrShow (p1, p2)
    putStrLn $ "Final network parameters:\n" ++ arrShow (finalLearner ^. p)

    return finalLearner


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
-- |
-- Module      : Graphics.Image.Processing.Complex.Internal
-- Copyright   : (c) Alexey Kuleshevich 2016-2018
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.Image.Processing.Complex.Internal
  (
  -- * Pixel
    Complex(..)
  -- ** Rectangular form
  , (+:)
  , realPart
  , imagPart
  -- ** Polar form
  , mkPolar
  , cis
  , polar
  , magnitude
  , phase
  -- ** Conjugate
  , conjugate
  -- * Image
  -- ** Rectangular form
  , (+:!)
  , realPartI
  , imagPartI
  -- ** Polar form
  , mkPolarI
  , cisI
  , polarI
  , magnitudeI
  , phaseI
  -- ** Conjugate
  , conjugateI
  -- * Re-export
  ) where

import Control.Applicative
import Data.Complex (Complex(..))
import qualified Data.Complex as C
import Graphics.Image.Internal
import Prelude hiding (map, zipWith)

infix 6 +:!, +:

(+:) :: Applicative (Color cs) => Pixel cs e -> Pixel cs e -> Pixel cs (Complex e)
(+:) = liftA2 (:+)
{-# INLINE (+:) #-}

realPart :: Functor (Color cs) => Pixel cs (Complex e) -> Pixel cs e
realPart = fmap C.realPart
{-# INLINE realPart #-}

imagPart ::  Functor (Color cs) => Pixel cs (Complex e) -> Pixel cs e
imagPart = fmap C.imagPart
{-# INLINE imagPart #-}

mkPolar :: (Applicative (Color cs), Floating e) => Pixel cs e -> Pixel cs e -> Pixel cs (Complex e)
mkPolar = liftA2 C.mkPolar
{-# INLINE mkPolar #-}

cis :: (Functor (Color cs), Floating e) => Pixel cs e -> Pixel cs (Complex e)
cis = fmap C.cis
{-# INLINE cis #-}

polar :: (Functor (Color cs), RealFloat e) => Pixel cs (Complex e) -> (Pixel cs e, Pixel cs e)
polar z = (magnitude z, phase z)
{-# INLINE polar #-}

magnitude :: (Functor (Color cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs e
magnitude = fmap C.magnitude
{-# INLINE magnitude #-}

phase :: (Functor (Color cs), RealFloat e) => Pixel cs (Complex e) -> Pixel cs e
phase = fmap C.phase
{-# INLINE phase #-}

conjugate :: (Functor (Color cs), Num e) => Pixel cs (Complex e) -> Pixel cs (Complex e)
conjugate = fmap C.conjugate
{-# INLINE conjugate #-}

-- | Construct a complex image from two images representing real and imaginary parts.
--
-- >>> frog <- readImageRGB "images/frog.jpg"
-- >>> frog !+! 0
-- <Image RGB Complex Double: 200x320>
-- >>> frog !+! frog
-- <Image RGB Complex Double: 200x320>
--
(+:!) :: (ColorModel cs e, ColorModel cs (Complex e)) =>
         Image cs e -> Image cs e -> Image cs (Complex e)
(+:!) = zipWith (+:)
{-# INLINE (+:!) #-}

-- | Extracts the real part of a complex image.
realPartI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
             Image cs (Complex e) -> Image cs e
realPartI = map realPart
{-# INLINE realPartI #-}

-- | Extracts the imaginary part of a complex image.
imagPartI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
             Image cs (Complex e) -> Image cs e
imagPartI = map imagPart
{-# INLINE imagPartI #-}

-- | Form a complex image from polar components of magnitude and phase.
mkPolarI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
            Image cs e -> Image cs e -> Image cs (Complex e)
mkPolarI = zipWith mkPolar
{-# INLINE mkPolarI #-}

-- | @'cisI' t@ is a complex image with magnitude 1 and phase t (modulo @2*'pi'@).
cisI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
        Image cs e -> Image cs (Complex e)
cisI = map cis
{-# INLINE cisI #-}

-- | The function @'polar''@ takes a complex image and returns a (magnitude, phase)
-- pair of images in canonical form: the magnitude is nonnegative, and the phase
-- in the range @(-'pi', 'pi']@; if the magnitude is zero, then so is the phase.
polarI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
          Image cs (Complex e) -> (Image cs e, Image cs e)
polarI !zImg = (magnitudeI zImg, phaseI zImg)
{-# INLINE polarI #-}

-- | The nonnegative magnitude of a complex image.
magnitudeI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
              Image cs (Complex e) -> Image cs e
magnitudeI = map magnitude
{-# INLINE magnitudeI #-}

-- | The phase of a complex image, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
phaseI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
          Image cs (Complex e) -> Image cs e
phaseI = map phase
{-# INLINE phaseI #-}

-- | The conjugate of a complex image.
conjugateI :: (RealFloat e, ColorModel cs e, ColorModel cs (Complex e)) =>
              Image cs (Complex e) -> Image cs (Complex e)
conjugateI = map conjugate
{-# INLINE conjugateI #-}


{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module Utils.Types.Num
  ( NumType(..)
  )
where

import           RIO                     hiding ( toRational )
import           Data.Complex                   ( Complex(..) )
import           Data.Ratio                     ( denominator
                                                , numerator
                                                , (%)
                                                )

data NumType = Complex (Complex Double)
             | Real Double
             | Rational Rational
             | Integer Integer
             deriving(Eq, Ord, Data)

instance Ord a => Ord (Complex a) where
  compare (ar :+ ai) (br :+ bi) = case compare ar br of
    EQ    -> compare ai bi
    other -> other

instance Show NumType where
  show (Integer  num       ) = show num
  show (Rational num) = show (numerator num) <> "/" <> show (denominator num)
  show (Real     num       ) = show num
  show (Complex  (re :+ im)) = show re <> "+" <> show im <> "i"

toComplex :: NumType -> NumType
toComplex (    Integer  num) = Complex $ fromInteger num :+ 0
toComplex num@(Rational _  ) = toComplex $ toReal num
toComplex (    Real     num) = Complex $ num :+ 0
toComplex num@(Complex  _  ) = num

toReal :: NumType -> NumType
toReal (Integer num) = Real $ fromInteger num
toReal (Rational num) =
  Real $ fromInteger (numerator num) / fromInteger (denominator num)
toReal num@(Real _) = num
toReal _            = undefined

toRational :: NumType -> NumType
toRational (    Integer  num) = Rational $ fromInteger num
toRational num@(Rational _  ) = num
toRational _                  = undefined

instance Num NumType where
  (Integer a) + (Integer b) = Integer $ a + b
  a           + b           = numTypeOp (+) a b
  (Integer a) * (Integer b) = Integer $ a * b
  a           * b           = numTypeOp (*) a b
  abs         = numTypeFun abs
  signum      = numTypeFun signum
  negate      = numTypeFun negate
  fromInteger = Integer

instance Fractional NumType where
  (Integer a) / (Integer b) = Rational $ a % b
  a           / b           = numTypeOp (/) a b
  fromRational = Rational

instance Floating NumType where
  pi    = Real pi
  exp   = realOp exp
  log   = realOp log
  sin   = realOp sin
  cos   = realOp cos
  asin  = realOp asin
  acos  = realOp acos
  atan  = realOp atan
  sinh  = realOp sinh
  cosh  = realOp cosh
  asinh = realOp asinh
  acosh = realOp acosh
  atanh = realOp atanh
  sqrt  = realOp sqrt

realOp :: (Double -> Double) -> NumType -> NumType
realOp op (Real a) = Real $ op a
realOp op a        = realOp op (toReal a)

numTypeOp
  :: (forall  a . (Num a, Fractional a) => a -> a -> a)
  -> NumType
  -> NumType
  -> NumType
numTypeOp op (  Real     a) (Real     b)   = Real $ op a b
numTypeOp op (  Rational a) (Rational b)   = Rational $ op a b
numTypeOp op (  Complex  a) (Complex  b)   = Complex $ op a b
numTypeOp op a@(Complex  _) b              = numTypeOp op a (toComplex b)
numTypeOp op a              b@(Complex _)  = numTypeOp op (toComplex a) b
numTypeOp op a@(Real _)     b              = numTypeOp op a (toReal b)
numTypeOp op a              b@(Real _)     = numTypeOp op (toReal a) b
numTypeOp op a@(Rational _) b              = numTypeOp op a (toRational b)
numTypeOp op a              b@(Rational _) = numTypeOp op (toRational a) b
numTypeOp _  _              _              = undefined

numTypeFun :: (forall  a . Num a => a -> a) -> NumType -> NumType
numTypeFun fun (Integer  a) = Integer $ fun a
numTypeFun fun (Real     a) = Real $ fun a
numTypeFun fun (Rational a) = Rational $ fun a
numTypeFun fun (Complex  a) = Complex $ fun a


module State where

import           Protolude

import           Data.Complex (Complex (..))
import qualified Data.Vector  as V

import           Instruction  (Instruction)

data MachineState =
  MachineState
    { quantumState :: QuantumState
    , instructions :: [Instruction]
    , pc           :: Integer
    , running      :: Bool
    }

type QuantumState = V.Vector (Complex Double)

-- Construct a quantum state given number of qubits
makeQuantumState :: Int -> QuantumState
makeQuantumState len =
  V.generate
    (2 ^ len)
    (\i ->
       if i == 0
         then (1 :+ 0)
         else (0 :+ 0))


{-# LANGUAGE FlexibleContexts #-}

module Learning
    ( ols
    , gls
    , kmeans
    , Centroids
    , Clusters
    , Clustering
    , Initializer
    , average
    , invertMap
    ) where

import Debug.Trace
import Data.Set
import Data.Map.Strict
import Numeric.LinearAlgebra hiding ((!))

--covariance :: (Field n, Ord n, Num (Vector n))
--           => Set (Vector n)
--           -> Matrix n

--covariance set = exxt - uut where
--  exxt = mean $ Prelude.map (\a -> asColumn a <> asRow a) (Data.Set.elems set)
--  uut  = asColumn mu <> asRow mu
--  mu   = average set

-- | Ordinary least squares
ols :: (Field n, Monoid (Matrix n))
    => Matrix n
    -> Vector n
    -> Vector n

ols x y = head $ toColumns $ optimiseMult [l', tr l', tr x, asColumn y] where
  xTx = mTm x
  l   = chol xTx
  l'  = inv l

-- | Generalized least squares
gls :: (Eq n, Num (Vector n), Monoid (Matrix n), Field n)
    => Matrix n
    -> Matrix n
    -> Vector n
    -> Vector n

gls c x y = t where
  x' = tr x
  ic = inv c
  r  = inv (optimiseMult [x', ic, x])
  s  = optimiseMult [r, x, ic]
  t  = s #> y

type Centroids n = Set (Vector n)

type Clusters n = Map (Vector n) (Set (Vector n))

type Clustering n = (Centroids n, Clusters n)

type Initializer m n = Int -> Set (Vector n) -> m (Centroids n)

-- | KMeans clustering algorithm
kmeans :: (Functor m, Container Vector n, Ord n, Num (Vector n), Normed (Vector n), Field n)
       => Initializer m n
       -> Int
       -> Set (Vector n)
       -> m (Clustering n)

-- KMeans runs by first getting an original clustering, then reclustering
-- until a fixed point is reached.
kmeans init k set = (untilEqual fst (recluster k) . (\centroids -> (centroids, reassign centroids set))) <$> init k set

untilEqual :: Eq a => (e -> a) -> (e -> e) -> e -> e
untilEqual g f e = if g e' == g e then e else untilEqual g f e' where
  e' = f e

-- Reclustering means taking each cluster, computing a new centroid based on
-- the contents, and then reassigning each vector to the cluster which has a
-- centroid closest to it.
recluster :: (Container Vector n, Ord n, Fractional n, Normed (Vector n), Num (Vector n), Field n)
          => Int
          -> Clustering n
          -> Clustering n
recluster k (centroids, clusters) = (centroids', reassign centroids' (Data.Set.unions $ Data.Map.Strict.elems clusters)) where
  centroids' = Data.Set.map (\centroid -> average (clusters ! centroid)) centroids

average :: (Container c n, Fractional n, Ord n, Num (c n), Linear n c)
        => Set (c n) -> c n

average set = scale (1 / (fromInteger . toInteger $ Data.Set.size set)) $ Data.Set.foldr (+) one set' where
  (one, set') = Data.Set.deleteFindMin set

-- Reassignment goes like this: given a set of vectors and a set of centroids,
-- return a map from centroid to sets of vectors which have been found to be within this
-- centroid.
reassign :: (Normed (Vector n), Ord n, Eq n, Num (Vector n), Field n)
         => Set (Vector n)
         -> Set (Vector n)
         -> Clusters n

reassign centroids points = invertMap $ fromSet (getClosest centroids) points

-- | Given a map from x to y, returns a map from y to the preimage of y in the
--   original map
invertMap :: (Ord y, Ord x) => Map x y -> Map y (Set x)
invertMap map = st where
  st = fromListWith Data.Set.union [ (e, Data.Set.singleton k)
                                   | k <- ks, e <- es, map ! k == e ]
  es = Data.Map.Strict.elems map
  ks = keys map

getClosest :: (Normed (Vector n), Num (Vector n), Field n, Eq n) => Set (Vector n) -> Vector n -> Vector n
getClosest points point = if Data.Set.null points then error "getClosest" else fst $ Data.Set.foldr
  (\candidate (c, d) -> let d' = norm_2 (candidate - point)
                        in if d > d' then (candidate, d') else (c, d))
  (candidate, norm_2 (candidate - point))
   points' where
  (candidate, points') = Data.Set.deleteFindMin points


{-# LANGUAGE BangPatterns #-}
module FourierSeriesOfPinwheels where

import           Control.Monad.Parallel          as MP
import           Data.Array.IArray               as IA
import           Data.Array.Repa                 as R
import           Data.Array.Repa.IO.Binary
import           Data.Array.Repa.Repr.ForeignPtr
import           Data.Complex
import           Data.List                       as L
import           Data.Vector.Storable            as VS
import           Data.Vector.Unboxed             as VU
import           FokkerPlanck.GreensFunction
import           FokkerPlanck.Histogram
import           Image.IO
import           Pinwheel.Base
import           Pinwheel.FourierSeries2D
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Parallel
import           Utils.Time
import FourierMethod.FourierSeries2D

main = do
  args@(gpuIDsStr:lenStr:deltaStr:inverseDeltaStr:numR2FreqsStr:periodStr:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:coefFilePath:stdStr:batchSizePolarFreqsStr:batchSizeR2FreqsStr:batchSizeR2Str:numThreadStr:_) <-
    getArgs
  let gpuIDs = read gpuIDsStr :: [Int]
      len = read lenStr :: Double
      delta = read deltaStr :: Double
      inverseDelta = read inverseDeltaStr :: Double
      numR2Freqs = read numR2FreqsStr :: Int
      period = read periodStr :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/FourierSeriesOfPinwheels"
      std = read stdStr :: Double
      batchSizePolarFreqs = read batchSizePolarFreqsStr :: Int
      batchSizeR2Freqs = read batchSizeR2FreqsStr :: Int
      batchSizeR2 = read batchSizeR2Str :: Int
  removePathForcibly (folderPath </> "Recons")
  createDirectoryIfMissing True folderPath
  let !numAngularFreqs = 2 * (phiFreq + thetaFreq) + 1
      !angularFreqsCenter = div numAngularFreqs 2
      !numRadialFreqs = 2 * (rhoFreq + scaleFreq) + 1
      !numPoints' = round $ len / delta :: Int
      !numPoints =
        if odd numPoints'
          then numPoints'
          else numPoints' + 1
  coefficients <-
    computeFourierCoefficients
      gpuIDs
      numR2Freqs
      numPoints
      period
      delta
      batchSizeR2Freqs
      batchSizePolarFreqs
      numAngularFreqs
      numRadialFreqs
  writeArrayToStorableFile coefFilePath coefficients
  coefficientsArr' <-
    readArrayFromStorableFile
      coefFilePath
      (Z :. numRadialFreqs :. numAngularFreqs :. numR2Freqs :. numR2Freqs) :: IO (R.Array F DIM4 (Complex Double))
  coefficientsArr <- applyGaussian 1 std coefficientsArr'
  let coefficientsVecs =
        L.map
          (\(!rf, !af) ->
             VU.convert . toUnboxed . computeS . R.slice coefficientsArr $
             (Z :. rf :. af :. All :. All))
          [ (rf, af)
          | rf <- [0 .. numRadialFreqs - 1]
          , af <- [0 .. numAngularFreqs - 1]
          ]
  createDirectoryIfMissing True (folderPath </> "Coefficients")
  MP.mapM_
    (\(rf, af) ->
       plotImageRepaComplex
         (folderPath </> "Coefficients" </>
          printf
            "Coef(%d,%d).png"
            (rf - div numRadialFreqs 2)
            (af - div numAngularFreqs 2)) .
       ImageRepa 8 .
       computeS . extend (Z :. (1 :: Int) :. All :. All) . R.slice coefficients $
       (Z :. rf :. af :. All :. All)) .
    L.take 10 $
    [ (rf, af)
    | rf <- [0 .. numRadialFreqs - 1]
    , af <- [0 .. numAngularFreqs - 1]
    ]
  recons <-
    computeFourierSeries
      gpuIDs
      numR2Freqs
      (round len)
      period
      inverseDelta
      batchSizePolarFreqs
      batchSizeR2
      numAngularFreqs
      numRadialFreqs
      coefficientsVecs
  createDirectoryIfMissing True (folderPath </> "Recons")
  MP.mapM_
    (\(!rf, !af) ->
       plotImageRepaComplex
         (folderPath </> "Recons" </>
          printf
            "Recon(%d,%d).png"
            (rf - div numRadialFreqs 2)
            (af - div numAngularFreqs 2)) .
       ImageRepa 8 .
       computeS . extend (Z :. (1 :: Int) :. All :. All) . R.slice recons $
       (Z :. rf :. af :. All :. All)) .
    L.take 10 $
    [ (rf, af)
    | rf <- [0 .. numRadialFreqs - 1]
    , af <- [0 .. numAngularFreqs - 1]
    ]


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveAnyClass      #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}
{-|
Module      : Grenade.Core.WeightInitialization
Description : Defines the Weight Initialization methods of Grenade.
Copyright   : (c) Manuel Schneckenreither, 2018
License     : BSD2
Stability   : experimental

This module defines the weight initialization methods.

-}

module Grenade.Core.WeightInitialization
    ( getRandomVector
    , getRandomVectorV
    , getRandomMatrix
    , getRandomList
    , WeightInitMethod (..)
    ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Primitive         (PrimBase, PrimState)
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits
import qualified Data.Vector.Storable            as V
import           GHC.Generics                    (Generic)
import           GHC.TypeLits                    hiding (natVal)
import           Numeric.LinearAlgebra.Static
import           System.IO.Unsafe                (unsafePerformIO)
import           System.Random.MWC
import           System.Random.MWC.Distributions

import           Grenade.Types
import           Grenade.Utils.Vector

import           Debug.Trace

-- ^ Weight initialization method.
data WeightInitMethod
  = UniformInit -- ^ W_l,i ~ U(-1/sqrt(n_l),1/sqrt(n_l))                   where n_l is the number of nodes in layer l
  | Xavier      -- ^ W_l,i ~ U(-sqrt (6/n_l+n_{l+1}),sqrt (6/n_l+n_{l+1})) where n_l is the number of nodes in layer l
  | HeEtAl      -- ^ W_l,i ~ N(0,sqrt(2/n_l))                              where n_l is the number of nodes in layer l
  deriving (Show, Eq, Ord, Enum, Bounded, NFData, Serialize, Generic)

type LayerInput = Integer
type LayerOutput = Integer

getRandomList ::
     (PrimBase m)
  => LayerInput
  -> LayerOutput
  -> Int -- ^ Length of list
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m [RealNum]
getRandomList i o n method gen = do
  let unifRands = replicateM n (uniformR (-1, 1) gen)
      gaussRands = replicateM n (realToFrac <$> standard gen)
  case method of
    UniformInit -> map (1 / sqrt (fromIntegral i) *) <$> unifRands
    Xavier      -> map ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl      -> map (sqrt (2 / fromIntegral i) *) <$> gaussRands

getRandomVectorV ::
     (PrimBase m)
  => LayerInput
  -> LayerOutput
  -> Int -- ^ Length of list
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m (V.Vector RealNum)
getRandomVectorV i o n method gen = do
  let mkVec :: [RealNum] -> V.Vector RealNum
      mkVec xs = V.imap (\idx (_ :: RealNum) -> xs' V.! idx) (createVectorUnsafe n)
        where
          xs' = V.fromList xs -- convert first to vector, otherwise this function is quadratic!
  let unifRands = replicateM n (uniformR (-1, 1) gen)
      gaussRands = replicateM n (realToFrac <$> standard gen)
  case method of
    UniformInit -> mkVec . map (1 / sqrt (fromIntegral i) *) <$> unifRands
    Xavier -> mkVec . map ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl -> mkVec . map (sqrt (2 / fromIntegral i) *) <$> gaussRands

-- test :: IO [Double]
-- test = withSystemRandom $ asGenIO $ \gen -> replicateM 10000 (uniformR (-1, 1) gen)

-- test :: Int -> IO (V.Vector Double)
-- test n = withSystemRandom $ asGenIO $ getRandomVectorV 10 100 n UniformInit

-- | Get a random vector initialized according to the specified method.
getRandomVector ::
     forall m n. (PrimBase m, KnownNat n)
  => Integer
  -> Integer
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m (R n)
getRandomVector i o method gen = do
  let unifRands = vector <$> replicateM n (uniformR (-1, 1) gen)
      gaussRands = vector <$> replicateM n (realToFrac <$> standard gen)
  case method of
    UniformInit -> ((1 / sqrt (fromIntegral i)) *) <$> unifRands
    Xavier      -> ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl      -> (sqrt (2 / fromIntegral i) *) <$> gaussRands
  where
    n = fromIntegral $ natVal (Proxy :: Proxy n)

-- | Get a matrix with weights initialized according to the specified method.
getRandomMatrix ::
     forall m r n. (PrimBase m, KnownNat r, KnownNat n, KnownNat (n * r))
  => Integer
  -> Integer
  -> WeightInitMethod
  -> Gen (PrimState m)
  -> m (L r n)
getRandomMatrix i o method gen = do
  let unifRands = matrix <$> replicateM nr (uniformR (-1, 1) gen)
      gaussRands = matrix <$> replicateM nr (realToFrac <$> standard gen)
  case method of
    UniformInit -> ((1 / sqrt (fromIntegral i)) *) <$> unifRands
    Xavier      -> ((sqrt 6 / sqrt (fromIntegral i + fromIntegral o)) *) <$> unifRands
    HeEtAl      -> (sqrt (2 / fromIntegral i) *) <$> gaussRands
  where
    nr = fromIntegral $ natVal (Proxy :: Proxy (n * r))


{-# OPTIONS_GHC -Wall                   #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

{-# LANGUAGE DataKinds #-}

module HaddockExample ( main ) where

import           Numeric.LinearAlgebra.Static
                 ( R, vector, Sym,
                   headTail, matrix, sym
                 )
import           Data.Random.Distribution.MultiNormal
import           Numeric.Kalman
import           Numeric.LinearAlgebra.Static

import Graphics.Rendering.Chart hiding ( Vector )
import Graphics.Rendering.Chart.Backend.Diagrams
import Diagrams.Backend.Cairo.CmdLine
import Diagrams.Prelude hiding ( render, Renderable, sample )
import Diagrams.Backend.CmdLine

import Data.Csv
import System.IO hiding ( hGetContents )
import Data.ByteString.Lazy ( hGetContents )
import qualified Data.Vector as V

deltaT, g :: Double
deltaT = 0.01
g  = 9.81

bigQ :: Sym 2
bigQ = sym $ matrix bigQl

qc1 :: Double
qc1 = 0.01

bigQl :: [Double]
bigQl = [ qc1 * deltaT^3 / 3, qc1 * deltaT^2 / 2,
           qc1 * deltaT^2 / 2,       qc1 * deltaT
         ]

bigR :: Sym 1
bigR  = sym $ matrix [0.1]

observe :: R 2 -> R 1
observe a = vector [sin x] where x = fst $ headTail a

linearizedObserve :: R 2 -> L 1 2
linearizedObserve a = matrix [cos x, 0.0] where x = fst $ headTail a

stateUpdate :: R 2 -> R 2
stateUpdate u =  vector [x1 + x2 * deltaT, x2 - g * (sin x1) * deltaT]
  where
    (x1, w) = headTail u
    (x2, _) = headTail w

linearizedStateUpdate :: R 2 -> Sq 2
linearizedStateUpdate u = matrix [1.0,                    deltaT,
                                  -g * (cos x1) * deltaT,    1.0]
  where
    (x1, _) = headTail u

singleEKF :: MultiNormal (R 2) -> R 1 -> MultiNormal (R 2)
singleEKF = runEKF (const observe) (const linearizedObserve) (const bigR)
             (const stateUpdate) (const linearizedStateUpdate) (const bigQ)
             undefined

singleUKF :: MultiNormal (R 2) -> R 1 -> MultiNormal (R 2)
singleUKF = runUKF (const observe) (const bigR) (const stateUpdate) (const bigQ)
             undefined

initialDist :: MultiNormal (R 2)
initialDist = MultiNormal (vector [1.6, 0.0])
                          (sym $ matrix [0.1, 0.0,
                                         0.0, 0.1])

multiEKF :: [ℝ] -> [MultiNormal (R 2)]
multiEKF obs = scanl singleEKF initialDist (map (vector . pure) obs)

multiUKF :: [ℝ] -> [MultiNormal (R 2)]
multiUKF obs = scanl singleUKF initialDist (map (vector . pure) obs)

denv :: IO (DEnv Double)
denv = defaultEnv vectorAlignmentFns 600 500

diagEstimated :: String ->
                 [(Double, Double)] ->
                 [(Double, Double)] ->
                 [(Double, Double)] ->
                 IO (Diagram Cairo)
diagEstimated t l xs es = do
  env <- denv
  return $ fst $ runBackend env (render (chartEstimated t l xs es) (600, 500))

chartEstimated :: String ->
              [(Double, Double)] ->
              [(Double, Double)] ->
              [(Double, Double)] ->
              Renderable ()
chartEstimated title acts obs ests = toRenderable layout
  where

    actuals = plot_lines_values .~ [acts]
            $ plot_lines_style  . line_color .~ opaque red
            $ plot_lines_title .~ "Actual Trajectory"
            $ plot_lines_style  . line_width .~ 1.0
            $ def

    measurements = plot_points_values .~ obs
                 $ plot_points_style  . point_color .~ opaque blue
                 $ plot_points_title .~ "Measurements"
                 $ def

    estimas = plot_lines_values .~ [ests]
            $ plot_lines_style  . line_color .~ opaque black
            $ plot_lines_title .~ "Inferred Trajectory"
            $ plot_lines_style  . line_width .~ 1.0
            $ def

    layout = layout_title .~ title
           $ layout_plots .~ [toPlot actuals, toPlot measurements, toPlot estimas]
           $ layout_y_axis . laxis_title .~ "Angle / Horizontal Displacement"
           $ layout_y_axis . laxis_override .~ axisGridHide
           $ layout_x_axis . laxis_title .~ "Time"
           $ layout_x_axis . laxis_override .~ axisGridHide
           $ def

displayHeader :: FilePath -> Diagram B -> IO ()
displayHeader fn =
  mainRender ( DiagramOpts (Just 900) (Just 700) fn
             , DiagramLoopOpts False Nothing 0
             )

main :: IO ()
main = do
  h <- openFile "matlabRNGs.csv" ReadMode
  cs <- hGetContents h
  let df = (decode NoHeader cs) :: Either String (V.Vector (Double, Double))
  case df of
    Left _ -> error "Whatever"
    Right generatedSamples -> do
      let xs = take 500 (multiEKF $ V.toList $ V.map fst generatedSamples)
      let mus = map (fst . headTail . mu) xs
      let obs = V.toList $ V.map fst generatedSamples
      let acts = V.toList $ V.map snd generatedSamples
      de1 <- diagEstimated "Fitted Pendulum"
             (zip [0,1..] acts)
             (zip [0,1..] obs)
             (zip [0,1..] mus)
      displayHeader "diagrams/PendulumFittedEkf.png" de1
      let ys = take 500 (multiUKF $ V.toList $ V.map fst generatedSamples)
      let nus = map (fst . headTail . mu) ys
      de2 <- diagEstimated "Fitted Pendulum"
             (zip [0,1..] acts)
             (zip [0,1..] obs)
             (zip [0,1..] nus)
      displayHeader "diagrams/PendulumFittedUkf.png" de2


{-|
Module: MachineLearning.Classification.OneVsAll
Description: One-vs-All Classification.
Copyright: (c) Alexander Ignatyev, 2016-2017
License: BSD-3
Stability: experimental
Portability: POSIX

One-vs-All Classification.
-}

module MachineLearning.Classification.OneVsAll
(
    Opt.MinimizeMethod(..)
  , module Log
  , module Model
  , predict
  , learn
  , MLC.calcAccuracy
  , Regularization(..)
)

where

import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Regularization (Regularization(..))
import qualified MachineLearning.Optimization as Opt
import qualified MachineLearning.LogisticModel as Log
import qualified MachineLearning.Model as Model
import qualified MachineLearning.Classification.Internal as MLC
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

-- | One-vs-All Classification prediction function.
-- Takes a matrix of features X and a list of vectors theta,
-- returns predicted class number assuming that class numbers start at 0.
predict :: Matrix -> [Vector] -> Vector
predict x thetas = predictions'
  where predict = Model.hypothesis Log.Logistic x
        predictions = LA.toRows . LA.fromColumns $ map predict thetas
        predictions' = LA.vector $ map (fromIntegral . LA.maxIndex) predictions

-- | Learns One-vs-All Classification
learn :: Opt.MinimizeMethod -- ^ (e.g. BFGS2 0.1 0.1)
         -> R                  -- ^ epsilon, desired precision of the solution;
         -> Int                -- ^ maximum number of iterations allowed;
         -> Regularization     -- ^ regularization parameter lambda;
         -> Int                -- ^ number of labels
         -> Matrix             -- ^ matrix X;
         -> Vector             -- ^ vector y
         -> [Vector]             -- ^ initial theta list;
         -> ([Vector], [Matrix])  -- ^ solution vector and optimization path.
learn mm eps numIters lambda nLabels x y initialThetaList =
  let ys = MLC.processOutputOneVsAll nLabels y
      minimize = Opt.minimize mm Log.Logistic eps numIters lambda x
  in unzip $ zipWith minimize ys initialThetaList


module HLearn.Optimization.Amoeba
    where

import Prelude
import Control.Monad
import Control.Monad.ST
import Data.List
import Debug.Trace
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Algorithms.Intro as Intro
import Numeric.LinearAlgebra hiding ((<>))
import qualified Numeric.LinearAlgebra as LA

findMinAmoeba f x0 = runST $ do

    -- initialize simplex
    vec <- VM.new (VG.length x0+1)
    VGM.write vec 0 (f x0,x0)
    forM [1..VGM.length vec-1] $ \i -> do
        e_i <- VGM.replicate (VG.length x0) 0
        VGM.write e_i (i-1) 1
        e_i' <- VG.freeze e_i
        let x_i = x0 `LA.add` e_i'
        VGM.write vec i (f x_i,x_i)

    -- iterate
    vec' <- itrM 1000 (stepAmoeba f) vec

    -- return
    (_,ret) <- VGM.read vec 0
    return ret

stepAmoeba f vec = stepAmoebaRaw 1 2 (-1/2) (1/2) f vec

stepAmoebaRaw ::
    ( Fractional b
    , Ord b
    , VGM.MVector vec (b,a)
--     , a ~ LA.Matrix b
    , a ~ LA.Vector b
    , Field b
    , vec ~ VM.MVector
    , b ~ Double
    ) => b
      -> b
      -> b
      -> b
      -> (a -> b)
      -> vec s (b,a)
      -> ST s (vec s (b,a))
stepAmoebaRaw alpha gamma ro sigma f vec = do

    Intro.sortBy (\a b -> compare (fst a) (fst b)) vec

    (f_1,x_1) <- VGM.read vec 0
    (f_2,x_2) <- VGM.read vec 1
    (f_n1,x_n1) <- VGM.read vec $ VGM.length vec -1

    x_0 <- liftM ( scale (1/fromIntegral (VGM.length vec-1))
                 . foldl1' (LA.add)
                 . init
                 . map snd
                 . V.toList
                 ) $ VG.unsafeFreeze vec

    let x_r = x_0 `LA.add` (scale alpha $ x_0 `LA.sub` x_n1)
        f_r = f x_r

        x_e = x_0 `LA.add` (scale gamma $ x_0 `LA.sub` x_n1)
        f_e = f x_e

        x_c = x_0 `LA.add` (scale ro $ x_0 `LA.sub` x_n1)
        f_c = f x_c

    -- check reflection
    if f_1 <= f_r && f_r < f_1
        then VGM.write vec (VGM.length vec-1) (f_r,x_r)

        -- check expansion
        else if f_r < f_1
            then if f_e < f_r
                then VGM.write vec (VGM.length vec-1) (f_e,x_e)
                else VGM.write vec (VGM.length vec-1) (f_r,x_r)

            -- check contraction
            else if f_c < f_n1
                then VGM.write vec (VGM.length vec-1) (f_c,x_c)

                -- reduction
                else forM_ [1..VGM.length vec-1] $ \i -> do
                    (f_i,x_i) <- VGM.read vec i
                    let x_i' = x_1 `LA.add` (scale sigma $ x_i `LA.sub` x_1)
                        f_i' = f x_i'
                    VGM.write vec i (f_i',x_i')

    return vec
--     refMinVal <- newSTRef (-infinity)
--     refMinIndex <- newSTRef 0
--     forM [0..VGM.length vec-1] $ \i -> do
--         ival <- VGM.read vec i
--         minVal <- readSTRef refMinVal
--         if minVal < fst ival
--             then return ()
--             else do
--                 writeSTRef refMinVal $ fst ival
--                 writeSTRef refMinIndex i
--     undefined

itrM :: Monad m => Int -> (a -> m a) -> a -> m a
itrM 0 f a = return a
itrM i f a = do
    a' <- f a
--     if a' == a
--         then trace ("no movement\n   a="++show a++"\n   a'="++show a') $ return ()
--         else trace ("itrM i="++show i++"; f a ="++show a'++"; a="++show a) $ return ()
    itrM (i-1) f a'


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
import System.IO
import Control.Monad
import Data.Complex
import Data.List
import Data.ByteString (ByteString)

data Obj2D = Obj2D [String]
data Obj3D = Obj3D [String]

class Render2D a where
  render2d :: a -> [Complex Double]

class Render3D a where
  render3d :: a -> Obj3D

data Gear = Gear {
  gAlpha :: Double
, gTooth :: Double
, gModule :: Double
, gWidth :: Double
} deriving Show

data Bar = Bar {
  bWidth :: Double
} deriving Show

involute :: Double -> Double -> Complex Double
involute r theta = (r*(cos theta + theta * sin theta)) :+ (r*(sin theta - theta * cos theta))

yrev :: Complex Double -> Complex Double
yrev (a:+b) = (a:+(-b))

alpha :: Gear -> Double
alpha Gear{..} = 2 * pi * gAlpha / 360

radius :: Gear -> Double
radius g@Gear{..} = (cos (alpha g)) * radius' g

radius' :: Gear -> Double
radius' Gear{..} = gModule * gTooth / 2

instance Render2D Gear where
  render2d g@Gear{..} =
    let dat :: [Complex Double]
        dat = ((radius g) - 0.5) :+ 0 : map (\i-> involute (radius g) (i*2.0*pi/100.0))  [0..10]
        dat''' :: [Complex Double]
        dat''' = map (\i-> mkPolar (radius g) (i*2.0*pi/100.0))  [0..100]
        gear''  =  concat $
                   flip map [0..(gTooth-1)] $
                   \t -> map (\i -> i * (cis (t*2*pi/gTooth) )) $
                   dat ++
                   map (\t -> t * cis (0.9*pi/gTooth))
                   (reverse (map yrev dat))
    in gear''

instance Render3D Gear where
  render3d g@Gear{..} =
    let gear2d = render2d g
        gear2d' = intercalate "," $ map (\(r:+i) -> "[" ++ show r ++ "," ++ show i ++ "]") gear2d
    in Obj3D ["linear_extrude(height = " ++ show gWidth ++ ") { polygon( points= [",
              gear2d',
              "] );}"]

writePlot :: String -> [Complex Double] -> IO ()
writePlot file dat = do
  withFile file WriteMode $ \h -> do
    forM_ dat $ \(r:+i) -> do
      hPutStrLn h $ show r ++ " " ++ show i

write3DPlot :: String -> Obj3D -> IO ()
write3DPlot file (Obj3D dat) = do
  withFile file WriteMode $ \h -> do
    forM_ dat $ \str -> do
      hPutStr h $ str

main = do
  writePlot "involute.dat" $ render2d $ Gear 20 8 1 6
  write3DPlot "involute.scad" $ render3d $ Gear 20 8 1 6


module Math.FROG.Tools
( mkTrace
, mkSigGate
, shift
, mkInitialGuess
, flatPolar2Complex
, center
) where

import qualified System.Random as R
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.GSL.Fourier as F
import qualified Foreign

import Math.FROG.Types

mkTrace :: ComplexSignal -> ComplexSignal -> Trace
mkTrace field gate = LA.fromColumns $ map atEachDelayApply [upper,upper-1..lower]
    where
    atEachDelayApply = (^2) . (LA.mapVector LA.magnitude)
                            . (flip shift upper)
                            .  F.fft
                            . (*field)
                            . (shift gate)
    sz = LA.dim field
    upper = (sz `quot` 2)
    lower = (-1) * (sz `quot` 2) + 1

mkSigGate :: Nonlinearity -> ComplexSignal -> ComplexSignal -> (ComplexSignal, ComplexSignal)
mkSigGate PG f g = (id f, LA.mapVector ((^2) . abs) g)
mkSigGate SHG f g = (id f, id g)
mkSigGate SD f g = (LA.mapVector (^2) f, LA.mapVector LA.conjugate g)
mkSigGate THG f g = (LA.mapVector (^2) f, id g)

shift :: ComplexSignal -> Int -> ComplexSignal
shift gate delay = LA.buildVector sz (builder delay)
    where
    sz = LA.dim gate
    builder d k = if any (<0) [(sz-1-d-k), (k+d)] then (getElem (k+d-(signum d)*sz)) else (getElem (k+d))
        where
        getElem = (LA.@>) gate

mkInitialGuess :: Int -> IO (LA.Vector Double)
mkInitialGuess sz = do
    stdgen <- R.getStdGen
    let randMags = R.randomRs (0, 1.0) stdgen
    let randAngs = R.randomRs (-pi, pi) stdgen
    return $ LA.buildVector (2*sz) (\k -> if k < sz then (randMags !! k) else (randAngs !! k))

flatPolar2Complex :: LA.Vector Double -> LA.Vector (LA.Complex Double)
flatPolar2Complex v = LA.buildVector sz (\k -> LA.mkPolar (mags LA.@> k) (angs LA.@> k))
    where
    sz = (LA.dim v) `quot` 2
    mags = LA.subVector 0 sz v
    angs = LA.subVector sz sz v

center :: ComplexSignal -> ComplexSignal
center f = shift f $ getDelay . getMaxInd $ f
    where
    sz = LA.dim f
    getMaxInd = LA.maxIndex . (LA.mapVector LA.magnitude)
    getDelay = (*(-1)) . ((-) (sz `quot` 2))


-----------------------------------------------------------------------------
-- |
-- Module      :  Finance.Hqfl.Pricer.BlackScholes
-- Copyright   :  (C) 2016 Mika'il Khan
-- License     :  (see the file LICENSE)
-- Maintainer  :  Mika'il Khan <co.kleisli@gmail.com>
-- Stability   :  stable
-- Portability :  portable
--
----------------------------------------------------------------------------
{-# LANGUAGE FlexibleInstances #-}

module Finance.Hqfl.Pricer.BlackScholes where

import Finance.Hqfl.Instrument
import Statistics.Distribution.Normal
import Data.Random

class BS a where
  price :: a -> Double -> Double -> Double

instance BS (Option Equity) where
  price (Option (Equity s q) m European k t) r v =
    case m of
       Call -> s * exp (-q * t) * cdf normal d1 - k * exp (-r * t) * cdf normal d2
       Put  -> k * exp (-r * t) * cdf normal (-d2) - s * exp (-q * t) * cdf normal (-d1)
    where d1 = (log (s / k) + (r - q + (v * v) / 2) * t) / (v * sqrt t)
          d2 = d1 - v * sqrt t
          normal = Normal (0 :: Double) 1

instance BS (Option StockIndex) where
  price (Option (StockIndex s q) m European k t) r v =
    case m of
       Call -> s * exp (-q * t) * cdf normal d1 - k * exp (-r * t) * cdf normal d2
       Put  -> k * exp (-r * t) * cdf normal (-d2) - s * exp (-q * t) * cdf normal (-d1)
    where d1 = (log (s / k) + (r - q + (v * v) / 2) * t) / (v * sqrt t)
          d2 = d1 - v * sqrt t
          normal = Normal (0 :: Double) 1

-- Garman & Kohlhagen (1983)
--instance BS (Option Currency) where
--  price (Option (Currency s rf) m European k t) r v =
--    case m of
--       Call -> s * exp (-rf * t) * cdf normal d1 - k * exp (-r * t) * cdf normal d2
--       Put  -> k * exp (-r * t) * cdf normal (-d2) - s * exp (-rf * t) * cdf normal (-d1)
--    where d1 = (log (s / k) + (r - rf + (v * v) / 2) * t) / (v * sqrt t)
--          d2 = d1 - v * sqrt t
--          normal = Normal (0 :: Double) 1

-- cdi :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- cdi s k r v t q h = s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal y)
--                     - k * exp(-r * t) * (h / s)**(2 * lambda - 2) * (cdf normal (y - (v * sqrt t)))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- cdo :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- cdo s k r v t q h = s * (cdf normal x1) * exp (-q * t) - k * exp(-r * t) * (cdf normal (x1 - v * sqrt t)) -
--                     s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal y1) + k * exp(-r * t) * (h/ s)**(2 * lambda - 2) * (cdf normal (y1 - v * sqrt t))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- cui :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- cui s k r v t q h = s * (cdf normal x1) * exp(-q * t) - k * exp(-r * t) * (cdf normal (x1 - v * sqrt t)) -
--       s * exp(-q *t) * (h / s)**(2 * lambda) * ((cdf normal (-y)) - (cdf normal (-y1))) +
--       k * exp(-r * t) * (h / s)**(2 * lambda - 2) * ((cdf normal (-y + v * sqrt t)) - (cdf normal (-y1 + v * sqrt t)))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- pui :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- pui s k r v t q h = -s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal (-y)) + k * exp(-r * t) * (h / s)**(2 * lambda - 2) * (cdf normal ((-y) + (v * sqrt t)))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- pdi :: Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double
-- pdi s k r v t q h = -s * (cdf normal (-x1)) * exp(-q * t)
--                     + k * exp(-r * t) * (cdf normal (-x1 + v * sqrt t))
--                     + s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal y - cdf normal y1)
--                     - k * exp(-r * t) * (h / s)**(2 * lambda -2) * (cdf normal (y - v * sqrt t) - cdf normal (y1 - v * sqrt t))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- puo s k r v t q h = -s * (cdf normal (-x1)) * exp(-q * t) + k * exp(-r * t) * (cdf normal (-x1 + v * sqrt t)) +
--                     s * exp(-q * t) * (h / s)**(2 * lambda) * (cdf normal (-y1)) - k * exp(-r * t) * (h / s)**(2 * lambda -2) * (cdf normal ((-y1) + v * sqrt t))
--   where lambda = ((r - q + (v * v) / 2)) / (v * v)
--         y1 = (log (h / s) / v * sqrt t) + (lambda * v * sqrt t)
--         x1 = (log ( s / h ) / (v * sqrt t)) + (lambda * v * sqrt t)
--         y = (log ( (h*h) / (s * k)) / v * sqrt t) + (lambda * v * sqrt t)
--         normal = Normal (0 :: Double) 1

-- instance BS (Option StockIndex) where
--   blackscholesprice (Option (StockIndex s) Call European k t) r v = call s k r v t q
--   blackscholesprice (Option (StockIndex s) Put European k t) r v = put s k r v t q

-- instance BS (Option Future) where
--   blackscholesprice (Option (Future f) Call European k t) r v = call f k r v t q
--   blackscholesprice (Option (Future f) Put European k t) r v = put f k r v t q

-- TODO : factor in rebates
-- instance BS (BarrierOption Equity) where
--   blackscholesprice (BarrierOption (Equity s q) x European y h k t) r v =
--     case (x, y) of
--       (DownAndIn, Call)
-- --        | s <= h -> Left "Contract is converted to a vanilla call"
--         | h <= k ->  cdi'
--         | h >= k ->   call' - cdo'
--       (DownAndOut, Call)
-- --        | s <= h -> Left "Contract is void"
--         | h <= k ->   call' - cdi'
--         | h >= k ->  cdo'
--       (UpAndIn, Call)
-- --        | h <= s -> Left "Contract is converted to a vanilla call"
--         | h <= k ->  call'
--         | h >= k ->   cui'
--       (UpAndOut, Call)
-- --        | s <= h -> Left "Contract is void"
--         | h <= k ->   call' - cdi'
--         | h >= k ->  cdo'
--       (DownAndIn, Put)
-- --        | s <= h -> Left "Contract is converted to a vanilla call"
--         | h <  k ->  pdi'
--         | h >= k ->  put'
--       (DownAndOut, Put)
-- --        | s <= h -> Left "Contract is void"
--         | h <  k ->   put' - pdi'
--         | h >=  k ->  0
--       (UpAndIn, Put)
-- --        | h <= s -> Left "Contract is converted to a vanilla call"
--         | h <= k ->   put' - puo'
--         | h >= k ->  pui'
--       (UpAndOut, Put)
-- --        | h <= s -> Left "Contract is void"
--         | h <= k ->   puo'
--         | h >= k ->   put' - pui'
--     where
--       call' = call s k r v t q
--       cdi'  = cdi  s k r v t q h
--       cdo'  = cdo  s k r v t q h
--       cui'  = cui  s k r v t q h
--       pdi'  = pdi  s k r v t q h
--       put'  = put  s k r v t q
--       pui'  = pui  s k r v t q h
--       puo'  = puo  s k r v t q h


module Tests.Function ( tests ) where

import Statistics.Function
import Test.Tasty
import Test.Tasty.QuickCheck
import Test.QuickCheck
import Tests.Helpers
import qualified Data.Vector.Unboxed as U

tests :: TestTree
tests = testGroup "S.Function"
  [ testProperty  "Sort is sort"                p_sort
  , testAssertion "nextHighestPowerOfTwo is OK" p_nextHighestPowerOfTwo
  ]

p_sort :: [Double] -> Property
p_sort xs =
  not (null xs) ==> U.all (uncurry (<=)) (U.zip v $ U.tail v)
    where
      v = sort $ U.fromList xs

p_nextHighestPowerOfTwo :: Bool
p_nextHighestPowerOfTwo
  = all (\(good, is) -> all ((==good) . nextHighestPowerOfTwo) is) lists
  where
    pows  = [1 .. 17 :: Int]
    lists = [ (2^m, [2^n+1 .. 2^m]) | (n,m) <- pows `zip` tail pows ]


module Supervised.LinearRegressionSpec where

import Prelude
import Test.Hspec
import Supervised.LinearRegression
import Datasets.Iris
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra as LA

v2la :: V.Vector IrisClass -> LA.Vector Double
v2la = LA.vector . V.toList . fmap (fromIntegral . fromEnum)

main :: IO ()
main = do
  ((trainX, trainY'), (testX, testY')) <- splitCV 0.6 <$> loadIris
  let
    !trainY = v2la trainY' :: LA.Vector Double
    !testY  = v2la testY' :: LA.Vector Double

  let
    trainedModel :: Model
    trainedModel = fit mkModel trainX trainY

  let
    predicted :: Labels
    predicted = trainedModel `predict` testX

  let
    score :: Double
    score = rSquare testY predicted

  print $ score
--  hspec spec

spec :: Spec
spec =
  describe "Linear Regression" $ undefined