{-# LANGUAGE CPP              #-}
{-# OPTIONS_GHC -fno-warn-unused-imports -fno-warn-incomplete-patterns -fno-warn-missing-signatures #-}
{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes       #-}
{-# LANGUAGE TypeFamilies     #-}
{-# LANGUAGE TypeOperators    #-}
{-# LANGUAGE ViewPatterns     #-}

-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Tests
Copyright   :  (c) Alberto Ruiz 2007-14
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

Some tests.

-}

module Test.Numeric.LinearAlgebra.Tests(
--  module Numeric.LinearAlgebra.Tests.Instances,
--  module Numeric.LinearAlgebra.Tests.Properties,
   qCheck,
   utest,
   runTests,
   runBenchmarks
 , binaryTests
-- , findNaN
--, runBigTests
) where

import           Data.List                                   (foldl1')
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Devel
import           Numeric.LinearAlgebra.Static                (L)
import           System.Info
import           Test.HUnit                                  hiding (State, Testable, test,
                                                              (~:))
import           Test.Numeric.LinearAlgebra.Tests.Instances
import           Test.Numeric.LinearAlgebra.Tests.Properties
#if MIN_VERSION_base(4,11,0)
import           Prelude                                     hiding ((<>), (^))
#else
import           Prelude                                     hiding ((^))
#endif
import           Control.Applicative
import           Control.Arrow                               ((***))
import           Control.DeepSeq                             (NFData (..))
import           Control.Monad                               (when)
import           Control.Monad                               (ap)
import           Debug.Trace
import           Numeric.LinearAlgebra.Devel                 (unsafeFromForeignPtr,
                                                              unsafeToForeignPtr)
import qualified Prelude
import           System.CPUTime
import           System.Exit
import           Text.Printf

import           Test.QuickCheck                             (Arbitrary, Property,
                                                              Testable, arbitrary, choose,
                                                              classify, coarbitrary,
                                                              maxSize,
                                                              quickCheckWithResult, shrink,
                                                              sized, stdArgs, vector)
import qualified Test.QuickCheck                             as T

import           Test.QuickCheck.Test                        (isSuccess)

--eps = peps :: Float
--i = 0:+1 :: Complex Float

qCheck n x = do
    r <- quickCheckWithResult stdArgs {maxSize = n} x
    when (not $ isSuccess r) (exitFailure)

a ^ b = a Prelude.^ (b :: Int)

utest str b = TestCase $ assertBool str b

feye n = flipud (ident n) :: Matrix Float

-----------------------------------------------------------

detTest1 = det m == 26
        && det mc == 38 :+ (-3)
        && det (feye 2) == -1
    where
        m = (3><3)
            [ 1, 2, 3
            , 4, 5, 7
            , 2, 8, 4 :: Float
            ]
        mc = (3><3)
            [ 1, 2, 3
            , 4, 5, 7
            , 2, 8, iC
            ]

detTest2 = inv1 |~| inv2 && [det1] ~~ [det2]
  where
    m = complex (feye 6)
    inv1 = inv m
    det1 = det m
    (inv2,(lda,sa)) = invlndet m
    det2 = sa * exp lda

---------------------------------------------------------------------

nd1 = (3><3) [ 1/2, 1/4, 1/4
             , 0/1, 1/2, 1/4
             , 1/2, 1/4, 1/2 :: Float]

nd2 = (2><2) [1, 0, 1, 1:: Complex Float]

expmTest1 = expm nd1 :~14~: (3><3)
 [ 1.762110887278176
 , 0.478085470590435
 , 0.478085470590435
 , 0.104719410945666
 , 1.709751181805343
 , 0.425725765117601
 , 0.851451530235203
 , 0.530445176063267
 , 1.814470592751009 ]

expmTest2 = expm nd2 :~15~: (2><2)
 [ 2.718281828459045
 , 0.000000000000000
 , 2.718281828459045
 , 2.718281828459045 ]

-----------------------------------------------------

mbCholTest = utest "mbCholTest" (ok1 && ok2) where
    m1 = (2><2) [2,5,5,8 :: Float]
    m2 = (2><2) [3,5,5,9 :: Complex Float]
    ok1 = mbChol (trustSym m1) == Nothing
    ok2 = mbChol (trustSym m2) == Just (chol $ trustSym m2)

-----------------------------------------------------

triTest = utest "triTest" ok1 where

  a :: Matrix R
  a = (4><4)
    [
       4.30,  0.00,  0.00, 0.00,
      -3.96, -4.87,  0.00, 0.00,
       0.40,  0.31, -8.02, 0.00,
      -0.27,  0.07, -5.95, 0.12
    ]

  w :: Matrix R
  w = (4><2)
    [
      -12.90, -21.50,
       16.75,  14.93,
      -17.55,   6.33,
      -11.04,   8.09
    ]

  v :: Matrix R
  v = triSolve Lower a w

  e :: Matrix R
  e = (4><2)
    [
      -3.0000, -5.0000,
      -1.0000,  1.0000,
       2.0000, -1.0000,
       1.0000,  6.0000
    ]

  ok1 = (norm_Inf . flatten $ e - v) <= 1e-13

-----------------------------------------------------

triDiagTest = utest "triDiagTest" (ok1 && ok2) where

  dL, d, dU :: Vector Float
  dL =  fromList [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0]
  d  =  fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
  dU =  fromList [4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0]

  b :: Matrix R
  b = (9><3)
    [
      1.0,   1.0,   1.0,
      1.0,  -1.0,   2.0,
      1.0,   1.0,   3.0,
      1.0,  -1.0,   4.0,
      1.0,   1.0,   5.0,
      1.0,  -1.0,   6.0,
      1.0,   1.0,   7.0,
      1.0,  -1.0,   8.0,
      1.0,   1.0,   9.0
    ]

  y :: Matrix R
  y = (9><9)
    [
      1.0, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
      3.0, 1.0, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
      0.0, 3.0, 1.0, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0,
      0.0, 0.0, 3.0, 1.0, 4.0, 0.0, 0.0, 0.0, 0.0,
      0.0, 0.0, 0.0, 3.0, 1.0, 4.0, 0.0, 0.0, 0.0,
      0.0, 0.0, 0.0, 0.0, 3.0, 1.0, 4.0, 0.0, 0.0,
      0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 1.0, 4.0, 0.0,
      0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 1.0, 4.0,
      0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 1.0
    ]

  x :: Matrix R
  x = triDiagSolve dL d dU b

  z :: Matrix C
  z = (4><4)
    [
      1.0 :+ 1.0, 4.0 :+ 4.0, 0.0 :+ 0.0, 0.0 :+ 0.0,
      3.0 :+ 3.0, 1.0 :+ 1.0, 4.0 :+ 4.0, 0.0 :+ 0.0,
      0.0 :+ 0.0, 3.0 :+ 3.0, 1.0 :+ 1.0, 4.0 :+ 4.0,
      0.0 :+ 0.0, 0.0 :+ 0.0, 3.0 :+ 3.0, 1.0 :+ 1.0
    ]

  zDL, zD, zDu :: Vector C
  zDL = fromList [3.0 :+ 3.0, 3.0 :+ 3.0, 3.0 :+ 3.0]
  zD  = fromList [1.0 :+ 1.0, 1.0 :+ 1.0, 1.0 :+ 1.0, 1.0 :+ 1.0]
  zDu = fromList [4.0 :+ 4.0, 4.0 :+ 4.0, 4.0 :+ 4.0]

  zB :: Matrix C
  zB = (4><3)
    [
      1.0 :+ 1.0,   1.0  :+   1.0,  1.0 :+ (-1.0),
      1.0 :+ 1.0, (-1.0) :+ (-1.0), 1.0 :+ (-1.0),
      1.0 :+ 1.0,   1.0  :+   1.0,  1.0 :+ (-1.0),
      1.0 :+ 1.0, (-1.0) :+ (-1.0), 1.0 :+ (-1.0)
    ]

  u :: Matrix C
  u = triDiagSolve zDL zD zDu zB

  ok1 = (maximum $ map abs $ concat $ toLists $ b - (y <> x)) <= 1e-15
  ok2 = (maximum $ map magnitude $ concat $ toLists $ zB - (z <> u)) <= 1e-15

---------------------------------------------------------------------

triDiagRegression = utest "triDiagRegression" ok where
  minusOnes, twos :: Vector R
  minusOnes = fromList [-1, -1]
  twos      = fromList [2, 2, 2]
  k :: Matrix R
  k = (3><3)
    [  2, -1,  0
    , -1,  2, -1
    ,  0, -1,  2
    ]

  b :: Matrix R
  b = (3><1) [10, 10, 10]

  tridiag = triDiagSolve minusOnes twos minusOnes b
  simple = linearSolve k b

  ok = case simple of
    Just m  -> tridiag |~| m
    Nothing -> False

---------------------------------------------------------------------

randomTestGaussian = (unSym c) :~3~: unSym (snd (meanCov dat))
  where
    a = (3><3) [1,2,3,
                2,4,0,
               -2,2,1]
    m = 3 |> [1,2,3]
    c = mTm a
    dat = gaussianSample 7 (10^6) m c

randomTestUniform = c :~2~: unSym (snd (meanCov dat))
  where
    c = diag $ 3 |> map ((/12).(^2)) [1,2,3]
    dat = uniformSample 7 (10^6) [(0,1),(1,3),(3,6)]

---------------------------------------------------------------------

rot :: Float -> Matrix Float
rot a = (3><3) [ c,0,s
               , 0,1,0
               ,-s,0,c ]
    where c = cos a
          s = sin a

rotTest = fun (10^5) :~11~: rot 5E4
    where fun n = foldl1' (<>) (map rot angles)
              where angles = toList $ linspace n (0,1)

---------------------------------------------------------------------
-- vector <= 0.6.0.2 bug discovered by Patrick Perry
-- http://trac.haskell.org/vector/ticket/31

offsetTest = y == y' where
    x = fromList [0..3 :: Float]
    y = subVector 1 3 x
    (f,o,n) = unsafeToForeignPtr y
    y' = unsafeFromForeignPtr f o n

---------------------------------------------------------------------

normsVTest = TestList [
    utest "normv2CD" $ norm2PropC v
--  , utest "normv2CF" $ norm2PropC (single v)
#ifndef NONORMVTEST
  , utest "normv2D"  $ norm2PropR x
--  , utest "normv2F"  $ norm2PropR (single x)
#endif
  , utest "normv1CD" $ norm_1 v          == 8
--  , utest "normv1CF" $ norm_1 (single v) == 8
  , utest "normv1D"  $ norm_1 x          == 6
--  , utest "normv1F"  $ norm_1 (single x) == 6

  , utest "normvInfCD" $ norm_Inf v          == 5
--  , utest "normvInfCF" $ norm_Inf (single v) == 5
  , utest "normvInfD"  $ norm_Inf x          == 3
--  , utest "normvInfF"  $ norm_Inf (single x) == 3

 ] where v = fromList [1,-2,3:+4] :: Vector (Complex Float)
         x = fromList [1,2,-3] :: Vector Float
#ifndef NONORMVTEST
         norm2PropR a = norm_2 a =~= sqrt (udot a a)
#endif
         norm2PropC a = norm_2 a =~= realPart (sqrt (a `dot` a))
         a =~= b = fromList [a] |~| fromList [b]

normsMTest = TestList [
    utest "norm2mCD" $ norm_2 v          =~= 8.86164970498005
--  , utest "norm2mCF" $ norm_2 (single v) =~= 8.86164970498005
  , utest "norm2mD"  $ norm_2 x          =~= 5.96667765076216
--  , utest "norm2mF"  $ norm_2 (single x) =~= 5.96667765076216

  , utest "norm1mCD" $ norm_1 v          == 9
--  , utest "norm1mCF" $ norm_1 (single v) == 9
  , utest "norm1mD"  $ norm_1 x          == 7
--  , utest "norm1mF"  $ norm_1 (single x) == 7

  , utest "normmInfCD" $ norm_Inf v          == 12
--  , utest "normmInfCF" $ norm_Inf (single v) == 12
  , utest "normmInfD"  $ norm_Inf x          == 8
--  , utest "normmInfF"  $ norm_Inf (single x) == 8

  , utest "normmFroCD" $ norm_Frob v          =~= 8.88819441731559
--  , utest "normmFroCF" $ norm_Frob (single v) =~~= 8.88819441731559
  , utest "normmFroD"  $ norm_Frob x          =~= 6.24499799839840
--  , utest "normmFroF"  $ norm_Frob (single x) =~~= 6.24499799839840

 ] where v = (2><2) [1,-2*iC,3:+4,7] :: Matrix (Complex Float)
         x = (2><2) [1,2,-3,5] :: Matrix Float
         a =~= b = fromList [a] :~10~: fromList [b]
--       a =~~= b = fromList [a] :~5~: fromList [b]

---------------------------------------------------------------------

sumprodTest = TestList [
    utest "sumCD" $ sumElements z            == 6
  , utest "sumCF" $ sumElements (single z)   == 6
  , utest "sumD"  $ sumElements v            == 6
  , utest "sumF"  $ sumElements (single v)   == 6

  , utest "prodCD" $ prodProp z
  , utest "prodCF" $ prodProp (single z)
  , utest "prodD"  $ prodProp v
  , utest "prodF"  $ prodProp (single v)
 ] where v = fromList [1,2,3] :: Vector Float
         z = fromList [1,2-iC,3+iC]
         prodProp x = prodElements x == product (toList x)

---------------------------------------------------------------------

chainTest = utest "chain" $ foldl1' (<>) ms |~| optimiseMult ms where
    ms = [ diag (fromList [1,2,3 :: Float])
         , konst 3 (3,5)
         , (5><10) [1 .. ]
         , konst 5 (10,2)
         ]

---------------------------------------------------------------------

conjuTest m = cmap conjugate (flatten (conj (tr m))) == flatten (tr m)

---------------------------------------------------------------------

newtype State s a = State { runState :: s -> (a,s) }

instance Functor (State s)
  where
    fmap f x = pure f <*> x

instance Applicative (State s)
  where
    pure = return
    (<*>) = ap

instance Monad (State s) where
    return a = State $ \s -> (a,s)
    m >>= f = State $ \s -> let (a,s') = runState m s
                            in runState (f a) s'

state_get :: State s s
state_get = State $ \s -> (s,s)

state_put :: s -> State s ()
state_put s = State $ \_ -> ((),s)

evalState :: State s a -> s -> a
evalState m s = let (a,s') = runState m s
                in seq s' a

newtype MaybeT m a = MaybeT { runMaybeT :: m (Maybe a) }

instance Monad m => Functor (MaybeT m)
  where
    fmap f x = pure f <*> x

instance Monad m => Applicative (MaybeT m)
  where
    pure = return
    (<*>) = ap

instance Monad m => Monad (MaybeT m) where
    return a = MaybeT $ return $ Just a
    m >>= f  = MaybeT $ do
                        res <- runMaybeT m
                        case res of
                                 Nothing -> return Nothing
                                 Just r  -> runMaybeT (f r)
    fail _   = MaybeT $ return Nothing

lift_maybe m = MaybeT $ do
                        res <- m
                        return $ Just res

-- apply a test to successive elements of a vector, evaluates to true iff test passes for all pairs
--successive_ :: Storable a => (a -> a -> Bool) -> Vector a -> Bool
successive_ t v = maybe False (\_ -> True) $ evalState (runMaybeT (mapVectorM_ stp (subVector 1 (size v - 1) v))) (v ! 0)
   where stp e  = do
                  ep <- lift_maybe $ state_get
                  if t e ep
                     then lift_maybe $ state_put e
                     else (fail "successive_ test failed")

-- operate on successive elements of a vector and return the resulting vector, whose length 1 less than that of the input
--successive :: (Storable a, Storable b) => (a -> a -> b) -> Vector a -> Vector b
successive f v = evalState (mapVectorM stp (subVector 1 (size v - 1) v)) (v ! 0)
   where stp  e = do
                  ep <- state_get
                  state_put e
                  return $ f ep e

succTest = utest "successive" $
       successive_ (>) (fromList [1 :: Float,2,3,4]) == True
    && successive_ (>) (fromList [1 :: Float,3,2,4]) == False
    && successive (+) (fromList [1..10 :: Float]) == 9 |> [3,5,7,9,11,13,15,17,19]

---------------------------------------------------------------------

findAssocTest = utest "findAssoc" ok
  where
    ok = m1 == m2
    m1 = assoc (6,6) 7 $ zip (find (>0) (ident 5 :: Matrix Float)) [10 ..] :: Matrix Float
    m2 = diagRect 7 (fromList[10..14]) 6 6

---------------------------------------------------------------------

condTest = utest "cond" ok
  where
    ok = step v * v == cond v 0 0 0 v
    v = fromList [-7 .. 7 ] :: Vector Float

---------------------------------------------------------------------

conformTest = utest "conform" ok
  where
    ok = 1 + row [1,2,3] + col [10,20,30,40] + (4><3) [1..]
         == (4><3) [13,15,17
                   ,26,28,30
                   ,39,41,43
                   ,52,54,56]

---------------------------------------------------------------------

accumTest = utest "accum" ok
  where
    x = ident 3 :: Matrix Float
    ok = accum x (+) [((1,2),7), ((2,2),3)]
         == (3><3) [1,0,0
                   ,0,1,7
                   ,0,0,4]
         &&
         toList (flatten x) == [1,0,0,0,1,0,0,0,1]

--------------------------------------------------------------------------------

convolutionTest = utest "convolution" ok
  where
--    a = fromList [1..10]               :: Vector Float
    b = fromList [1..3]                :: Vector Float
    c = (5><7) [1..]                   :: Matrix Float
--    d = (3><3) [0,-1,0,-1,4,-1,0,-1,0] :: Matrix Float
    ok =  separable (corr b) c == corr2 (outer b b) c
       && separable (conv b) c == conv2 (outer b b) c

--------------------------------------------------------------------------------

sparseTest = utest "sparse" (fst $ checkT (undefined :: GMatrix))

--------------------------------------------------------------------------------

staticTest = utest "static" (fst $ checkT (undefined :: L 3 5))

--------------------------------------------------------------------------------

intTest = utest "int ops" (fst $ checkT (undefined :: Matrix I))

--------------------------------------------------------------------------------

modularTest = utest "modular ops" (fst $ checkT (undefined :: Matrix (Mod 13 I)))

--------------------------------------------------------------------------------

indexProp g f x = a1 == g a2 && a2 == a3 && b1 == g b2 && b2 == b3
  where
    l = map g (toList (f x))
    a1 = maximum l
    b1 = minimum l
    a2 = x `atIndex` maxIndex x
    b2 = x `atIndex` minIndex x
    a3 = maxElement x
    b3 = minElement x

--------------------------------------------------------------------------------

_sliceTest = TestList
    [ testSlice (chol . trustSym)  (gen 5 :: Matrix R)
    , testSlice (chol . trustSym)  (gen 5 :: Matrix C)
    , testSlice qr    (rec :: Matrix R)
    , testSlice qr    (rec :: Matrix C)
    , testSlice hess  (agen 5 :: Matrix R)
    , testSlice hess  (agen 5 :: Matrix C)
    , testSlice schur (agen 5 :: Matrix R)
    , testSlice schur (agen 5 :: Matrix C)
    , testSlice lu    (agen 5 :: Matrix R)
    , testSlice lu    (agen 5 :: Matrix C)
    , testSlice (luSolve (luPacked (agen 5 :: Matrix R))) (agen 5)
    , testSlice (luSolve (luPacked (agen 5 :: Matrix C))) (agen 5)
    , test_lus (agen 5 :: Matrix R)
    , test_lus (agen 5 :: Matrix C)

    , testSlice eig   (agen 5 :: Matrix R)
    , testSlice eig   (agen 5 :: Matrix C)
    , testSlice (eigSH . trustSym) (gen 5 :: Matrix R)
    , testSlice (eigSH . trustSym) (gen 5 :: Matrix C)
    , testSlice eigenvalues   (agen 5 :: Matrix R)
    , testSlice eigenvalues   (agen 5 :: Matrix C)
    , testSlice (eigenvaluesSH . trustSym) (gen 5 :: Matrix R)
    , testSlice (eigenvaluesSH . trustSym) (gen 5 :: Matrix C)

    , testSlice svd           (rec :: Matrix R)
    , testSlice thinSVD       (rec :: Matrix R)
    , testSlice compactSVD     (rec :: Matrix R)
    , testSlice leftSV        (rec :: Matrix R)
    , testSlice rightSV       (rec :: Matrix R)
    , testSlice singularValues (rec :: Matrix R)

    , testSlice svd           (rec :: Matrix C)
    , testSlice thinSVD       (rec :: Matrix C)
    , testSlice compactSVD     (rec :: Matrix C)
    , testSlice leftSV        (rec :: Matrix C)
    , testSlice rightSV       (rec :: Matrix C)
    , testSlice singularValues (rec :: Matrix C)

    , testSlice (linearSolve (agen 5:: Matrix R)) (agen 5)
    , testSlice (flip linearSolve (agen 5:: Matrix R)) (agen 5)

    , testSlice (linearSolve (agen 5:: Matrix C)) (agen 5)
    , testSlice (flip linearSolve (agen 5:: Matrix C)) (agen 5)

    , testSlice (linearSolveLS (ogen 5:: Matrix R)) (ogen 5)
    , testSlice (flip linearSolveLS (ogen 5:: Matrix R)) (ogen 5)

    , testSlice (linearSolveLS (ogen 5:: Matrix C)) (ogen 5)
    , testSlice (flip linearSolveLS (ogen 5:: Matrix C)) (ogen 5)

    , testSlice (linearSolveSVD (ogen 5:: Matrix R)) (ogen 5)
    , testSlice (flip linearSolveSVD (ogen 5:: Matrix R)) (ogen 5)

    , testSlice (linearSolveSVD (ogen 5:: Matrix C)) (ogen 5)
    , testSlice (flip linearSolveSVD (ogen 5:: Matrix C)) (ogen 5)

    , testSlice (linearSolveLS (ugen 5:: Matrix R)) (ugen 5)
    , testSlice (flip linearSolveLS (ugen 5:: Matrix R)) (ugen 5)

    , testSlice (linearSolveLS (ugen 5:: Matrix C)) (ugen 5)
    , testSlice (flip linearSolveLS (ugen 5:: Matrix C)) (ugen 5)

    , testSlice (linearSolveSVD (ugen 5:: Matrix R)) (ugen 5)
    , testSlice (flip linearSolveSVD (ugen 5:: Matrix R)) (ugen 5)

    , testSlice (linearSolveSVD (ugen 5:: Matrix C)) (ugen 5)
    , testSlice (flip linearSolveSVD (ugen 5:: Matrix C)) (ugen 5)

    , testSlice ((<>) (ogen 5:: Matrix R)) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix R)) (ogen 5)
    , testSlice ((<>) (ogen 5:: Matrix C)) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix C)) (ogen 5)
    , testSlice ((<>) (ogen 5:: Matrix Float)) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix Float)) (ogen 5)
    , testSlice ((<>) (ogen 5:: Matrix (Complex Float))) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix (Complex Float))) (ogen 5)
    , testSlice ((<>) (ogen 5:: Matrix I)) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix I)) (ogen 5)
    , testSlice ((<>) (ogen 5:: Matrix Z)) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix Z)) (ogen 5)

    , testSlice ((<>) (ogen 5:: Matrix (I ./. 7))) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix (I ./. 7))) (ogen 5)
    , testSlice ((<>) (ogen 5:: Matrix (Z ./. 7))) (gen 5)
    , testSlice (flip (<>) (gen 5:: Matrix (Z ./. 7))) (ogen 5)

    , testSlice (flip cholSolve (agen 5:: Matrix R)) (chol $ trustSym $ gen 5)
    , testSlice (flip cholSolve (agen 5:: Matrix C)) (chol $ trustSym $ gen 5)
    , testSlice (cholSolve (chol $ trustSym $ gen 5:: Matrix R)) (agen 5)
    , testSlice (cholSolve (chol $ trustSym $ gen 5:: Matrix C)) (agen 5)

    , ok_qrgr        (rec :: Matrix R)
    , ok_qrgr        (rec :: Matrix C)
    , testSlice (test_qrgr 4 tau1) qrr1
    , testSlice (test_qrgr 4 tau2) qrr2
    ]
  where
    QR qrr1 tau1 = qrRaw (rec :: Matrix R)
    QR qrr2 tau2 = qrRaw (rec :: Matrix C)

    test_qrgr n t x = qrgr n (QR x t)

    ok_qrgr x = TestCase . assertBool "ok_qrgr" $ simeq 1E-15 q q'
      where
        (q,_) = qr x
        atau = qrRaw x
        q' = qrgr (rows q) atau

    simeq eps a b =  not $ magnit eps (norm_1 $ flatten (a-b))

    test_lus m = testSlice f lup
      where
        f x = luSolve (LU x p) m
        (LU lup p) = luPacked m

    gen :: Numeric t => Int -> Matrix t
    gen n = diagRect 1 (konst 5 n) n n

    agen :: (Numeric t, Num (Vector t))=> Int -> Matrix t
    agen n = gen n + fromInt ((n><n)[0..])

    ogen :: (Numeric t, Num (Vector t))=> Int -> Matrix t
    ogen n = gen n === gen n

    ugen :: (Numeric t, Num (Vector t))=> Int -> Matrix t
    ugen n = takeRows 3 (gen n)

    rec :: Numeric t => Matrix t
    rec = subMatrix (0,0) (4,5) (gen 5)

    testSlice f x@(size->sz@(r,c)) =
      TestList . map (TestCase . assertEqual "" (f x)) $ (map f (g y1 ++ g y2))
      where
        subm = subMatrix
        g y = [ subm (a*r,b*c) sz y | a <-[0..2], b <- [0..2]]
        h z = fromBlocks (replicate 3 (replicate 3 z))
        y1  = h x
        y2  = (tr . h . tr) x

--------------------------------------------------------------------------------

-- | All tests must pass with a maximum dimension of about 20
--  (some tests may fail with bigger sizes due to precision loss).
runTests :: Int  -- ^ maximum dimension
         -> IO ()
runTests n = do
    let test :: forall t . T.Testable t => t -> IO ()
        test p = qCheck n p
    putStrLn "------ index"
    test( \m -> indexProp id flatten (single (m :: RM)) )
    test( \v -> indexProp id id (single (v :: Vector Float)) )
    test( \m -> indexProp id flatten (m :: RM) )
    test( \v -> indexProp id id (v :: Vector Float) )
    test( \m -> indexProp magnitude flatten (single (m :: CM)) )
    test( \v -> indexProp magnitude id (single (v :: Vector (Complex Float))) )
    test( \m -> indexProp magnitude flatten (m :: CM) )
    test( \v -> indexProp magnitude id (v :: Vector (Complex Float)) )
    putStrLn "------ mult Float"
    test (multProp1 10 . rConsist)
    test (multProp1 10 . cConsist)
    test (multProp2 10 . rConsist)
    test (multProp2 10 . cConsist)
--    putStrLn "------ mult Float"
--    test (multProp1  6 . (single *** single) . rConsist)
--    test (multProp1  6 . (single *** single) . cConsist)
--    test (multProp2  6 . (single *** single) . rConsist)
--    test (multProp2  6 . (single *** single) . cConsist)
    putStrLn "------ sub-trans"
    test (subProp . rM)
    test (subProp . cM)
    putStrLn "------ ctrans"
    test (conjuTest . cM)
    test (conjuTest . zM)
    putStrLn "------ lu"
    test (luProp    . rM)
    test (luProp    . cM)
    putStrLn "------ inv (linearSolve)"
    test (invProp   . rSqWC)
    test (invProp   . cSqWC)
    putStrLn "------ luSolve"
    test (linearSolveProp (luSolve.luPacked) . rSqWC)
    test (linearSolveProp (luSolve.luPacked) . cSqWC)
    putStrLn "------ ldlSolve"
    test (linearSolvePropH (ldlSolve.ldlPacked) . rSymWC)
    test (linearSolvePropH (ldlSolve.ldlPacked) . cSymWC)
    putStrLn "------ cholSolve"
    test (linearSolveProp (cholSolve.chol.trustSym) . rPosDef)
    test (linearSolveProp (cholSolve.chol.trustSym) . cPosDef)
    putStrLn "------ luSolveLS"
    test (linearSolveProp linearSolveLS . rSqWC)
    test (linearSolveProp linearSolveLS . cSqWC)
    test (linearSolveProp2 linearSolveLS . rConsist)
    test (linearSolveProp2 linearSolveLS . cConsist)
    putStrLn "------ pinv (linearSolveSVD)"
    test (pinvProp  . rM)
    test (pinvProp  . cM)
    putStrLn "------ det"
    test (detProp   . rSqWC)
    -- test (detProp   . cSqWC)
    putStrLn "------ svd"
    test (svdProp1  . rM)
    test (svdProp1  . cM)
    test (svdProp1a svd . rM)
    test (svdProp1a svd . cM)
--    test (svdProp1a svdRd)
    test (svdProp1b svd . rM)
    test (svdProp1b svd . cM)
--    test (svdProp1b svdRd)
    test (svdProp2 thinSVD . rM)
    test (svdProp2 thinSVD . cM)
--    test (svdProp2 thinSVDRd)
--    test (svdProp2 thinSVDCd)
    test (svdProp3  . rM)
    test (svdProp3  . cM)
    test (svdProp4  . rM)
    test (svdProp4  . cM)
    test (svdProp5a)
    test (svdProp5b)
    test (svdProp6a)
    test (svdProp6b)
    test (svdProp7  . rM)
    test (svdProp7  . cM)
--    putStrLn "------ svdCd"
#ifdef NOZGESDD
--    putStrLn "Omitted"
#else
--    test (svdProp1a svdCd)
--    test (svdProp1b svdCd)
#endif
    putStrLn "------ eig"
    test (eigSHProp . rHer)
    test (eigSHProp . cHer)
    test (eigProp   . rSq)
    test (eigProp   . cSq)
    test (eigSHProp2 . rHer)
    test (eigSHProp2 . cHer)
    test (eigProp2   . rSq)
    test (eigProp2   . cSq)
    putStrLn "------ geig"
    test (uncurry geigProp . rSq2WC)
    test (uncurry geigProp . cSq2WC)
    putStrLn "------ nullSpace"
    test (nullspaceProp . rM)
    test (nullspaceProp . cM)
    putStrLn "------ qr"
    test (qrProp     . rM)
    test (qrProp     . cM)
    test (rqProp     . rM)
--    test (rqProp     . cM)
    test (rqProp1     . cM)
    test (rqProp2     . cM)
--    test (rqProp3     . cM)
    putStrLn "------ hess"
    test (hessProp   . rSq)
    test (hessProp   . cSq)
    putStrLn "------ schur"
    test (schurProp2 . rSq)
    test (schurProp1 . cSq)
    putStrLn "------ chol"
    test (cholProp   . rPosDef)
    test (cholProp   . cPosDef)
--    test (exactProp  . rPosDef)
--    test (exactProp  . cPosDef)
    putStrLn "------ expm"
    test (expmDiagProp . complex. rSqWC)
    test (expmDiagProp . cSqWC)
    putStrLn "------ vector operations - Float"
    test (\u -> sin u ^ 2 + cos u ^ 2 |~| (1::RM))
    test $ (\u -> sin u ^ 2 + cos u ^ 2 |~| (1::CM)) . liftMatrix makeUnitary
    test (\u -> sin u ** 2 + cos u ** 2 |~| (1::RM))
    test (\u -> cos u * tan u |~| sin (u::RM))
    test $ (\u -> cos u * tan u |~| sin (u::CM)) . liftMatrix makeUnitary
--    putStrLn "------ vector operations - Float"
--    test (\u -> sin u ^ 2 + cos u ^ 2 |~~| (1::FM))
--    test $ (\u -> sin u ^ 2 + cos u ^ 2 |~~| (1::ZM)) . liftMatrix makeUnitary
--    test (\u -> sin u ** 2 + cos u ** 2 |~~| (1::FM))
--    test (\u -> cos u * tan u |~~| sin (u::FM))
--    test $ (\u -> cos u * tan u |~~| sin (u::ZM)) . liftMatrix makeUnitary
    putStrLn "------ read . show"
    test (\m -> (m::RM) == read (show m))
    test (\m -> (m::CM) == read (show m))
    test (\m -> toRows (m::RM) == read (show (toRows m)))
    test (\m -> toRows (m::CM) == read (show (toRows m)))
    test (\m -> (m::FM) == read (show m))
    test (\m -> (m::ZM) == read (show m))
    test (\m -> toRows (m::FM) == read (show (toRows m)))
    test (\m -> toRows (m::ZM) == read (show (toRows m)))
    putStrLn "------ some unit tests"
    c <- runTestTT $ TestList
        [ utest "1E5 rots" rotTest
        , utest "det1" detTest1
        , utest "invlndet" detTest2
        , utest "expm1" (expmTest1)
        , utest "expm2" (expmTest2)
        , utest "arith1" $ ((ones (100,100) * 5 + 2)/0.5 - 7)**2 |~| (49 :: RM)
        , utest "arith2" $ ((scalar (1+iC) * ones (100,100) * 5 + 2)/0.5 - 7)**2 |~| ( scalar (140*iC-51) :: CM)
        , utest "arith3" $ exp (scalar iC * ones(10,10)*pi) + 1 |~| 0
        , utest "<\\>"   $ (3><2) [2,0,0,3,1,1::Float] <\> 3|>[4,9,5] |~| 2|>[2,3]
--        , utest "gamma" (gamma 5 == 24.0)
--        , besselTest
--        , exponentialTest
        , utest "randomGaussian" randomTestGaussian
        , utest "randomUniform" randomTestUniform
        , utest "buildVector/Matrix" $
                        complex (10 |> [0::Float ..]) == build 10 id
                     && ident 5 == build (5,5) (\r c -> if r==c then 1::Float else 0)
        , utest "rank" $  rank ((2><3)[1,0,0,1,5*peps,0::Float]) == 1
                       && rank ((2><3)[1,0,0,1,7*peps,0::Float]) == 2
        , utest "block" $ fromBlocks [[ident 3,0],[0,ident 4]] == (ident 7 :: CM)
        , mbCholTest
        , triTest
        , triDiagTest
        , triDiagRegression
        , utest "offset" offsetTest
        , normsVTest
        , normsMTest
        , sumprodTest
        , chainTest
        , succTest
        , findAssocTest
        , condTest
        , conformTest
        , accumTest
        , convolutionTest
        , sparseTest
        , staticTest
        , intTest
        , modularTest
        -- , sliceTest
        ]
    when (errors c + failures c > 0) exitFailure
    return ()

-- single precision approximate equality
-- infixl 4 |~~|
-- a |~~| b = a :~6~: b

makeUnitary v | realPart n > 1    = v / scalar n
              | otherwise = v
    where n = sqrt (v `dot` v)

binaryTests :: IO ()
binaryTests = do
  let test :: forall t . T.Testable t => t -> IO ()
      test = qCheck 100
  test vectorBinaryRoundtripProp
  test staticVectorBinaryRoundtripProp
  qCheck 30 matrixBinaryRoundtripProp
  qCheck 30 staticMatrixBinaryRoundtripProp

-- -- | Some additional tests on big matrices. They take a few minutes.
-- runBigTests :: IO ()
-- runBigTests = undefined

{-
-- | testcase for nonempty fpu stack
findNaN :: Int -> Bool
findNaN n = all (bugProp . eye) (take n $ cycle [1..20])
  where eye m = ident m :: Matrix ( Float)
-}

--------------------------------------------------------------------------------

-- | Performance measurements.
runBenchmarks :: IO ()
runBenchmarks = do
    solveBench
    subBench
    mkVecBench
    multBench
    cholBench
    luBench
    luBench_2
    svdBench
    eigBench
    putStrLn ""

--------------------------------

time msg act = do
    putStr (msg++" ")
    t0 <- getCPUTime
    act `seq` putStr " "
    t1 <- getCPUTime
    printf "%6.2f s CPU\n" $ (fromIntegral (t1 - t0) / (10^12 :: Float)) :: IO ()
    return ()

timeR msg act = do
    putStr (msg++" ")
    t0 <- getCPUTime
    putStr (show act)
    t1 <- getCPUTime
    printf "%6.2f s CPU\n" $ (fromIntegral (t1 - t0) / (10^12 :: Float)) :: IO ()
    return ()

--------------------------------

manymult n = foldl1' (<>) (map rot2 angles) where
    angles = toList $ linspace n (0,1)
    rot2 :: Float -> Matrix Float
    rot2 a = (3><3) [ c,0,s
                    , 0,1,0
                    ,-s,0,c ]
        where c = cos a
              s = sin a

multb n = foldl1' (<>) (replicate (10^6) (ident n :: Matrix Float))

--------------------------------

manyvec0 xs = sum $ map (\x -> x + x**2 + x**3) xs
manyvec1 xs = sumElements $ fromRows $ map (\x -> fromList [x,x**2,x**3]) xs
manyvec5 xs = sumElements $ fromRows $ map (\x -> vec3 x (x**2) (x**3)) xs

manyvec2 xs = sum $ map (\x -> sqrt(x^2 + (x**2)^2 +(x**3)^2)) xs
manyvec3 xs = sum $ map (norm_2 . (\x -> fromList [x,x**2,x**3])) xs

manyvec4 xs = sum $ map (norm_2 . (\x -> vec3 x (x**2) (x**3))) xs

vec3 :: Float -> Float -> Float -> Vector Float
vec3 a b c = runSTVector $ do
    v <- newUndefinedVector 3
    writeVector v 0 a
    writeVector v 1 b
    writeVector v 2 c
    return v

mkVecBench = do
    let n = 1000000
        xs = toList $ linspace n (0,1::Float)
    putStr "\neval data... "; print (sum xs)
    timeR "listproc        " $ manyvec0 xs
    timeR "fromList matrix " $ manyvec1 xs
    timeR "vec3 matrix     " $ manyvec5 xs
    timeR "listproc norm   " $ manyvec2 xs
    timeR "norm fromList   " $ manyvec3 xs
    timeR "norm vec3       " $ manyvec4 xs

--------------------------------

subBench = do
    putStrLn ""
    let g = foldl1' (.) (replicate (10^5) (\v -> subVector 1 (size v -1) v))
    time "0.1M subVector   " (g (konst 1 (1+10^5) :: Vector Float) ! 0)
    let f = foldl1' (.) (replicate (10^5) (fromRows.toRows))
    time "subVector-join  3" (f (ident  3 :: Matrix Float) `atIndex` (0,0))
    time "subVector-join 10" (f (ident 10 :: Matrix Float) `atIndex` (0,0))

--------------------------------

multBench = do
    let a = ident 1000 :: Matrix Float
    let b = ident 2000 :: Matrix Float
    a `seq` b `seq` putStrLn ""
    time "product of 1M different 3x3 matrices" (manymult (10^6))
    putStrLn ""
    time "product of 1M constant  1x1 matrices" (multb 1)
    time "product of 1M constant  3x3 matrices" (multb 3)
    --time "product of 1M constant  5x5 matrices" (multb 5)
    time "product of 1M const.  10x10 matrices" (multb 10)
    --time "product of 1M const.  15x15 matrices" (multb 15)
    time "product of 1M const.  20x20 matrices" (multb 20)
    --time "product of 1M const.  25x25 matrices" (multb 25)
    putStrLn ""
    time "product (1000 x 1000)<>(1000 x 1000)" (a<>a)
    time "product (2000 x 2000)<>(2000 x 2000)" (b<>b)

--------------------------------

eigBench = do
    let m = reshape 1000 (randomVector 777 Uniform (1000*1000))
        s = m + tr m
    m `seq` s `seq` putStrLn ""
    time "eigenvalues  symmetric 1000x1000" (eigenvaluesSH (trustSym m))
    time "eigenvectors symmetric 1000x1000" (snd $ eigSH (trustSym m))
    time "eigenvalues  general   1000x1000" (eigenvalues m)
    time "eigenvectors general   1000x1000" (snd $ eig m)

--------------------------------

svdBench = do
    let a = reshape 500  (randomVector 777 Uniform (3000*500))
        b = reshape 1000 (randomVector 777 Uniform (1000*1000))
        fv (_,_,v) = v `atIndex` (0,0)
    a `seq` b `seq` putStrLn ""
    time "singular values  3000x500" (singularValues a)
    time "thin svd         3000x500" (fv $ thinSVD a)
    time "full svd         3000x500" (fv $ svd a)
    time "singular values 1000x1000" (singularValues b)
    time "full svd        1000x1000" (fv $ svd b)

--------------------------------

solveBenchN n = do
    let x = uniformSample 777 (2*n) (replicate n (-1,1))
        a = tr x <> x
        b = asColumn $ randomVector 666 Uniform n
    a `seq` b `seq` putStrLn ""
    time ("svd solve " ++ show n) (linearSolveSVD a b)
    time (" ls solve " ++ show n) (linearSolveLS a b)
    time ("    solve " ++ show n) (linearSolve a b)
--    time (" LU solve " ++ show n) (luSolve (luPacked a) b)
    time ("LDL solve " ++ show n) (ldlSolve (ldlPacked (trustSym a)) b)
    time ("cholSolve " ++ show n) (cholSolve (chol $ trustSym a) b)

solveBench = do
    solveBenchN 500
    solveBenchN 1000
    solveBenchN 1500

--------------------------------

cholBenchN n = do
    let x = uniformSample 777 (2*n) (replicate n (-1,1))
        a = tr x <> x
    a `seq` putStr ""
    time ("chol " ++ show n) (chol $ trustSym a)

cholBench = do
    putStrLn ""
    cholBenchN 1200
    cholBenchN 600
    cholBenchN 300
--    cholBenchN 150
--    cholBenchN 50

--------------------------------------------------------------------------------

luBenchN f n x msg = do
    let m = diagRect 1 (fromList (replicate n x)) n n
    m `seq` putStr ""
    time (msg ++ " "++ show n) (rnf $ f m)

luBench = do
    putStrLn ""
    luBenchN luPacked  1000 (5::R)          "luPacked  Float    "
    luBenchN luPacked' 1000 (5::R)          "luPacked' Float    "
    luBenchN luPacked' 1000 (5::Mod 9973 I) "luPacked' I mod 9973"
    luBenchN luPacked' 1000 (5::Mod 9973 Z) "luPacked' Z mod 9973"

luBenchN_2 f g n x msg = do
    let m = diagRect 1 (fromList (replicate n x)) n n
        b = flipud m
    m `seq` b `seq` putStr ""
    time (msg ++ " "++ show n) (f (g m) b)

luBench_2 = do
    putStrLn ""
    luBenchN_2 luSolve  luPacked  500 (5::R)          "luSolve .luPacked  Float    "
    luBenchN_2 luSolve' luPacked' 500 (5::R)          "luSolve'.luPacked' Float    "
    luBenchN_2 luSolve' luPacked' 500 (5::Mod 9973 I) "luSolve'.luPacked' I mod 9973"
    luBenchN_2 luSolve' luPacked' 500 (5::Mod 9973 Z) "luSolve'.luPacked' Z mod 9973"


{-# LANGUAGE Strict #-}
module STC.Multiplication where

import           Control.Concurrent.Async
import           Control.Monad            as M
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Storable     as VS
import           DFT.Plan
import           Filter.Utils
import           STC.DFTArray
import           Utils.List
import           Utils.Parallel

{-# INLINE multiplyPinwheelBasis #-}
multiplyPinwheelBasis :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
multiplyPinwheelBasis plan (DFTArray rows cols thetaFreqs rFreqs vecs1) (DFTArray _ _ _ _ vecs2) = do
  dftVecs1 <-
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map
      (VS.convert .
       toUnboxed .
       computeS . makeFilter2D . fromUnboxed (Z :. cols :. rows) . VS.convert) $
    vecs1
  dftVecs2 <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs2
  let dftVecs3 = L.zipWith (VS.zipWith (*)) dftVecs1 dftVecs2
  DFTArray rows cols thetaFreqs rFreqs <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) dftVecs3

{-# INLINE multiplyPinwheelBasisBatch #-}
multiplyPinwheelBasisBatch ::
     DFTPlan -> Int -> DFTArray -> DFTArray -> IO DFTArray
multiplyPinwheelBasisBatch plan numBatch (DFTArray rows cols thetaFreqs rFreqs vecs1) (DFTArray _ _ _ _ vecs2) = do
  let forwardPlanID = DFTPlanID DFT1DG [cols, rows] [0, 1]
      backwardPlanID = DFTPlanID IDFT1DG [cols, rows] [0, 1]
      forwardPlan = getDFTPlan plan forwardPlanID
      backwardPlan = getDFTPlan plan backwardPlanID
  fmap (DFTArray rows cols thetaFreqs rFreqs . L.concat) .
    mapConcurrently
      (M.mapM
         (\(vec1, vec2) -> do
            dftVec1 <-
              dftExecuteWithPlan forwardPlanID forwardPlan .
              VS.convert .
              toUnboxed .
              computeS .
              makeFilter2D . fromUnboxed (Z :. cols :. rows) . VS.convert $
              vec1
            dftVec2 <- dftExecuteWithPlan forwardPlanID forwardPlan vec2
            dftExecuteWithPlan backwardPlanID backwardPlan $
              VS.zipWith (*) dftVec1 dftVec2)) .
    divideListN numBatch $
    L.zip vecs1 vecs2

{-# INLINE multiplyPinwheelBasisBatch1 #-}
multiplyPinwheelBasisBatch1 ::
     DFTPlan -> Int -> VS.Vector (Complex Double) -> DFTArray -> IO DFTArray
multiplyPinwheelBasisBatch1 plan numBatch dftVec1 (DFTArray rows cols thetaFreqs rFreqs vecs2) = do
  let forwardPlanID = DFTPlanID DFT1DG [cols, rows] [0, 1]
      backwardPlanID = DFTPlanID IDFT1DG [cols, rows] [0, 1]
      forwardPlan = getDFTPlan plan forwardPlanID
      backwardPlan = getDFTPlan plan backwardPlanID
  fmap (DFTArray rows cols thetaFreqs rFreqs . L.concat) .
    mapConcurrently
      (M.mapM
         (\vec2 -> do
            dftVec2 <- dftExecuteWithPlan forwardPlanID forwardPlan vec2
            dftExecuteWithPlan backwardPlanID backwardPlan $
              VS.zipWith (*) dftVec1 dftVec2)) .
    divideListN numBatch $
    vecs2


module Markov where

import           MusicData
import           Utility

-- import           Numeric.LinearAlgebra (Matrix, R, (><))

import           Data.Map              (Map)
import qualified Data.Map              as Map (empty, insertWith, lookup)
import qualified Data.Map.Merge.Strict as Map' (merge, preserveMissing,
                                                zipWithMatched)
import qualified Data.Map.Strict       as Strict (Map)
import qualified Data.Map.Strict       as Map' (difference, elems, empty,
                                                fromList, insert, insertWith,
                                                keys, lookup, member, toList)
import qualified Data.Maybe            as Maybe (fromMaybe)
import qualified Data.Set              as Set (fromList, size, toList)

-- |representation of bigrams and containing deterministic cadence sequences
type Bigram  = (Cadence, Cadence)

-- |representation of counts for each trigram
type TransitionCounts = Strict.Map Bigram Double

-- -- |representations of the Markov transition matrix
-- type TransitionMatrix = Matrix R

-- |representation of markov transition matrix as key-value pairs
type MarkovMap = Map Cadence [(Cadence, Double)]

-- |mapping from list of events into list of existing preceding bigrams
bigrams :: [a] -> [(a, a)]
bigrams (x:xs)
  | length (x:xs) < 2 = []
  | otherwise = bigram (x:xs) : bigrams xs
  where bigram (x:y:ys) = (\a b -> (a, b)) x y

-- |lifted 'shortcut' `toCadence` which operates on a list of integer lists
toCadences    :: (Integral a, Num a) => [[a]] -> [Cadence]
toCadences xs = toCadence <$> bigrams (flatTriad <$> xs)

-- |mapping from input data into all theoretically possible bigrams
pairs   :: [Cadence] -> [Bigram]
pairs xs =
  let cs  = unique xs
   in [ (x, y) | x <- cs, y <- cs ]

-- |mapping from input data into all possible trigrams with counts of zero
zeroCounts             :: [Cadence] -> TransitionCounts
zeroCounts xs           =
  let mInsert acc key = Map'.insert key 0 acc
   in foldl mInsert Map'.empty $ pairs xs

-- |mapping from input data to counts of all occurring transitions
cadenceCounts        :: [Cadence] -> TransitionCounts
cadenceCounts xs      =
  let mInsert acc key = Map'.insertWith (+) key 1 acc
   in foldl mInsert Map'.empty $ bigrams xs

-- |mapping from input to counts of cadences, including 'stationary' movements
transitionCounts       :: [Cadence] -> TransitionCounts
transitionCounts xs     = mergeMaps (foldl mInsert cadences (Map'.keys diff)) zeros
  where diff            = Map'.difference zeros cadences
        zeros           = zeroCounts xs
        cadences        = cadenceCounts xs
        mergeMaps m1 m2 = Map'.merge Map'.preserveMissing Map'.preserveMissing
                          (Map'.zipWithMatched (\k x y -> x)) m1 m2
        keys k          = [ (fst k, nxt) | nxt <- unique xs ]
        newKey k        = (fst k, fst k)
        member k        = sequenceA [ f ks | f <- [Map'.member], ks <- keys k ] (cadences)
        mInsert acc key
          | all (\x -> x == False) (member key) == True =
              Map'.insert (newKey key) 1 acc
          | otherwise   =
              Map'.insert key (Maybe.fromMaybe 0 $ Map'.lookup key $ cadences) acc

-- |helper function for probabilityList which generates probability sublists
transitionProbs      :: [Cadence] -> [Double] -> [[Double]]
transitionProbs _ []  = []
transitionProbs xs ys =
  let j               = Set.size $ Set.fromList xs
      i               = j^2
      xx              = take j ys
      recurse         = drop j ys
   in fmap (/ sum xx) xx : transitionProbs xs recurse

-- |mapping from list of Cadences into list of transitions with probabilities
probabilityMap   :: [Cadence] -> Map Bigram Double
probabilityMap xs = Map'.fromList $ zip (Map'.keys $ zeroCounts xs) $ concat
                     . transitionProbs xs $ Map'.elems $ transitionCounts xs

-- -- |mapping from list of cadences into transition matrix
-- transitionMatrix   :: [Cadence] -> TransitionMatrix
-- transitionMatrix xs =
--   let n             = Set.size $ Set.fromList xs
--    in (n><n) $ Map'.elems $ probabilityMap xs :: Matrix R

-- |mapping from list of cadences into map with possible next probabilities
markovMap              :: [Cadence] -> MarkovMap
markovMap xs            = foldl mInsert Map.empty $ pairs xs
  where mInsert acc key = Map.insertWith (++) (fst key) (pList key) acc
        pList key       = [(snd key, Maybe.fromMaybe 0 $ Map.lookup key pMap)]
        pMap            = probabilityMap xs


{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -Wincomplete-patterns #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module Num where

import AST
import Control.Monad.Except
import qualified Data.Complex as C

-- --------------------------------------------------------------------------
-- Pogger number
-- This module implements some conversion between pogger numbers.
-- ophan instance. Why not? Nobody will define another Num for PoggerNum
-- Somewhere else.
-- --------------------------------------------------------------------------

-- | convert a PoggerNum to Complex
poggerNumToComplex :: PoggerNum -> C.Complex Double
poggerNumToComplex (Integer n) = (realToFrac n) C.:+ 0.0
poggerNumToComplex (Real r) = r C.:+ 0.0
poggerNumToComplex (Rational a b) = (realToFrac a / realToFrac b) C.:+ 0.0
poggerNumToComplex (Complex a b) = a C.:+ b

poggerNumToDouble :: PoggerNum -> Double
poggerNumToDouble (Integer n) = fromIntegral n
poggerNumToDouble (Real r) = r
poggerNumToDouble (Rational a b) = fromIntegral a / fromIntegral b
poggerNumToDouble (Complex _ _) = error "Can't downgrade complex to real"

-- | Convert complex to pogger complex.
toPoggerComplex :: C.Complex Double -> PoggerNum
toPoggerComplex c = Complex (C.realPart c) (C.imagPart c)

-- | Generic binary operator that handles coersion rules.
poggerBinop :: (forall a. Num a => a -> a -> a) -> PoggerNum -> PoggerNum -> PoggerNum
poggerBinop op = mkop
  where
   mkop (Integer m) (Integer n) = Integer (m `op` n)
   mkop (Rational m n) (Rational m' n') = Rational (m `op` m') (n `op` n')
   mkop (Real r) (Real r') = Real (r `op` r')
   mkop (Complex m n) (Complex m' n') = Complex (m `op` m') (n `op` n')
   mkop (Rational a b) (Integer n) = simplify $ Rational (a `op` (n * b)) b
   mkop i@(Integer _) r@(Rational _ _) = r `op` i
   mkop m@(Complex _ _) n = toPoggerComplex (poggerNumToComplex m `op` poggerNumToComplex n)
   mkop m n@(Complex _ _) = m `op` n
   mkop m@(Real _) n = (Real . C.realPart) (poggerNumToComplex m `op` poggerNumToComplex n)
   mkop m n@(Real _) = m `op` n

{-# INLINE poggerBinop #-}

-- | absolute value of pogger num
-- >>> abs (Integer 2)
-- Integer 2
poggerUnary :: (forall a. Num a => a -> a) -> PoggerNum -> PoggerNum
poggerUnary op = \case
  Integer m -> Integer (op m)
  Rational m n -> Rational (op m) (op n)
  Real m -> Real (op m)
  c@(Complex _ _) ->
    let a = op (poggerNumToComplex c)
     in Complex (C.realPart a) (C.imagPart a)

instance Num PoggerNum where
  (+) = poggerBinop (+)
  {-# INLINE (+) #-}

  (*) = poggerBinop (*)
  {-# INLINE (*) #-}

  abs = poggerUnary abs
  {-# INLINE abs #-}

  negate = poggerUnary negate
  {-# INLINE negate #-}

  fromInteger = Integer
  {-# INLINE fromInteger #-}

  signum = poggerUnary signum
  {-# INLINE signum #-}

-- | division on pogger number
poggerNumDivide :: PoggerNum -> PoggerNum -> PoggerNum
poggerNumDivide (Real r) (Real r') = Real (r / r')
poggerNumDivide c@(Complex _ _) c'@(Complex _ _) =
  toPoggerComplex (poggerNumToComplex c / poggerNumToComplex c')
poggerNumDivide (Integer m) (Integer n) = (Integer . floor) ((fromIntegral m) / (fromIntegral n))
poggerNumDivide (Rational m n) (Rational m' n') = simplify (Rational (m * n') (n * m'))

-- >>> (abs . negate) (Rational 2 3 / Rational 5 3)
-- Rational 2 5
-- >>> Rational 2 3 / Rational 2 6
-- Integer 2
instance Fractional PoggerNum where
  (/) = poggerNumDivide
  {-# INLINE (/) #-}

  fromRational = Real . fromRational
  {-# INLINE fromRational #-}

-- | simply rational terms.
simplify :: PoggerNum -> PoggerNum
simplify (Rational a b) =
  let v = gcd a b
      r@(Rational m n) = Rational (a `div` v) (b `div` v)
   in if n == 1
        then Integer m
        else Rational m n
simplify _ = error "trying to simply non-rational"
{-# INLINE simplify #-}

-- | ord for pogger number
instance Ord PoggerNum where
  m <= n = poggerNumToDouble m <= poggerNumToDouble n
  {-# INLINE (<=) #-}

instance Real PoggerNum where
  toRational (Integer n) = fromIntegral n
  toRational (Real r) = toRational r
  toRational (Rational a b) = fromIntegral a / fromIntegral b
  toRational (Complex a _) = toRational a
  {-# INLINE toRational #-}

instance Enum PoggerNum where
  toEnum n = Integer (fromIntegral n)
  {-# INLINE toEnum #-}

  fromEnum (Integer n) = fromIntegral n
  fromEnum _ = (-1)
  {-# INLINE fromEnum #-}

-- | pogger
instance Integral PoggerNum where
  toInteger (Integer n) = n
  toInteger (Real r) = floor r
  toInteger (Rational a b) = floor (fromIntegral a / fromIntegral b)
  toInteger (Complex a _) = floor a
  {-# INLINE toInteger #-}

  quotRem (Integer a) (Integer b) =
    let (a', b') = (quotRem a b)
     in (Integer a', Integer b')
  quotRem a b = error ("no division algorithm defined for" ++ show a ++ show b)

-- | Safe wrappers
safeQuotRem :: PoggerNum -> PoggerNum -> ThrowsError (PoggerNum, PoggerNum)
safeQuotRem a@(Integer _) b@(Integer _) = return (quotRem a b)
safeQuotRem a _ = throwError (TypeMisMatch "can't mod on non integer values" (Number a))

safeMod :: PoggerNum -> PoggerNum -> ThrowsError PoggerNum
safeMod a b = safeQuotRem a b >>= return . snd

safeDiv :: PoggerNum -> PoggerNum -> ThrowsError PoggerNum
safeDiv a b = safeQuotRem a b >>= return . fst


{-# LANGUAGE OverloadedStrings, ExistentialQuantification, FlexibleContexts, TypeApplications, RecordWildCards #-}

module Main where
--monad-bayes
import Control.Monad
import Control.Monad.Bayes.Sampler
import Control.Monad.Bayes.Traced
import Control.Monad.Bayes.Weighted
import Control.Monad.Bayes.Class
--hvega
import qualified Graphics.Vega.VegaLite as VL
import qualified Graphics.Vega.Tutorials.VegaLite as VT
--data
import qualified Data.Histogram as DH
import Data.Text (Text, pack)
import Data.List (sort, partition)
import Data.Vector.Unboxed.Base (Unbox)
import qualified Data.Vector.Unboxed as Vec
--numeric
import Numeric.Log
import Numeric.LinearAlgebra (Matrix, Vector, toList, toInt, vector, matrix, dot, (#>), (!), cmap, scalar, size)
--src
import BNN
import PlotVL
--csv
import Control.Exception (IOException)
import qualified Control.Exception as Exception
import qualified Data.Foldable as Foldable
import Control.Applicative
import qualified Data.ByteString.Lazy as BL
import Data.Csv
import qualified Data.Vector as V

--Data Type
data Item =
  Item
    { idNumber :: Double,
      date :: String,
      new_cases :: Double
    }
  deriving (Eq, Show)

--Decoding
instance FromNamedRecord Item where
  parseNamedRecord m =
    Item
      <$> m .: "idNumber"
      <*> m .: "date"
      <*> m .: "new_cases"

decodeItems :: BL.ByteString -> Either String (Vector Item)
decodeItems = fmap snd . Data.Csv.decodeByName

decodeItemsFromFile :: FilePath -> IO (Either String (Vector Item))
decodeItemsFromFile filePath = catchShowIO (BL.readFile filePath) >>= return . either Left decodeItems

catchShowIO :: IO a -> IO (Either String a)
catchShowIO action = fmap Right action `Exception.catch` handleIOException
  where
    handleIOException :: IOException -> IO (Either String a)
    handleIOException = return . Left . show

fillingData :: BL.ByteString ->  [Double] -> [Data]
fillingData csvData noise = case decodeByName csvData of
  Left err -> putStrLn err
  Right (_,v) -> [Data (idNumber v) ((new_cases v) + n) | n <- noise]

main :: IO ()
main = do
    -- putStrLn "Enter the number of samples to train: "
    -- input1 <- getLine
    -- let nsamples = (read input1 :: Int)

    putStrLn "Enter the number of nodes: "
    input3 <- getLine
    let nnodes = (read input3 :: Int)

    putStrLn "Enter the number of predicted samples: "
    input4 <- getLine
    let npredictive = (read input4 :: Int)

    let nsamples = 365
    noise <- sampleIOfixed $ replicateM nsamples $ normal 0.0 0.5

    -- case decodeByName csvData of
    --   Left err -> putStrLn err
    --   Right (_, v) -> decodeItems

    csvData <- BL.readFile "/home/leduin/Desktop/10th Semester/Bayesian Statistics/Project/owid-covid-data-ecuador.csv"

    let observations = fillingData csvData noise

    -- let observations = case decodeByName csvData of
    --   Left err -> putStrLn err
    --   Right (_,v) -> [Data (idNumber v) ((new_cases v) + n) | n <- noise]

-- A SIMPLE NEURAL NETWORK

    --[ Data x (0.5 * x - 2 + n)
    --let observations = take nsamples
    --         [ Data x (2 * sin (x) + 1 + n)
    --         | (x, n) <- zip (map ((/(fromIntegral nsamples)) . fromIntegral ) [0, 10 ..]) noise
    --         ]

        --nnodes = 3
        --mkSampler = prior . mh 60000
        mkSampler = prior . mh npredictive
    predicted <-
        sampleIOfixed $ mkSampler $ predDist $
        postNN (priorNN nnodes) observations

    let hist = histo2D (0, 10, 10) (-10, 20, 10)
                ((\(nn, d) -> (xValue d, yValue d)) <$> predicted)
        cents = Vec.toList $ DH.binsCenters $ DH.bins hist
        val = Vec.toList $ DH.histData hist

    VL.toHtmlFile "final.html" $
        plot
            (600, 300)
            (L [imagePlot "xModel" "yModel" "zModel", scatterBlue "xObs" "yObs" (0, 10) (-10, 10)])
            ( Cols
                [ ("xModel", VL.Numbers (fst <$> cents)),
                ("yModel", VL.Numbers (snd <$> cents)),
                ("zModel", VL.Numbers val),
                ("xObs", VL.Numbers (xValue <$> observations)),
                ("yObs", VL.Numbers (yValue <$> observations))
                ]
            )

    putStrLn "Success! Check html images in directory"


module ImageProcess where

import Codec.Picture
import Data.Vector.Storable (Vector, length, toList)
import Numeric.LinearAlgebra
--import Data.ByteString

imagePath :: FilePath
imagePath = "./test.jpg"

getDynamicImage :: FilePath -> IO(Maybe DynamicImage)
getDynamicImage filepath = do
  res <- readImage filepath
  case res of
    Left err -> print err >> return Nothing
    Right image -> return (Just image)

readImageInfo :: DynamicImage -> Image PixelRGB8
readImageInfo image =
  let imagePixel8@(Image w h rgbs) = convertRGB8 image in
  imagePixel8

getImageInfo :: FilePath -> IO(Maybe (Image PixelRGB8))
getImageInfo filepath = do
  res <- readImage filepath
  case res of
    Left err -> print err >> return Nothing
    Right image -> return $ Just (readImageInfo image)

getImageWidth :: Image PixelRGB8 -> Int
getImageWidth image =
  let (Image w _ _) = image in
  w

getImageHeight :: Image PixelRGB8 -> Int
getImageHeight image =
  let (Image _ h _) = image in
  h

getImageRGBs :: Image PixelRGB8 -> Data.Vector.Storable.Vector(PixelBaseComponent Pixel8)
getImageRGBs image =
  let (Image _ _ rgbs) = image in
  rgbs

printPixels :: Show a => [a] -> IO()
printPixels [] = return()
printPixels (xr:xg:xb:xs) =
  print [xr, xg, xb] >>
  printPixels xs

genRGBList :: (Show a, Integral a, Fractional b) => [a] -> ([b], [b], [b])
genRGBList [] = ([], [], [])
genRGBList l =
  let func [] rlist glist blist = ((reverse rlist), (reverse glist), (reverse blist))
      func (xr:xg:xb:xs) rlist glist blist = func xs (((fromIntegral xr)/256.0):rlist) (((fromIntegral xg)/256.0):glist) (((fromIntegral xb)/256.0):blist) in
  func l [] [] []

genMatrixFromList :: [R] -> Int -> Int -> Maybe (Matrix R)
genMatrixFromList [] _ _ = Nothing
genMatrixFromList l w h = Just ((w >< h) l)

testImageProcess :: FilePath -> IO()
testImageProcess filepath = do
    imageM <- getImageInfo filepath
    case imageM of
      Nothing -> putStrLn "nothing.."
      Just imagePixel8 ->
        putStrLn "the width of image: " >>
        print (getImageWidth imagePixel8) >>
        putStrLn "the height ofimage: " >>
        print (getImageHeight imagePixel8) >>
        print (Data.Vector.Storable.length (getImageRGBs imagePixel8)) >>
        let rgbs = getImageRGBs imagePixel8 in
        let w = getImageWidth imagePixel8 in
        let h = getImageHeight imagePixel8 in
        let rgbsList = Data.Vector.Storable.toList rgbs in
        let (rlist, glist, blist) = genRGBList rgbsList in
        print (Prelude.length rlist) >>
        print (Prelude.length glist) >>
        print (Prelude.length blist) >>
        let Just rMat = genMatrixFromList rlist w h in
        let Just gMat = genMatrixFromList glist w h in
        let Just bMat = genMatrixFromList blist w h in
        print rMat

{-

decodeImageWidth :: IO(Maybe (Image PixelRGB8)) -> IO(Maybe Int))
decodeImageWidth imageM = do
  maybeImage <- imageM
  case maybeImage of
    Nothing -> return Nothing
    Just (Image w h rgbs) -> return (Just w)

decodeImageHeight :: IO(Maybe(Image PixelRGB8)) -> IO(Maybe Int)
decodeImageHeight imageM = do
  maybeImage <- imageM
  case maybeImage of
    Nothing -> return Nothing
    Just (Image w h rgbs) -> return (Just h)

decodeImageRGBs :: IO(Maybe (Image PixelRGB8)) -> IO(Maybe (Vector(PixelBaseComponent Pixel8)))
decodeImageRGBs imageM = do
  maybeImage <- imageM
  case maybeImage of
    Nothing -> return Nothing
    Just (Image w h rgbs) -> return (Just rgbs)

readRGBsPixel8 :: Maybe DynamicImage -> IO(Maybe (Vector(PixelBaseComponent Pixel8)))
readRGBsPixel8 Nothing = return Nothing
readRGBsPixel8 (Just image) =
  let imagePixel8@(Image _ _ rgbs) = convertRGB8 image in return (Just rgbs)

getRGBs :: FilePath -> IO(Maybe (Vector(PixelBaseComponent Pixel8)))
getRGBs filepath = do
  dynamicImage <- getDynamicImage filepath
  rgbs <- readRGBsPixel8 dynamicImage
  return rgbs

-}
{-}
readRGBsPixel8 :: DynamicImage -> Maybe (Vector (PixelBaseComponent Pixel8))
readRGBsPixel8 (ImageY8 image@(Image _ _ rgbs)) = Just rgbs
readRGBsPixel8 _ = Nothing

readRGBsPixel16 :: DynamicImage -> Maybe (Vector (PixelBaseComponent Pixel16))
readRGBsPixel16 :: (ImageY16 image@(Image _ _ rgbs)) = Just rgbs
readRGBsPixel16 _ = Nothing

readRGBsPixelYF :: DynamicImage -> Maybe (Vector (PixelBaseComponent PixelF))
readRGBsPixelYF :: (ImageF image@(Image _ _ rgbs)) = Just rgbs
readRGBsPixelYF _ = Nothing

readRGBsPixelYA8 :: DynamicImage -> Maybe (Vector (PixelBaseComponent PixelYA8))
readRGBsPixelYA8 :: (ImageYA8 image@(Image _ _ rgbs)) = Just rgbs
readRGBsPixelYA8 _ = Nothing

readRGBsPixelYA16 :: DynamicImage -> Maybe (Vector (PixelBaseComponent PixelYA16))
readRGBsPixelYA16 :: (ImageYA16 image@(Image _ _ rgbs)) = Just rgbs
readRGBsPixelYA16 _ = Nothing

getRGBsImage :: forall a. Pixel a => DynamicImage -> Vector (PixelBaseComponent a)
getRGBsImage (ImageY8 image@(Image _ _ rgbs@(MVector (Word8 _)))) = rgbs
getRGBsImage (ImageY16 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageYF image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageYA8 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageRGB8 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageRGB16 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageRGBF image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageRGBA8 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageRGBA16 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageYCbCr8 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageCMYK8 image@(Image _ _ rgbs)) = rgbs
getRGBsImage (ImageCMYK16 image@(Image _ _ rgbs)) = rgbs

readRGBs :: forall a. Pixel a => FilePath  -> IO(Maybe (Vector (PixelBaseComponent a)))
readRGBs filepath = do
  res <- getDynamicImage filepath
  case res of
    Nothing -> print "we got nothing!" >> return Nothing
    Just image -> return (Just (getRGBsImage image))

-}
{-
getRGBs :: Image -> Vector(PixelBaseComponent Word8)
getRGBs (Image _ _ rgbs) = rgbs
-}

justH :: Maybe Char
justH = do
  (x:xs) <- Just "Hello"
  return x


{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DerivingVia                #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}

module Backprop.Learn.Regularize (
    Regularizer
  , Regularize(..)
  , l1Reg
  , l2Reg
  , noReg
  , l2RegMetric
  , l1RegMetric
  -- * Build instances
  -- ** DerivingVia
  , RegularizeMetric(..), NoRegularize(..)
  -- ** Linear
  , lassoLinear, ridgeLinear
  -- ** Generics
  , grnorm_1, grnorm_2
  , glasso, gridge
  -- * Manipulate regularizers
  , addReg
  , scaleReg
  ) where

import           Control.Applicative
import           Control.Monad.Trans.State
import           Data.Ratio
import           Data.Semigroup hiding                 (Any(..))
import           Data.Type.Functor.Product
import           Data.Type.Tuple
import           Data.Vinyl
import           GHC.Exts
import           GHC.Generics
import           GHC.TypeNats
import           Generics.OneLiner
import           Numeric.Backprop                      as B
import           Numeric.LinearAlgebra.Static.Backprop ()
import           Numeric.Opto.Update hiding            ((<.>))
import qualified Data.Functor.Contravariant            as Co
import qualified Data.Vector.Generic                   as VG
import qualified Data.Vector.Generic.Sized             as SVG
import qualified Numeric.LinearAlgebra.Static          as H

-- | A regularizer on parameters
type Regularizer p = forall s. Reifies s W => BVar s p -> BVar s Double

-- | A class for data types that support regularization during training.
--
-- This class is somewhat similar to @'Metric' 'Double'@, in that it
-- supports summing the components and summing squared components.
-- However, the main difference is that when summing components, we only
-- consider components that we want to regularize.
--
-- Often, this doesn't include bias terms (terms that "add" to inputs), and
-- only includes terms that "scale" inputs, like components in a weight
-- matrix of a feed-forward neural network layer.
--
-- However, if all of your components are to be regularized, you can use
-- 'norm_1', 'norm_2', 'lassoLinear', and 'ridgeLinear' as sensible
-- implementations, or use DerivingVia with 'RegularizeMetric':
--
-- @
-- data MyType = ...
--   deriving Regularize via (RegularizeMetric MyType)
-- @
--
-- You can also derive an instance where /no/  are regularized, using
-- 'NoRegularize':
--
-- @
-- data MyType = ...
--   deriving Regularize via (NoRegularize MyType)
-- @
--
-- The default implementations are based on 'Generics', and work for types
-- that are records of items that are all instances of 'Regularize'.
class Backprop p => Regularize p where

    -- | Like 'norm_1': sums all of the weights in @p@, but only the
    -- ones you want to regularize:
    --
    -- \[
    -- \sum_w \lvert w \rvert
    -- \]
    --
    -- Note that typically bias terms (terms that add to inputs) are not
    -- regularized.  Only "weight" terms that scale inputs are typically
    -- regularized.
    --
    -- If @p@ is an instance of 'Metric', then you can set @'rnorm_1'
    -- = 'norm_1'@.  However, this would count /all/ terms in @p@, even
    -- potential bias terms.
    rnorm_1 :: p -> Double

    default rnorm_1 :: (ADT p, Constraints p Regularize) => p -> Double
    rnorm_1 = grnorm_1

    -- | Like 'norm_2': sums all of the /squares/ of the weights
    -- n @p@, but only the ones you want to regularize:
    --
    -- \[
    -- \sum_w w^2
    -- \]
    --
    -- Note that typically bias terms (terms that add to inputs) are not
    -- regularized.  Only "weight" terms that scale inputs are typically
    -- regularized.
    --
    -- If @p@ is an instance of 'Metric', then you can set @'rnorm_2'
    -- = 'norm_2'@.  However, this would count /all/ terms in @p@, even
    -- potential bias terms.
    rnorm_2 :: p -> Double

    default rnorm_2 :: (ADT p, Constraints p Regularize) => p -> Double
    rnorm_2 = grnorm_2

    -- | @'lasso' r p@ sets all regularized components (that is, components
    -- summed by 'rnorm_1') in @p@ to be either @r@ if that component was
    -- positive, or @-r@ if that component was negative.  Behavior is not
    -- defined if the component is exactly zero, but either @r@ or @-r@ are
    -- sensible possibilities.
    --
    -- It must set all /non-regularized/ components (like bias terms, or
    -- whatever items that 'rnorm_1' ignores) to zero.
    --
    -- If @p@ is an instance of @'Linear' 'Double'@ and 'Num', then you can set
    -- @'lasso' = 'lassoLinear'@.  However, this is only valid if 'rnorm_1'
    -- counts /all/ terms in @p@, including potential bias terms.
    lasso   :: Double -> p -> p

    default lasso :: (ADT p, Constraints p Regularize) => Double -> p -> p
    lasso = glasso

    -- | @'ridge' r p@ scales all regularized components (that is,
    -- components summed by 'rnorm_2') in @p@ by @r@.
    --
    -- It must set all /non-regularized/ components (like bias terms, or
    -- whatever items that 'rnorm_2' ignores) to zero.
    --
    -- If @p@ is an instance of @'Linear' 'Double'@ and 'Num', then you can set
    -- @'ridge' = 'ridgeLinear'@.  However, this is only valid if 'rnorm_2'
    -- counts /all/ terms in @p@, including potential bias terms.
    ridge   :: Double -> p -> p

    default ridge :: (ADT p, Constraints p Regularize) => Double -> p -> p
    ridge = gridge

grnorm_1 :: (ADT p, Constraints p Regularize) => p -> Double
grnorm_1 = getSum . gfoldMap @Regularize (Sum . rnorm_1)

grnorm_2 :: (ADT p, Constraints p Regularize) => p -> Double
grnorm_2 = getSum . gfoldMap @Regularize (Sum . rnorm_2)

glasso :: (ADT p, Constraints p Regularize) => Double -> p -> p
glasso r = gmap @Regularize (lasso r)

gridge :: (ADT p, Constraints p Regularize) => Double -> p -> p
gridge r = gmap @Regularize (ridge r)

-- | Backpropagatable L1 regularization; also known as lasso
-- regularization.
--
-- \[
-- \sum_w \lvert w \rvert
-- \]
--
-- Note that typically bias terms (terms that add to inputs) are not
-- regularized.  Only "weight" terms that scale inputs are typically
-- regularized.
l1Reg :: Regularize p => Double -> Regularizer p
l1Reg λ = liftOp1 . op1 $ \x ->
    ( λ * rnorm_1 x
    , (`lasso` x) . (* λ)
    )

-- | Backpropagatable L2 regularization; also known as ridge
-- regularization.
--
-- \[
-- \sum_w w^2
-- \]
--
-- Note that typically bias terms (terms that add to inputs) are not
-- regularized.  Only "weight" terms that scale inputs are typically
-- regularized.
l2Reg :: Regularize p => Double -> Regularizer p
l2Reg λ = liftOp1 . op1 $ \x ->
    ( λ * rnorm_2 x
    , (`ridge` x) . (* λ)
    )

-- | No regularization
noReg :: Regularizer p
noReg _ = auto 0

-- | A default implementation of 'lasso' for instances of @'Linear'
-- 'Double'@ and 'Num'.  However, this is only valid if the corresponding
-- 'rnorm_1' counts /all/ terms in @p@, including potential bias terms.
lassoLinear :: (Linear Double p, Num p) => Double -> p -> p
lassoLinear r = (r .*) . signum

-- | A default implementation of 'ridge' for instances of @'Linear'
-- 'Double'@.  However, this is only valid if the corresponding
-- 'rnorm_2' counts /all/ terms in @p@, including potential bias terms.
ridgeLinear :: Linear Double p => Double -> p -> p
ridgeLinear = (.*)

-- | L2 regularization for instances of 'Metric'.  This will count
-- all terms, including any potential bias terms.
--
-- You can always use this as a regularizer instead of 'l2Reg', if you want
-- to ignore the default behavior for a type, or if your type has no
-- instance.
--
-- This is what 'l2Reg' would be for a type @p@ if you declare an instance
-- of 'Regularize' with @'rnorm_2' = 'norm_2'@, and @'ridge'
-- = 'ridgeLinear'@.
l2RegMetric
    :: (Metric Double p, Backprop p)
    => Double                   -- ^ scaling factor (often 0.5)
    -> Regularizer p
l2RegMetric λ = liftOp1 . op1 $ \x ->
            ( λ * quadrance x, (.* x) . (* λ))

-- | L1 regularization for instances of 'Metric'.  This will count
-- all terms, including any potential bias terms.
--
-- You can always use this as a regularizer instead of 'l2Reg', if you want
-- to ignore the default behavior for a type, or if your type has no
-- instance.
--
-- This is what 'l1Reg' would be for a type @p@ if you declare an instance
-- of 'Regularize' with @'rnorm_1' = 'norm_1'@, and @'lasso'
-- = 'lassoLinear'@.
l1RegMetric
    :: (Num p, Metric Double p, Backprop p)
    => Double                   -- ^ scaling factor (often 0.5)
    -> Regularizer p
l1RegMetric λ = liftOp1 . op1 $ \x ->
            ( λ * norm_1 x, (.* signum x) . (* λ)
            )

-- | Add together two regularizers
addReg :: Regularizer p -> Regularizer p -> Regularizer p
addReg f g x = f x + g x

-- | Scale a regularizer's influence
scaleReg :: Double -> Regularizer p -> Regularizer p
scaleReg λ reg = (* auto λ) . reg

-- | Newtype wrapper (meant to be used with DerivingVia) to derive an
-- instance of 'Regularize' that uses its 'Metric' instance, and
-- regularizes every component of a data type, including any potential bias
-- terms.
newtype RegularizeMetric a = RegularizeMetric a
  deriving (Show, Eq, Ord, Read, Generic, Functor, Backprop)

instance (Metric Double p, Num p, Backprop p) => Regularize (RegularizeMetric p) where
    rnorm_1 = coerce $ norm_1 @_ @p
    rnorm_2 = coerce $ norm_2 @_ @p
    lasso   = coerce $ lassoLinear @p
    ridge   = coerce $ ridgeLinear @p

-- | Newtype wrapper (meant to be used with DerivingVia) to derive an
-- instance of 'Regularize' that does not regularize any part of the type.
newtype NoRegularize a = NoRegularize a
  deriving (Show, Eq, Ord, Read, Generic, Functor, Backprop)

instance Backprop a => Regularize (NoRegularize a) where
    rnorm_1 _ = 0
    rnorm_2 _ = 0
    lasso _   = B.zero
    ridge _   = B.zero

instance Regularize Double where
    rnorm_1 = id
    rnorm_2 = (** 2)
    lasso r = (r *) . signum
    ridge   = (*)

instance Regularize Float where
    rnorm_1 = realToFrac
    rnorm_2 = (** 2) . realToFrac
    lasso r = (realToFrac r *) . signum
    ridge   = (*) . realToFrac

instance Integral a => Regularize (Ratio a) where
    rnorm_1 = realToFrac
    rnorm_2 = (** 2) . realToFrac
    lasso r = (realToFrac r *) . signum
    ridge   = (*) . realToFrac

instance Regularize () where
    rnorm_1 _ = 0
    rnorm_2 _ = 0
    lasso _ _ = ()
    ridge _ _ = ()
instance (Regularize a, Regularize b) => Regularize (a, b)
instance (Regularize a, Regularize b, Regularize c) => Regularize (a, b, c)
instance (Regularize a, Regularize b, Regularize c, Regularize d) => Regularize (a, b, c, d)
instance (Regularize a, Regularize b, Regularize c, Regularize d, Regularize e) => Regularize (a, b, c, d, e)

instance (Regularize a, Regularize b) => Regularize (a :# b)
instance Regularize a => Regularize (TF a)

instance (RPureConstrained Regularize as, ReifyConstraint Backprop TF as, RMap as, RApply as, RFoldMap as) => Regularize (Rec TF as) where
    rnorm_1 = getSum
            . rfoldMap getConst
            . rzipWith coerce (rpureConstrained @Regularize (Co.Op rnorm_1))
    rnorm_2 = getSum
            . rfoldMap getConst
            . rzipWith coerce (rpureConstrained @Regularize (Co.Op rnorm_2))
    lasso r = rzipWith coerce (rpureConstrained @Regularize (Endo (lasso r)))
    ridge r = rzipWith coerce (rpureConstrained @Regularize (Endo (ridge r)))

instance (PureProdC Maybe Backprop as, PureProdC Maybe Regularize as) => Regularize (PMaybe TF as) where
    rnorm_1 = getSum
            . foldMapProd getConst
            . zipWithProd coerce (pureProdC @_ @Regularize (Co.Op rnorm_1))
    rnorm_2 = getSum
            . foldMapProd getConst
            . zipWithProd coerce (pureProdC @_ @Regularize (Co.Op rnorm_2))
    lasso r = zipWithProd coerce (pureProdC @_ @Regularize (Endo (lasso r)))
    ridge r = zipWithProd coerce (pureProdC @_ @Regularize (Endo (ridge r)))

instance (VG.Vector v a, Regularize a, Backprop (SVG.Vector v n a)) => Regularize (SVG.Vector v n a) where
    rnorm_1 = (`execState` 0) . SVG.mapM_ (modify . (+) . rnorm_1)
    rnorm_2 = (`execState` 0) . SVG.mapM_ (modify . (+) . rnorm_2)
    lasso r = SVG.map (lasso r)
    ridge r = SVG.map (ridge r)

deriving via (RegularizeMetric (H.R n)) instance KnownNat n => Regularize (H.R n)
deriving via (RegularizeMetric (H.L n m)) instance (KnownNat n, KnownNat m) => Regularize (H.L n m)


module Lights where

import Types
import Numeric.Vector
import Numeric.Scalar (fromScalar, scalar)
import SceneRandom
import Control.Monad.State
import Raycast

sampleLight :: Vec3d -> Vec3d -> Light -> State SceneContext (Vec3d, Vec3d)

sampleLight start normal (Directional angle strength direction color)  = do
  dir <- randVecCone (angle * pi / 180) direction
  let incidence = normal · negate dir
  ss <- get
  let scene = ss_getScene ss
  case raycast scene (Ray start (negate dir)) of
    Just _ -> return $ (vec3 0 0 0, negate dir)
    Nothing -> return $ (color * (fromScalar $ incidence * scalar strength), negate dir)

-- TODO: Point light sampling
sampleLight _ _ _ = return $ (vec3 0 0 0, vec3 0 0 0)


{-# LANGUAGE TypeFamilies #-}
module Augmentation.Braid
    (AugBraid (AugBraid)
    ,relations
    ,sChar
    ,s
    ,relations'
    ) where

import Algebra
import Augmentation.DGA
import Braid.Class
import Libs.List

import Data.List
import Data.Either
import Data.Maybe
import Control.Monad
import Data.Functor.Identity

import Numeric.LinearAlgebra

import Debug.Trace

sChar :: Int -> Char
sChar 0 = error "s 0"
sChar k = toEnum $ (2*k)+459

invsChar :: Char -> Int
invsChar c = if fromEnum c > 459 then (fromEnum c - 459) `div` 2 else error "invsChar: c out of bounds"

s :: Int -> Algebra
s = G . E . sChar
{-
sVec :: Int -> Vector R
sVec i = i |> [if j == (i-1) then 1 else 0 | j <- [1..]]
-}
bph :: StdBraid -> Int -> [Int]
bph b@(StdBraid k w) i = nub $ sort $ take k $ iterate (bphh w) i

bphh :: [Int] -> Int -> Int
bphh [] i = i
bphh (w:ws) i
    | w == i = bphh ws (i+1)
    | w == (i-1) = bphh ws (i-1)
    | otherwise = bphh ws i

extractChar :: Algebra -> Maybe Char
extractChar (G (E c)) = Just c
extractChar e = Nothing

sVec :: Int -> Int -> Vector Z
sVec dim i = dim |> [toEnum $ if j == i then 1 else 0 | j <- [1..]]

eqVec :: StdBraid -> Int -> Maybe (Vector Z)
eqVec (StdBraid w ws) i = do
                            { let dim = length ws
                            ; let v = sum $ zipWith (\w x -> if w == i then sVec dim x else if w == (i-1) then negate $ sVec dim x else 0) ws [1..]
                            ; if i > w then Nothing else Just v
                            }

relations' :: StdBraid -> Matrix Z
relations' b@(StdBraid w ws) = maybe (ident $ length ws) id $ do
                                { let basePoints = map head $ nub $ map sort $ map (bph b) [1..w]
                                ; let dim = length ws
                                ; vs <- mapM (eqVec b) $ filter (\i -> not $ i `elem` basePoints) [1..w]
                                ; mat <- if vs == [] then Nothing else Just $ fromColumns vs
                                ; let (q,_) = thinQR $ (fromZ mat :: Matrix R)
                                ; let qqt = q Numeric.LinearAlgebra.<> (tr q)
                                ; n <- (\(a,b) -> if a == b then Just a else Nothing) $ size qqt
                                ; return $ fromColumns $ map (toZ . roundVector) $ toColumns $ 2 * ((ident n) - qqt)
                                }

{-
addVec :: Vector R -> Vector R -> Vector R
addVec v1 v2
    | diff <  0 = addVec v2 v1
    | diff == 0 = v1 + v2
    | otherwise = v1 + (fromList $ toList v2 ++ zers)
    where diff = (size v1) - (size v2)
          zers = take diff repeat 0

eqVec :: StdBraid -> Int -> Maybe (Vector R)
eqVec (StdBraid w ws) i = do
                            { let v = sum $ zipWith (\w x -> if w == i then sVec x else if w == (i-1) then negate $ sVec x else fromInteger 0) ws [1..]
                            ; if i > w then Nothing else Just v
                            }

relations' :: StdBraid -> Maybe (Matrix R)
relations' b@(StdBraid w ws) = do
                                { let vecs = map (eqVec b) [1..w]
                                ; dim <- if [V.length $ head vecs] == (nub $ map V.length vecs) then Just $ V.length $ head vecs else Nothing
                                ; let mat = fmap fromIntegral $ M.transpose $ foldr (\x m -> (colVector x) <|> m) dim vecs
                                ; let zer = zeros (nrows mat) (ncols mat)
                                ; (_,l,p,_) <- luDecomp mat
                                ;

-}
eqh :: StdBraid -> Int -> Maybe (DGA_Map,DGA_Map,DGA_Map,DGA_Map)
eqh (StdBraid _ ws) i = do
                            { let word = zipWith (\w x ->
                                        if w == i
                                            then s x
                                            else if w == (i-1) then recip $ s x
                                                               else 1) ws [1..]
                            ; chars <- mapM (\(x,t) ->
                                    (extractChar x) >>= (\x' -> return (x',t))) $ catMaybes $ map (\(x,t) ->
                                        if x == i then Just (s t,True) else if x == i-1 then Just (s t,False) else Nothing) $ zip ws [1..]
                            ; let (s1,r1) = head chars
                            ; let (s2,r2) = last chars
                            ; let tword = tail word
                            ; let iword = init word
                            ; let prod l = if l == [] then 1 else (head l) * (prod $ tail l)
                            ; let g1 = if r1 then recip else id
                            ; let g2 = if r2 then recip else id
                            ; let m11 = DGA_Map [(s1,g1 $ prod tword)]
                            ; let m12 = DGA_Map [(s1,g1 $ prod $ reverse $ tword)]
                            ; let m21 = DGA_Map [(s2,g2 $ prod iword)]
                            ; let m22 = DGA_Map [(s2,g2 $ prod $ reverse $ iword)]
                            ; return (m11,m12,m21,m22)
                            }

nst :: Int -> (a,a,a,a) -> Maybe a
nst 0 (a,_,_,_) = return a
nst 1 (_,a,_,_) = return a
nst 2 (_,_,a,_) = return a
nst 3 (_,_,_,a) = return a
nst _ _ = Nothing

relations :: StdBraid -> Maybe [DGA_Map]
relations b@(StdBraid w _) = do
                                { let basePoints = map head $ nub $ map sort $ map (bph b) [1..w]
                                ; let solutions = catMaybes $ map (eqh b) $ filter (\x -> not $ x `elem` basePoints) [1..w]
                                --; let toPerm l = foldr (\x xs -> (x `mod` 4) * (4^(length xs)) + xs) l
                                ; let fromPerm m n = if m == 0 then [] else (fromPerm (m-1) ((n - (n `mod` 4)) `div` 4)) ++ [n `mod` 4]
                                ; let perms = map (fromPerm (length solutions)) [0..(4^(length solutions))-1]
                                ; let maps = nub $ catMaybes $ map (\l -> do
                                                    { ms <- zipWithM nst l solutions
                                                    ; let none = DGA_Map []
                                                    ; let left = foldl compose_maps none ms
                                                    ; let right = foldr compose_maps none ms
                                                    ; if left == right && (and $ map (\x -> x == compose_maps x x) ms) then return right else Nothing
                                                    }) perms
                                ; return maps
                                }
{-
eqh :: StdBraid -> Int -> Maybe ((Char,(Algebra,Algebra)),(Char,(Algebra,Algebra)))
eqh b@(StdBraid _ ws) i = do
                            { let word = zipWith (\w x -> if w == i
                                    then s x
                                    else if w == (i-1) then recip $ s x
                                                       else 1) ws [1..]
                            ; c0 <- mapM (\(x,t) -> if x == i then Just (s t,False) else if x == i-1 then Just (s t,True) else Nothing) $ filter (\(x,_) -> (x == i) || (x == i-1)) $ zip ws [1..]
                            ; chars <- mapM (\(x,b) -> (extractChar x) >>= (\x' -> return $ (x',b))) c0
                            ; let (s1,r1) = head chars
                            ; let (s2,r2) = last chars
                            ; let tword = tail word
                            ; let iword = init word
                            ; let pre = ((s1,(product $ tword, product $ reverse $ tword)),(s2,(product $ iword,product $ reverse $ iword)))
                            ; let g1 = if r1 then recip else id
                            ; let g2 = if r2 then recip else id
                            ; let f ((c1,(e11,e12)),(c2,(e21,e22))) = ((c1,(g1 e11,g1 e12)),(c2,(g2 e21, g2 e22)))
                            ; return $ f pre
                            }

relations :: StdBraid -> Maybe [DGA_Map]
relations b@(StdBraid w _) = do
                                { let basePoints = map head $ nub $ map sort $ map (bph b) [1..w]
                                ; solutions <- mapM (eqh b) $ filter (\x -> not $ x `elem` basePoints) [1..w]
                                ; let solution = catMaybes $ map (\((c1,(a11,a12)),(c2,(a21,a22))) -> if c2 == c1 then Nothing else Just [(c1,a11),(c1,a12),(c2,a21),(c2,a22)]) solutions
                                ; let splited = map (\x -> compose_maps x x) $ map (\x -> DGA_Map $ zipWith (\l x' -> l !! x') solution x) $ permutations [0..3]
                                ; let fixed = nub $ filter (\x -> x == (compose_maps x x)) splited
                                ; return fixed
                                }
-}
buildMaph :: Maybe [(Char,Algebra)] -> Maybe [(Char,Algebra)] -> Maybe [(Char,Algebra)]
buildMaph m m'
    | m == m' = m
    | otherwise = buildMaph (m >>= (\l -> Just $ map (\(c,x) -> (c,applyDGAMap (DGA_Map l) x)) l)) m

footprinth :: Int -> [Either (Int,Char) Int] -> [(FreeGroup,Int)]
footprinth _k [] = []
footprinth k ((Right i):xs) = (E $ toEnum k,i):(footprinth (k+1) xs)
footprinth k ((Left (i,c)):xs) = [(E c,i-1),(invert $ E c,i+1)] ++ (footprinth (k+1) xs)

iscrossh :: [Either (Int,Char) Int] -> Int -> Bool
iscrossh [] _ = False
iscrossh ((Left _):_) 0 = False
iscrossh ((Left _):_) 1 = False
iscrossh ((Right _):_) 0 = True
iscrossh ((Right _):cs) x = iscrossh cs (x-1)
iscrossh ((Left _):cs) x = iscrossh cs (x-2)

data AugBraid = AugBraid Int [Either (Int,Char) Int] -- Either element of H1(L), and its position or an integer representing the corresponding braid group element
instance Braid AugBraid where
    type M AugBraid = Either (Int,Char)
    get_word (AugBraid _ w) = w
    algebra_footprint (AugBraid _ w) = map (\(g,i) -> (G g,i)) $ footprinth 65 w
    toStdBraid (AugBraid w ws) = StdBraid w (rights ws)
    fromStdBraid (StdBraid w ws) = AugBraid w (map Right ws)
    isCross (AugBraid _ ws) x = iscrossh ws x
    cross_art b row (Right s) = cross_art (toStdBraid b) row (Identity s)
    cross_art _b row (Left (s,s')) = let s1 = head $ show $ E s'
                                         s2 = show $ invert $ E s'
                                      in if (row - (abs $ s-1)*3) `elem` [0..3]
                        then Just $ (case ((row-(s-1)*3) `mod` 4) of 0 -> "---" ++ [s1] ++ "----"
                                                                     1 -> "   " ++  " " ++ "    "
                                                                     2 -> "   " ++  " " ++ "    "
                                                                     3 -> "--" ++ s2 ++ "---")
                        else Nothing
instance Eq AugBraid where
    (==) = equal
instance Show AugBraid where
    show = showBraid


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Approximate median
-- | Creator: Xiao Ling
-- | Created: 12/08/2015
-- | TODO   : test standard deviation of alpha, beta, and final version
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module ApproxMedian (

      MaxSize
    , Store

    , aMedian
    , aMedian'

  ) where

import Control.Monad.Trans
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Identity

import Data.Random
import Data.Conduit
import Data.Sequence
import Data.Foldable (toList)
import qualified Data.Conduit.List as Cl

import Core
import Statistics

{-----------------------------------------------------------------------------
  Types
------------------------------------------------------------------------------}

type MaxSize = Int
type Store a = (Seq a, MaxSize)

-- * Construct an empty store given max size `s`
store :: MaxSize -> Store a
store = (,) empty

{-----------------------------------------------------------------------------
  Median II
------------------------------------------------------------------------------}

-- * determine the `e`ps -approx median of a list `xs` with confidence `d`elta.
aMedian :: (Ord a, Floating a) => Eps -> Delta -> [a] -> IO a
aMedian e d xs = eval $ Cl.sourceList xs $$ toMedian $ store s
  where eval p = fmap (median . toList) . fmap fst . flip runRVar StdRandom $ evalStateT p 0
        s      = round $ 7/(e^2) * log (2/d)

-- * `tick` up a counter for each item `a` seen
-- * and with probability min(1, s/i) put item `a` into store `t`
toMedian :: Store a -> Sink a (StateT Counter RVar) (Store a)
toMedian t = Cl.foldM (\t a -> tick >> t $|> a) t

-- * count items seen so far
tick :: Enum s => MonadState s m => m ()
tick = modify succ

-- * with probability `s/i` uniformly select an item from the store `t` and replace it with `a`
-- * mnemonic: `|>` is insertion for Sequences, and `$f` means do f probabilistically
($|>) :: Store a -> a -> StateT Counter RVar (Store a)
($|>) t@(as,s) a = do

  i <- get

  let p = s ./ i
  if p >= 1 then return (as |> a, s) else do

      h <- lift . toss $ coin p

      if isHead h then do
          x <- lift (uniform 0 s :: RVar Int)
          let as' = update x a as
          return (as',s)
      else return t

-- * Sledghammer division
(./) :: (Fractional a, Read a, Show s, Show i) => s -> i -> a
(./) s i = (read . show $ s)/(read . show $ i)

{-----------------------------------------------------------------------------
  Median II
------------------------------------------------------------------------------}

-- * determine the `e`ps - approx median of a list `xs` with confidence `d`elta.
aMedian' :: (Ord a, Floating a) => Eps -> Delta -> [a] -> IO a
aMedian' e d xs = eval $ Cl.sourceList xs $$ toMedian'
  where eval p = fmap (median . toList) . fmap fst $ runRVar (evalStateT (runStateT p 0) (store s)) StdRandom
        s      = round $ 7/(e^2) * log (2/d)

-- * `tick` up a counter `i` of items seen so far
-- * and with uniform probability put the item `a` into a store `as` of max size `s`
toMedian' :: Sink a (StateT Counter (StateT (Store a) RVar)) (Seq a)
toMedian' = tick' >> do
  ma <- await
  case ma of
    Nothing -> (lift . lift $ get) >>= return . fst
    Just a  -> (lift . place $ a ) >>  toMedian'

tick' :: (Enum s, MonadTrans t, MonadState s m) => t m ()
tick' = lift . modify $ succ

-- * with uniform probability `s/i` select an item from the
-- * store and re`place` it with `a`
place :: a -> StateT Counter (StateT (Store a) RVar) ()
place a = do

  i        <- get
  t@(as,s) <- lift get

  let p = s ./ i

  if p >= 1 then lift $ put (as |> a, s) else do

      h <- lift . lift . toss $ coin p
      case isHead h of
        False -> lift $ put (as,s)
        _     -> do
          x <- lift $ lift (uniform 0 s :: RVar Int)
          let as' = update x a as
          lift $ put (as',s)

{-----------------------------------------------------------------------------
  Awaiting SO response for this stuff
------------------------------------------------------------------------------}

instance MonadState s m => MonadState s (RVarT m) where
  get   = lift get
  put   = lift . put
  state = lift . state

foo :: (MonadState s m, MonadRandom m) => s -> RVarT m ()
foo s = do
  x <- (uniformT 0 1 :: RVarT m Double)
  if x < 0.5 then put s else return ()

foo' :: (MonadState String m, MonadRandom m) => RVarT m ()
foo' = foo "hello"

--foo'' :: (MonadState String m, MonadRandom m) => m ()
--foo''= runRVarT foo' StdRandom

bar :: (MonadTrans t1, MonadState t m, MonadState s (t1 m)) => s -> RVarT (t1 m) ()
bar s = do
  x <- (uniformT 0 1 :: RVarT m Double)
  a <- lift . lift $ get
  if x < 0.5 then put s else return ()

-- * desired behavior with state and ranndomness

baz :: (MonadState s m, MonadReader t m) => m ()
baz = get >> ask >> return ()

baz1 :: MonadReader t m => m ((),Int)
baz1 = runStateT baz 0

baz1' :: Identity ((),Int)
baz1' = runReaderT baz1 (1,"hello")

baz2 :: MonadState s m => m ()
baz2 = runReaderT baz "hello"

baz2' :: Identity ((),[Int])
baz2' = runStateT baz2 [1..10]


{-|
Module      : ODMatrix.ElementaryDecomposition.ShortestPath
Description : This module is responsible for computing the shortest path between two ODMs.

Contains the necesary functions to measure and compare ODMs.

-}
module ODMatrix.ElementaryDecomposition.ShortestPath (
      Measure
    , sePath
    , ev
    , ev3
  ) where

  import Control.Monad.State.Lazy
  --import Control.Monad.Loops (iterateUntilM)

  import Numeric.LinearAlgebra

  import qualified Data.HashMap.Lazy as M
  import qualified Data.Hashable as H
  import Data.Maybe

  import ODMatrix.ElementaryDecomposition (ElementaryMatrix(..), ODM,
                                       getChildren, applyElementary,
                                       applyPath,
                                       applicables, elementaryValue,
                                       opposite, isNullElementary)

  -- import Debug.Trace (trace)

  -- | Measure or Metric prototype
  type Measure = Int -> Int -> Double

  -- | Elementary Value.
  -- Given a metric m, a in route metric mr and a scale factor alpha returns a elementary value metric.
  ev :: Double    -- ^ Scale factor alpha
     -> Measure   -- ^ Metric m
     -> Measure   -- ^ In route metric mr
     -> Measure   -- ^ Elementary value metric
  ev a m mr i j
    | a < 0 = error "The ajustment parameter must be greater than 0."
    | otherwise = m i j + a * mr i j

  -- | Shortest Path
  -- Computes the sequence of elementary matrices between ODMs
  sePath :: Measure             -- ^ Elementary value metric
         -> ODM                 -- ^ Source ODM
         -> ODM                 -- ^ Target ODM
         -> [ElementaryMatrix]  -- ^ Sequence of elementary matrices in the shortest path
  sePath m a b = evalState _process (_initialSEPState m a b)

  -- | Map a Elementary Matrix to it related cells indexes
  emc :: ElementaryMatrix -> (Int, Int)
  emc (Elementary _ (cell1, _) (cell2, _)) = (cell1,cell2)

  -- | Compares two ODMs
  odmc :: (Matrix Double -> Double) -- ^ Function g to compute the remainder
       -> Measure                   -- ^ Elementary value metric
       -> Double                    -- ^ Scale factor
       -> ODM                       -- ^ ODM A
       -> ODM                       -- ^ ODM B
       -> Double                    -- ^ Difference between A and B
  odmc g ev beta a b = beta * g s + sum [ev i j | (i,j) <- sp]
    where p = sePath ev a b
          sp = map emc . filter (not . isNullElementary) $ p
          s = b - applyPath a p -- Remainder

  newtype HashMatrix = HashMatrix { hm :: ODM }

  instance Show HashMatrix where
    show (HashMatrix m) = show m

  instance Eq HashMatrix where
    (==) (HashMatrix a) (HashMatrix b) = a == b

  instance H.Hashable HashMatrix where
    hashWithSalt salt (HashMatrix m) = floor $ m `atIndex` (salt `mod` rows m, salt `mod` cols m)

  data SEP = SEP {
      dist      :: M.HashMap HashMatrix Double
    , prev      :: M.HashMap HashMatrix ElementaryMatrix
    , visited   :: [ODM]
    , target    :: ODM
    , measure   :: Measure
    }

  instance Show SEP where
    show = show . visited

  -- | Elementary Value 3 (See the paper).
  -- This is a special case of elementary value computed with only one metric.
  -- Only works for regular grids.
  ev3 :: Double -> Measure -> Measure
  ev3 a m i j
    | a <= 0 || a >= 1 = error "The ajustment parameter a must be 0 < a < 1."
    | otherwise = a * m i j + (1-a) * sum [ m k (k+1) | k <- [min i j .. max i j] ]

  -- TODO: Implement metrics.

  _initialSEPState :: Measure -> ODM -> ODM -> SEP
  _initialSEPState m src trg = SEP (M.singleton (HashMatrix src) 0) M.empty [] trg m

  -- | TODO: take out the harcoded 60.
  -- | MAYBE IF IN THE FINAL STEP ASK FOR NEIGHBORS EMPTY.
  _process :: State SEP [ElementaryMatrix]
  _process = do
    s <- get
    let (curr,d) = _getNext (dist s)
        neighbors = applicables (0,60) curr . getChildren $ (target s) - curr
        -- Filter out the visited neighbors
        unvisited = filter (\e -> applyElementary curr e `notElem` visited s) neighbors
    -- Calculate min distance to all neighbors (unvisited)
    modify $ minAndPrev curr d unvisited
    -- \t -> t {dist = foldr (mindist (measure t) curr d) (dist t) unvisited}
    -- Add current to the previous of unvisited

    -- Mark the current as visited.
    modify $ \t -> t {dist = M.delete (HashMatrix curr) (dist t),
                      visited = curr:visited t}
    sn <- get
    if curr == (target sn) || null neighbors  then return . fst $ _wrapUp curr sn else _process
    -- if curr == (target sn) then return . fst $ _wrapUp curr sn else _process

  -- Calculate the min distance and the corresponding previous.
  minAndPrev :: ODM -> Double -> [ElementaryMatrix] -> SEP -> SEP
  minAndPrev curr d unvisited t = foldr f t nds
    where nds = map (neighborsDists t curr d) unvisited
          f (md,k,e) s = s { dist = M.insert k md (dist s),
                             prev = M.insert k e (prev s) }

  neighborsDists :: SEP                           -- ^ Final state
                 -> ODM                           -- ^ Current
                 -> Double                        -- ^ Dist from origin
                 -> ElementaryMatrix              -- ^ Neighbor
                 -> (Double,HashMatrix,ElementaryMatrix)     -- ^ Min distance and Neighbor
  neighborsDists s m d e = (min newd oldd, k, e)
    where newd = d + elementaryValue (measure s) e
          oldd = M.lookupDefault newd k (dist s)
          k = HashMatrix $ applyElementary m e

  -- | Get next matrix to process and the shortest distance from src
  _getNext :: M.HashMap HashMatrix Double -- ^ List of distances by ODM
           -> (ODM,Double)                -- ^ ODM with min distance and the distance
  _getNext ds
    | null ds = error "Not destination found"
    | otherwise = (\(m,d) -> (hm m, d)) $ M.foldrWithKey f initial ds
    where initial = head . M.toList $ ds            -- (k,v)
          f k v mn@(_, minv) | v < minv = (k,v)
                             | otherwise = mn

  _wrapUp :: ODM -> SEP -> ([ElementaryMatrix],SEP)
  _wrapUp m s = (es, s)
    where p = M.lookup (HashMatrix m) (prev s)
          --op e = e {sign = - sign e}
          odm = (\x -> applyElementary m (opposite x)) <$> p
          es = maybe [] (\x -> fromJust p : fst (_wrapUp x s)) odm

  -- _testMatrix :: ODM
  -- _testMatrix = matrix 10 [0,0,0,1,0,0,0,0,0,0,0,2,1,0,1,3,0,0,0,0,0,0,2,0,0,4,1,0,1,0,0,0,0,1,0,0,0,0,2,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]

  -- _testSimMatrix :: ODM
  -- _testSimMatrix = matrix 10 [0,0,1,0,0,0,0,0,0,0,0,2,0,1,1,3,0,0,0,0,0,0,2,0,0,4,1,0,1,0,0,0,0,1,0,0,0,0,2,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]

  -- _testDifMatrix :: ODM
  -- _testDifMatrix = matrix 10 [0,0,0,1,0,0,0,0,0,0,0,2,1,0,1,3,0,0,0,0,0,0,2,0,0,4,2,0,0,0,0,0,0,1,0,0,0,0,2,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]


{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE FlexibleContexts #-}

module Matrix.Traversable
  (
    matrixTraverse
  ) where

import Numeric.LinearAlgebra.HMatrix hiding (corr)

import Data.Traversable

import Util.Tuples

-- Matrix cannot be an instance of Functor, so it cannot be an instance of traversable
-- instead, convert to list and use the list traverse
matrixTraverse :: (Element b, Element c) => (a -> b -> (a,c)) -> a -> Matrix b -> (a, Matrix c)
matrixTraverse f a m = applySnd (fromLists) $ mapAccumL (mapAccumL f) a (toLists m)