text:
{-# LANGUAGE GADTs, NoImplicitPrelude, TypeFamilies,
InstanceSigs, StandaloneDeriving, UndecidableInstances,
ScopedTypeVariables, FlexibleInstances, DataKinds,
FunctionalDependencies, PolyKinds,
TypeOperators, RankNTypes, ImpredicativeTypes, MultiParamTypeClasses,
AllowAmbiguousTypes, TypeApplications, FlexibleContexts #-}
---
module Main where

import qualified Data.Map.Strict as Map
import Prelude hiding ((.), id)
import Linear.Vector
import Data.Tuple
import Data.Complex
import Linear.Epsilon (nearZero)
import Control.Category
import Control.Monad ((<=<))
import GHC.TypeNats
import Data.Proxy
import Diag (mymain)
-- import qualified Text.PrettyPrint.ANSI.Leijen as PP
import qualified Text.PrettyPrint.Boxes as B
-- import Control.Arrow
{-
type MapVec r k = Map.Map k r

vadd :: (Ord k, Num r) => MapVec r k -> MapVec r k -> MapVec r k
vadd = Map.unionWith (+)

smul :: (Ord k, Num r) => k -> MapVec r k -> MapVec r k
smul

-- kmett alread did this stuff
-}

-- http://blog.sigfpe.com/2007/03/monads-vector-spaces-and-quantum.html

data W b a = W { runW :: [(a,b)] } deriving (Eq,Show,Ord)
instance Semigroup (W b a) where
  (W x) <> (W y) = W (x <> y)
instance Monoid (W b a) where
  mempty = W mempty
mapW f (W l) = W $ map (\(a,b) -> (a,f b)) l
instance Functor (W b) where
    fmap f (W a) = W $ map (\(a,p) -> (f a,p)) a

instance Num b => Applicative (W b) where
  pure x = W [(x,1)]
  (W fs) <*> (W xs) = W [(f x, a * b) | (f, a) <- fs, (x, b) <- xs]
  -- the fs is a a vector of rows kind of.
    -- error "I don't feel like it"

instance Num b => Monad (W b) where
   return x = W [(x,1)]
   l >>= f = W $ concatMap (\(W d,p) -> map (\(x,q)->(x,p*q)) d) (runW $ fmap f l)

a .* b = mapW (a*) b

instance (Eq a,Show a,Num b) => Num (W b a) where
     W a + W b = W $ (a ++ b)
     a - b = a + (-1) .* b
     _ * _ = error "Num is annoying"
     abs _ = error "Num is annoying"
     signum _ = error "Num is annoying"
     fromInteger a = if a==0 then W [] else error "fromInteger can only take zero argument"

collect :: (Ord a, Num b) => W b a -> W b a
collect = W . Map.toList . Map.fromListWith (+) . runW

trimZero = W . filter (\(k,v) -> not $ nearZero v) . runW
simplify :: Ord a => Q a -> Q a
simplify = trimZero . collect
-- filter (not . nearZero . snd)

type P a = W Double a

type Q a = W (Complex Double) a

v1 :: Map.Map Bool Double
v1 = zero

boolbase :: [Map.Map Bool Double]
boolbase = basis

e1 = Map.singleton True 1.0
e2 = Map.singleton False 1.0

star = mapW conjugate

{-
There are two types of particles

-}
data FibAnyon = Id | Tau
-- I started using DataKinds and it ended up being a problem.

data A a

data Tau
data Id
data FibTree root leaves where
   TTT :: FibTree Tau l -> FibTree Tau r -> FibTree Tau (l,r)
   ITT :: FibTree Tau l -> FibTree Tau r -> FibTree Id (l,r) -- we would like to enforce that isingany can split into a b, but thisis tough
   TIT :: FibTree Id l -> FibTree Tau r -> FibTree Tau (l,r)
   TTI :: FibTree Tau l -> FibTree Id r -> FibTree Tau (l,r)
   III :: FibTree Id l -> FibTree Id r -> FibTree Id (l,r)
   -- III :: FibTree 'Id l -> FibTree 'Id l' -> FibTree 'Id (l,l') -- maybe shouldn't exist. uneccessary. Does really help for monoidal instance
   TLeaf :: FibTree Tau Tau
   ILeaf :: FibTree Id Id

-- pretty printing would be hella nice
--deriving instance Show (FibTree a b)

instance Show (FibTree a b) where
  show = drawTree
deriving instance Eq (FibTree a b)
--deriving instance Ord (FibTree a b)
{-
r = PP.text "\\"
l = PP.text "/"
ls = PP.group $ (PP.text "  /" <> PP.line <> PP.text " /" <> PP.line <>PP.text "/")
rs = PP.group $ (PP.text "\\" <> PP.line <> PP.text " \\" <> PP.line <>PP.text "  \\")
-}

sp = B.char ' '
l = B.char '/'
r = B.char '\\'
h = B.char '-'
v = B.char '|'
t = B.char 'T'
i = B.char 'I'
rs x =  (r  B.<> sp B.<> sp) B.//
      (sp B.<> r  B.<> x) B.//
      (sp B.<> sp  B.<> r)
ls x =  (sp  B.<> sp B.<> l) B.//
      (x B.<> l  B.<> sp) B.//
      (l B.<> sp  B.<> sp)

{-
drawty = v B.// v B.// v B.// (h B.<> h B.<> h B.<> h)
drawb :: FibTree a b -> B.Box
drawb ILeaf = i
drawb TLeaf = t
drawb (TTT l r) = (drawb l) B.vcat   B.// ((drawb r)  B.<> v )
-}
-- PP.nest 2 l <> PP.line <> PP.nest 1 l <> PP.line <> l
-- text "Tau \\"
-- text "Id \\"

-- From Data.Tree

draw :: FibTree a b -> [String]
draw (ITT l r) = "I" : drawSubTrees (l,r)
draw (TIT l r) = "T" : drawSubTrees (l,r)
draw (TTT l r) = "T" : drawSubTrees (l,r)
draw (TTI l r) = "T" : drawSubTrees (l,r)
draw (III l r) = "I" : drawSubTrees (l,r)
draw  ILeaf    = "I" : []
draw  TLeaf    = "T" : []

-- drawSubTrees' [] = []
drawSubTrees' t =
    "|" : shift "`- " "   " (draw t)
drawSubTrees (t,ts) =
    "|" : shift "+- " "|  " (draw t) ++ drawSubTrees' ts

shift first other = zipWith (++) (first : repeat other)

{-
draw' ILeaf = ["I"]
draw' TLeaf = ["T"]
draw' (TTT l r) =  " -" ++ ((fmap . fmap) (\s -> s ++ "  |") (draw' l)) ++ (["-+"] ++ draw' r))
-}
-- drawex = drawb (TTT (TTT TLeaf TLeaf) TLeaf)

-- | Neat 2-dimensional drawing of a tree.
drawTree :: FibTree a b -> String
drawTree  = unlines . draw

-- | Neat 2-dimensional drawing of a forest.
-- drawForest :: Forest String -> String
-- drawForest  = unlines . map drawTree

instance Ord (FibTree a b) where
  compare (ITT l r) (ITT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (ITT _ _) _ = LT
  compare _ (ITT _ _) = GT
  compare (TTI l r) (TTI l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (TIT l r) (TIT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (TTT l r) (TTT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (III l r) (III l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'

  compare (TTI _ _) _ = LT
  compare _ (TTI _ _) = GT
  compare (TIT _ _) _ = LT
  compare _ (TIT _ _) = GT
  compare (TTT _ _) _ = LT
  compare _ (TTT _ _) = GT
  compare (III _ _) _ = LT
  compare _ (III _ _) = GT
  compare TLeaf TLeaf = EQ
  compare ILeaf ILeaf = EQ
  -- whoa. GHC is smart enough to realize some of these patterns can't happen
  -- because of the types.

{-
instance Enum (FibTree 'Tau (A 'Tau)) where
  toEnum _ = TLeaf
  fromEnum _ = 0
-}
{-
instance Enum (FibTree 'Id 'Id) where
  toEnum _ = TLeaf
  fromEnum _ = 0
  -}
-- and no inhabitants types?
--instance Enum (FibTree d b), Enum (FibTree e c) => Enum (FibTree a (b,c))
{-
instance (Enum (FibTree 'Tau b),
  Enum (FibTree 'Tau c),
  Enum (FibTree ' b),
  Enum (FibTree 'Id c)) => Enum (FibTree 'Tau (b,c))
  toEnum 0 = TTI
  toEnum
  fromEnum
-}

--instance Enum (FibTree a (b,c)) Enum (FibTree a (b,c))=> Enum (FibTree 'Id (b,c))

-- (FibTree a (b,c))

type TreeFun a b c d = FibTree a b -> FibTree c d
{-
lmap :: (forall a. FibTree a b -> FibTree (a) c) -> (forall e. FibTree e (b,d) -> FibTree e (c,d)) --  forall (a :: FibAnyon).
lmap f (ITT l r) = (ITT (f l) r)
lmap f (TTI l r) = (TTI (f l) r)
lmap f (TIT l r) = (TIT (f l) r)
lmap f (TTT l r) = (TTT (f l) r)
-}
lmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (b,d) -> Q (FibTree e (c,d)))
lmap f (ITT l r) = fmap (\l' -> ITT l' r) (f l)
lmap f (TTI l r) = fmap (\l' -> TTI l' r) (f l)
lmap f (TIT l r) = fmap (\l' -> TIT l' r) (f l)
lmap f (TTT l r) = fmap (\l' -> TTT l' r) (f l)
lmap f (III l r) = fmap (\l' -> III l' r) (f l)
{-
rmap :: (forall a. FibTree a b -> FibTree a c) -> (forall e. FibTree e (d,b) -> FibTree e (d,c)) --  forall (a :: FibAnyon).
rmap f (ITT l r) = ITT l (f r)
rmap f (TTI l r) = TTI l (f r)
rmap f (TIT l r) = TIT l (f r)
rmap f (TTT l r) = TTT l (f r)
-}
rmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (d,b) -> Q (FibTree e (d,c)))
rmap f (ITT l r) = fmap (\r' -> ITT l r') (f r)
rmap f (TTI l r) = fmap (\r' -> TTI l r') (f r)
rmap f (TIT l r) = fmap (\r' -> TIT l r') (f r)
rmap f (TTT l r) = fmap (\r' -> TTT l r') (f r)
rmap f (III l r) = fmap (\r' -> III l r') (f r)

fibswap :: FibTree a (l,l') -> FibTree a (l',l)
fibswap (ITT l r) = (ITT r l)
fibswap (TTI l r) = (TIT r l)
fibswap (TIT l r) = (TTI r l)
fibswap (TTT l r) = (TTT r l)

eye = 0 :+ 1

braid :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid (ITT l r) = W [(ITT r l,  cis $ 4 * pi / 5)]  -- different scalar factors for trivial and non trivial fusion
braid (TTT l r) = W [(TTT r l,  (cis $ - 3 * pi / 5))]
braid (TTI l r) = pure $ TIT r l-- exchange with trivial means nothing
braid (TIT l r) = pure $ TTI r l
braid (III l r) = pure $ III r l

-- The inverse of braid
braid' :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid' = star . braid

-- property
-- braid (braid' v) == v
-- braid' (braid v) == v

{-
fibassoc :: FibTree a ((c,d),e) -> FibTree a (c,(d,e))
fibassoc (ITT l r) = (ITT r l)
fibassoc (TTI l r) = (TIT r l)
fibassoc (TIT l r) = (TTI r l)
fibassoc (TTT l r) = (TTT r l)
-}
-- Looks like we need more constructors for the tree
{-
fmove :: FibTree a (c,(d,e)) -> FibTree a ((c,d),e)
-- fmove (ITT  a  (TTI b c)) = ITI ( TTT  a b) c
fmove (ITT  a  (TIT b c)) = ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = ITT ( TTT  a b) c

fmove (TTT  a  (TTI b c)) = TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = TTT ( TTI  a b) c
fmove (TTT  a  (TTT b c)) = TTT ( TTT  a b) c

fmove (TIT  a  (TTI b c)) = TTI ( TIT  a b) c
-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
fmove (TIT  a  (TTT b c)) = TTT ( TIT  a b) c
-}

fmove :: FibTree a (c,(d,e)) -> Q (FibTree a ((c,d),e))
-- fmove (ITT  a  (TTI b c)) = pure $ ITI ( TTT  a b) c -- pure (auto (auto a b) c) -- no maybe not. The internal one isn't auto
fmove (ITT  a  (TIT b c)) = pure $ ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = pure $ ITT ( TTT  a b) c
fmove (ITT  a  (TTI b c)) = pure $ III ( ITT  a b) c

fmove (TIT  a  (TTT b c)) = pure $ TTT ( TIT  a b) c
fmove (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
fmove (TIT  a  (TIT b c)) = pure $ TIT ( III  a b) c

-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root
-- internal I
fmove (TTI  a  (III b c)) = pure $ TTI ( TTI  a b) c
fmove (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, tau)         , (TTT ( TTT  a b) c, sqrt tau)]
-- internal T
fmove (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, sqrt tau)  ,   (TTT ( TTT  a b) c, - tau   )]
fmove (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c

fmove (III  a  (ITT b c)) = pure $ ITT ( TIT  a b) c
fmove (III  a  (III b c)) = pure $ III ( III  a b) c

-- largely just a tranpose of the above case.
fmove' :: FibTree a ((c,d),e) -> Q (FibTree a (c,(d,e)))
fmove' (ITT ( TTI  a b) c) = pure $ (ITT  a  (TIT b c))
fmove' (ITT ( TTT  a b) c) = pure $  (ITT  a  (TTT b c))
fmove' (ITT ( TIT  a b) c) = pure $  (III  a  (ITT b c))

--fmoveq (ITT  a  (TTT b c)) = pure $

fmove' (TTI ( TTT  a b) c) = pure $ (TTT  a  (TTI b c))
fmove' (TTI ( TTI  a b) c) = pure $ (TTI  a  (III b c))
fmove' (TTI ( TIT  a b) c) = pure $ TIT  a  (TTI b c)
--fmoveq (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c

--fmoveq (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c
fmove' (TIT ( ITT  a b) c) = W [(TTI  a  (ITT b c), tau)         , (TTT  a  (TTT b c) , sqrt tau)]
fmove' (TIT ( III  a b) c ) = pure $ TIT  a  (TIT b c)

fmove' (TTT ( TTI  a b) c ) = pure $ TTT  a  (TIT b c)
fmove' (TTT ( TIT  a b) c ) = pure $ TIT  a  (TTT b c)
fmove' (TTT ( TTT  a b) c) = W [(TTI  a  (ITT b c), sqrt tau)  , (TTT  a  (TTT b c),   - tau  )]

fmove' (III ( III  a b) c ) = pure $ III  a  (III b c)
fmove' (III ( ITT  a b) c ) = pure $ ITT  a  (TTI b c)

--fmoveq (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root

-- internal I

--fmoveq (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, recip tau)         , (TTT ( TTT  a b) c, recip $ sqrt tau)]
-- internal T

--fmoveq (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, recip $ sqrt tau)  , (TTT ( TTT  a b) c,   - recip tau  )]

checkf :: Q (FibTree Tau ((Tau, Tau),Tau))
checkf = simplify $ fmove' (TIT (ITT TLeaf TLeaf) TLeaf) >>= fmove
checkf' = simplify $ fmove' (TTT (TTT TLeaf TLeaf) TLeaf) >>= fmove
-- checkf''' = fmove' (TIT (ITT TLeaf TLeaf) TLeaf)
-- checkf'' = fmove' (ITT (TTT TLeaf TLeaf) TLeaf) >>= fmove

-- tau**2 + tau == 1
tau :: Complex Double
tau =  ((sqrt 5) - 1) / 2 :+ 0 -- 0.618 :+ 0

--test1 = lmap fibswap
-- test2 = lmap $ lmap fibswap
-- test3 = rmap fmove

type Vec1 b r = [(b, r)]

smul :: Num r => r -> Vec1 b r -> Vec1 b r
smul s = map (\(b,s') -> (b, s' * s))

linapply :: Num r => (a -> Vec1 b r) -> (Vec1 a r -> Vec1 b r) -- bind
linapply f [] = []
linapply f ((b,s) : xs) = smul s (f b) ++ linapply f xs

dTTT :: FibTree a (Tau, Tau) -> Q (FibTree a Tau)
dTTT (TTT _ _) = pure TLeaf
dTTT _ = mempty
{-
dITT :: FibTree a (Tau, Tau) -> Q (FibTree a Id)
dTTT (TTT _ _) = pure TLeaf
dTTT _ = mempty
-}
dTTI :: FibTree a (Tau, Id) -> Q (FibTree a Tau) -- A dual tree of type  'tau (tau,'tau)
dTTI (TTI _ _) =  pure TLeaf
dTTI _ = mempty

dTTT'''' = dot (TTT TLeaf TLeaf)
-- I could do this as a typeclass. NOPE. don't need that garbage

-- I think that incompatible FibTree a' a is empty by default vs unconstructible
-- which is better?
dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
dot x@(TTI _ _) y@(TTI _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TIT _ _) y@(TIT _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TTT _ _) y@(TTT _ _) | x == y = pure TLeaf
                            | otherwise = mempty
dot x@(III _ _) y@(III _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot x@(ITT _ _) y@(ITT _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot _ _ = mempty

test5 = (pure $ TTT TLeaf TLeaf) >>= dTTT
test6 = (pure $ ITT TLeaf TLeaf) >>= dTTT
test7 v = do y <- rmap dTTI v  -- wait. What the hell am I doing?
             lmap dTTT y
test8 v = rmap dTTT v  >>= dTTT
-- Building from the leaves up...?
-- you  hve to apply the deeper stuff before the upper stuff.
--- So this is a dual tree with TTT in the right branch. Then TTT above that.
-- In total (T,(T,T)) -> Tleaf.
-- maybe we need to cps it.
dTTT' :: (Q (FibTree a Tau) -> (Complex Double)) -> (FibTree a (Tau, Tau)) -> Complex Double
dTTT' f = f . dTTT

dTTT'' :: (FibTree a Tau -> Complex Double) -> FibTree a (Tau, Tau) -> Complex Double
dTTT'' f = undefined

-- sort of similar to bind in a way. lift function to a Q function
-- (Q FibTree a ) -> (FibTree a -> Q FibTree b)  -> Q FibTree b
dot' :: (FibTree a b -> Complex Double) -> Q (FibTree a b) -> Complex Double
dot' f (W v) = sum $ map (\(e,a) -> a * (f e)) v

dTTT''' = dot' . dTTT''

pentagon1 ::  FibTree a (e,(d,(c,b))) -> Q (FibTree a (((e,d),c),b))
pentagon1 v =  do
                 v1 <- fmove v
                 fmove v1

-- type annotations not necessary. For clarity.
pentagon2 :: FibTree a (b,(c,(d,e))) -> Q (FibTree a (((b,c),d),e))
pentagon2 v = do
                v1 :: FibTree a (b,((c,d),e)) <- rmap fmove v
                v2 :: FibTree a ((b,(c,d)),e) <- fmove v1
                lmap fmove v2

ex3 = TTT TLeaf (TTI TLeaf (ITT TLeaf TLeaf))
ex1 = TTT TLeaf (TTT TLeaf (TTT TLeaf TLeaf))
pentagon =  simplify $ ((pentagon1 ex1) - (pentagon2 ex1))
pentagon' =  simplify $ ((pentagon1 ex3) - (pentagon2 ex3))

hexagon1 :: FibTree a (b,(c,d)) -> Q (FibTree a ((d,b),c))
hexagon1 v = do
             v1 :: FibTree a ((b,c),d) <- fmove v
             v2 :: FibTree a (d,(b,c)) <- braid v1
             fmove v2

hexagon2 :: FibTree a (b,(c,d)) -> Q (FibTree a ((d,b),c))
hexagon2 v = do
             v1 :: FibTree a (b,(d,c)) <- rmap braid v
             v2 :: FibTree a ((b,d),c) <- fmove v1
             lmap braid v2

ex4 = (TTI TLeaf (ITT TLeaf TLeaf))
ex2 = (TTT TLeaf (TTT TLeaf TLeaf))
hexagon =  simplify $ ((hexagon1 ex2) - (hexagon2 ex2))
hexagon' =  simplify $ ((hexagon1 ex4) - (hexagon2 ex4))

-- hexagon2 :: FibTree a (b,(c,d)) -> Q (FibTree a (c,(d,b)))

-- quickcheck forall v, pentagon1 v == pentagon2 v

-- test9 = test8 $ pure $ TTT (TLeaf) ()
-- Vec1 (FibTree 'Tau l) r -> Vec1 (FibTree 'Tau l') r -> Vec1 (FibTree 'Id (l,l')) r

-- linearized tensor products
--linITT :: Num r => Vec1 (FibTree 'Tau l) r -> Vec1 (FibTree 'Tau l') r -> Vec1 (FibTree 'Id (l,l')) r -- we would like to enforce that isingany can split into a b, but thisis tough
--linITT = [ (ITT l r, s * s')  |  (l, s) <- v1 , (r, s') <- v2   ]

-- class AutoNode c a b where
      -- prod ::  FibTree a l -> FibTree b l' -> FibTree  c (l,l')
{-
 instance AutoProd 'Tau 'Tau 'Tau
    prod = TTT
 instance AutoProd 'Id 'Tau 'Tau
    prod = ITT
-- etc.

-- I can't make a fst. It extracts a skolemized object.
-- So I do need to do complete pattern matching? No. I CAN do this... ? Not always clear I'd want to?
fibfst :: AutoNode a e _ => FibTree a (b,c) -> FibTree e b

-}
class TProd a b c where
  tprod :: Vec1 (FibTree a l) r -> Vec1 (FibTree b l') r -> Vec1 (FibTree c (l,l')) r

-- tprod :: AutoNode a b c => Vec1 (FibTree a l) r -> Vec1 (FibTree b l') r -> Vec1 (FibTree c (l,l')) r
--class Index tree n  -- reference leaf by number?

-- class AutoFMove tree n

{-
linTTI :: FibTree 'Tau l -> FibTree 'Id l' -> FibTree 'Tau (l,l')
linTIT :: FibTree 'Id l -> FibTree 'Tau l' -> FibTree 'Tau (l,l')
libTTT :: FibTree 'Tau l -> FibTree 'Tau l' -> FibTree 'Tau (l,l')
-}
(~*) :: r -> b -> (b,r) -- **, `smul`
(~*) = flip (,)

(~+) :: a -> [a] -> [a] -- ++ instead?
(~+) = (:)

test4 = [1 ~* 'a', 2 ~* 'b'] -- eh. What's the point.
-- test3 = (lmap . lmap) fibswap
--lmap (ITT l r)
-- ^+^
-- zero
--

-- TODOS

-- Categorical Interface
-- lmap, rmap vs second first.

{-
pullLeft (Tau,x) = (Tau,x)
pullLeft (Id,x) = (Id,x)
pullLeft (a,b) = fmove ((a',c),b) where (a',c) = pullLeft a
-- ((a,b),c) = (, (b,c)) where (a',d) = pullLeft a
-}
{-
type family FMove' a where
  FMove' ((a,b),c)  = (a,(b,c))

type family PullLeft a where
  PullLeft Tau = Tau
  PullLeft Id = Id
  PullLeft (Tau,b) = (Tau,b)
  PullLeft (Id,b) =  (Id,b)
  PullLeft ((a,b),c)  = FMove' (PullLeft (a,b), c)

-- Auto F Moves. / auto braid.

pullLeft :: (PullLeft (b,c) ~ (b',c')) => FibTree a (b,c) -> Q (FibTree a (b',c'))
pullLeft x@(TTT TLeaf _) = pure x
pullLeft (TTT l@(TTT _ _ ) r) =  do
                      l' <- pullLeft l
                      fmove' (TTT l' r)
-}
class PullLeft a b | a -> b where -- | a -> b functional dependency causes errors?
  pullLeft :: FibTree c a -> Q (FibTree c b)

instance PullLeft (Tau,c) (Tau,c) where
  pullLeft = pure

instance PullLeft (Id,c) (Id,c) where
  pullLeft = pure

instance PullLeft Tau Tau where
  pullLeft = pure

instance PullLeft Id Id where
  pullLeft = pure

instance (PullLeft (a,b) (a',b'), r ~ (a',(b',c))) => PullLeft ((a, b),c) r where
  pullLeft t = do
           t' <- lmap pullLeft t
           fmove' t'
{-
instance (PullLeft a a', r ~ (a',(b,c))) => PullLeft (a, (b,c)) r where
  pullLeft t = lmap pullLeft t
-}
{-
instance (PullLeft a a', PullLeft (a',b) (a'',b''), ) => PullLeft (a,b) r where
  pullLeft = lmap pullLeft t
-}

{-  pullLeft (TTT l r) =  do
                        l' <- pullLeft l
                        fmove' (TTT l' r)
-}

class PullRight a b | a -> b where -- | a -> b functional dependency causes errors?
  pullRight :: FibTree c a -> Q (FibTree c b)
{-
instance (PullRight a a', r ~ ((b,c),a')) => PullRight ((b,c),a) r where
  pullRight t = rmap pullRight t
-}
instance PullRight Tau Tau where
  pullRight = pure

instance PullRight Id Id where
  pullRight = pure

instance PullRight (c,Tau) (c,Tau) where
  pullRight = pure

instance PullRight (c,Id) (c,Id) where
  pullRight = pure

instance (PullRight (a,b) (a',b'), r ~ ((c,a'),b')) => PullRight (c,(a, b)) r where
  pullRight t = do
           t' <- rmap pullRight t
           fmove t'

type family Count a where
  Count Tau = 1
  Count Id = 1
  Count (a,b) = (Count a) + (Count b)

type family LeftCount a where
  LeftCount (a,b) = Count a

-- The version without the explicit ordering supplied.
class LCA n a b c d | n a c -> b d where
  lcamap :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)
{-
lcamapP :: LCA n a b c d => Proxy n ->  (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)
lcamapP _ f x = lcamap f x
-}
instance (lc ~ (LeftCount a),
          gte ~ (CmpNat lc n),
         LCA' n gte a b c d) => LCA n a b c d where
  lcamap f x = lcamap' @n @gte f x

class LCA' n gte a b c d | n gte a c -> b d where
  lcamap' :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)

-- we find b at the lca and pass it back up. c gets passed all the way down, d gets computed by rebuilding out of c.
-- a drives the search.
instance (n' ~ (n - Count l), -- we're searching in the right subtree. Subtract the leaf number in the left subtree
        lc ~ (LeftCount r), -- dip one left down to order which way we have to go next
        gte ~ (CmpNat lc n'), -- Do we go left, right or havce we arrived in the next layer?
        LCA' n' gte r b c d',  -- recurive call
        d ~ (l,d') -- reconstruct total return type from recurive return type. left tree is unaffected by lcamapping
        ) => LCA' n 'LT (l,r) b c d where
    lcamap' f x = rmap (lcamap' @n' @gte f) x

instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LCA' n gte l b c d',
          d ~ (d',r)
          ) => LCA' n 'GT (l,r) b c d where
    lcamap' f x = lmap (lcamap' @n @gte f) x

instance (b ~ a, d ~ c) => LCA' n 'EQ a b c d where
  lcamap' f x = f x

class LeafMap n gte a b c d | n gte a c -> b d where
  leafmap :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)

instance (n' ~ (n - Count l), -- we're searching in the right subtree. Subtract the leaf number in the left subtree
        lc ~ (LeftCount r), -- dip one left down to order which way we have to go next
        gte ~ (CmpNat lc n'), -- Do we go left, right or havce we arrived in the next layer?
        LeafMap n' gte r b c d',  -- recurive call
        d ~ (l,d') -- reconstruct total return type from recurive return type. left tree is unaffected by lcamapping
        ) => LeafMap n 'LT (l,r) b c d where
    leafmap f x = rmap (leafmap @n' @gte f) x

instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LeafMap n gte l b c d',
          d ~ (d',r)
          ) => LeafMap n 'GT (l,r) b c d where
    leafmap f x = lmap (leafmap @n @gte f) x

-- In the equals case, we now continue onward.
instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LeafMap n gte l b c d',
          d ~ (d',r)
          ) => LeafMap n 'EQ (l,r) b c d where
    leafmap f x = lmap (leafmap @n @gte f) x

-- base case
instance (b ~ Tau, d ~ c) => LeafMap 1 'EQ Tau b c d where
  leafmap f x = f x

instance (b ~ Id, d ~ c) => LeafMap 1 'EQ Id b c d where
  leafmap f x = f x

-- split
-- leafmap @3 (const (TTT TLeaf TLeaf))
-- need one last arbitrary left or right fmove to put the two on the same stalk
-- neighbormap :: (LCA n a (l,r) (l',r') d, PullRight l (l',x), PullLeft r (y,r') ) => (FibTree.  -> Q FibTree) -> FibTree

-- I need this to also work for (Tau, Tau)
-- (Tau, yada yada)
-- I.e. I need to do different things depending on

-- Jesus. What a shitshow. But I don't see how to do better.
-- a is starting tree
-- b is extracted part
-- c is transformed b
-- d is recoustrcted a replacing b with c and rearranging.
class NeighborMap a b c d | a c -> b d where
  nmap :: (forall r. FibTree r b -> Q (FibTree r c)) -> FibTree e a -> Q (FibTree e d)

instance NeighborMap ((l,x),(y,r)) (x,y) c ((l,c),r) where
   nmap f x = do
              x'  <- fmove x -- (((l',x),y),r')
              x'' <- lmap fmove' x' -- ((l',(x,y)),r')
              lmap (rmap f) x''
instance NeighborMap (Tau, (y,r)) (Tau, y)  c  (c,r) where
   nmap f x = fmove x >>= lmap f
instance NeighborMap (Id, (y,r)) (Id, y) c (c,r) where
   nmap f x = fmove x >>= lmap f
instance NeighborMap ((l,x), Tau) (x,Tau) c (l,c) where
   nmap f x = fmove' x >>= rmap f
instance NeighborMap ((l,x), Id) (x,Id) c  (l,c) where
   nmap f x = fmove' x >>= rmap f
instance NeighborMap (Tau, Tau) (Tau,Tau) c  c where
   nmap f x = f x
instance NeighborMap (Id, Id) (Id,Id) c  c where
   nmap f x = f x
instance NeighborMap (Tau, Id) (Tau,Id) c  c where
   nmap f x = f x
instance NeighborMap (Id, Tau) (Id,Tau) c  c where
   nmap f x = f x

neighbormap :: forall n a b c d l l' r' x y e r z. (LCA n a b c d,
   b ~ (l,r),
   c ~ ((l',z),r'),
   PullRight l (l',x),
   PullLeft r (y,r')) =>
   (forall r. FibTree r (x,y) -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap f z = lcamap @n @a @b @c @d (helper f) z

helper ::  (c ~ ((l',z),r'),
   PullRight l (l',x),
   PullLeft r (y,r')) => (forall r. FibTree r (x,y) -> Q (FibTree r z)) -> FibTree e (l,r) -> Q (FibTree e c)
helper f x = do
            x' <- rootneighbor x
            lmap (rmap f) x'

rootneighbor :: (PullRight l (l',x), PullLeft r (y,r')) => FibTree e (l,r) -> Q (FibTree e ((l',(x,y)),r'))
rootneighbor x = do
                x' <- lmap pullRight x
                x'' <- rmap pullLeft x' -- ((l',x),(y,r'))
                x''' <- fmove x'' -- (((l',x),y),r')
                lmap fmove' x''' -- ((l',(x,y)),r')

neighbormap' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') b' z c) =>
   (forall r. FibTree r b' -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap' f z = lcamap @n @a @b @c @d (\x -> do
                                          x'  <- lmap pullRight x
                                          x'' <- rmap pullLeft x'
                                          nmap f x'') z

neighbormap'' :: forall n a b' z d c l l' r r' x y e b. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') b' z c) =>
   (forall r. FibTree r b' -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap'' f z = lcamap @n @a @b @c @d (\x -> do
                                          x'  <- lmap pullRight x
                                          x'' <- rmap pullLeft x'
                                          nmap f x'') z

t1 = neighbormap' @2 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t2 = neighbormap' @1 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t3 = neighbormap' @2 braid (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))

                                      {-
neighbormap p f z = let helper (x :: FibTree _ (l,r)) = do
                                x' <- rootneighbor x
                                lmap (rmap f) x'
                    in
                    lcamap helper z
                    -}
{-  lcamap
((\x -> do
                                x' <- rootneighbor x
                                lmap (rmap f) x'
                                ) :: (forall s. FibTree s (l,r) -> Q (FibTree s c))) z
-}

{-(\x -> do
                            x' <- lmap pullRight x
                            x'' <- rmap pullLeft x' -- ((l',x),(y,r'))
                            x''' <- fmove x'' -- (((l',x),y),r')
                            x4 <- lmap fmove' x''' -- ((l',(x,y)),r')
                            lmap (lmap f) x4 -- ((l',z),r')
                            ) z
-}
{-
-- neighbormap f z = lcamap @n (\x -> do
                            x' <- lmap pullRight x
                            x'' <- rmap pullLeft x'
                            x''' <- fmove x''
                            lmap? f x'''   ) z
-}
{-
abraid :: forall n a b c d l l' r' x y e r z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) => Q (FibTree e a) -> Q (FibTree e d)

   -}

abraid :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   Q (FibTree e a) -> Q (FibTree e d)
abraid x = x >>= (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid

abraid' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   Q (FibTree e a) -> Q (FibTree e d)
abraid' x = x >>= (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid'

{-
abraid
  :: forall n a b c d l l' r r' x y e z b'. (LCA n a (l, r) c d, PullRight l l', PullLeft r r',
      NeighborMap (l', r') (x, y) (y, x) c) =>
     (FibTree e a) -> Q (FibTree e d)
abraid x = neighbormap' @n @a @(l,r) @c @d @l @l' @r @r' braid x
-}
{-
t4 = abraid @2 $
     abraid @4 $
     abraid @3 $
     pure (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
-}
{-
ndot' :: forall n a b c d l l' r r' x y e z b' a'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) a' c) =>
   FibTree a' (x,y) -> FibTree e a -> Q (FibTree e d)
   -}
-- ndot' x = (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @a' @(x,y)) (dot' x)

ndot :: forall n a l r c d l' r' x y a' e. (LCA n a (l, r) c d, PullRight l l', PullLeft r r',
      NeighborMap (l', r') (x, y) a' c) =>
     FibTree a' (x, y) -> FibTree e a -> Q (FibTree e d)
ndot x = neighbormap' @n @a @(l,r) (dot x)

-- This a b c d pattern is something like a lens. We are replacing
-- piece of b with c which turns a into d. b is piece of a.
-- s t a b

-- quantified constraints
--  (forall b' c'. (g a b' c' d , f b' b c c') => Compose f g a b c d

-- (forall r. FibTree r b' -> Q (FibTree r z)) ->
{-
abraid' :: forall n a b c d l l' r r' x y e b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   FibTree e a -> Q (FibTree e d)
abraid' = (neighbormap'' @n @a @(x,y) @(y,x) @d) braid
-}

{-
abraid' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   FibTree e a -> Q (FibTree e d)
abraid' = (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid'
-}

{-
abraid :: (LCA n a (l1, r) c d, PullRight l1 l'1, PullLeft r r',
      NeighborMap (l'1, r') (l2, l'2) (l'2, l2) c) =>
     FibTree e a -> Q (FibTree e d)
abraid = neighbormap' braid
-}
-- autobraid = neighbormap braid
-- autodot x = neighbormap (dot x)

-- Build state monad to carry along the vector.
-- should be able to do in applicative style, since data does not change structre of computation

-- auto braid
-- find least common ancestor
-- pullLeft, and pullRight
-- then braid

-- stateful
-- Pretty Print tree

-- Operators.
newtype FibOp a b = FibOp (forall c. FibTree c a -> Q (FibTree c b))
type FibOp' c a b = FibTree c a -> Q (FibTree c b)

instance Category FibOp where
  id = FibOp pure
  (FibOp f) . (FibOp g) = FibOp (f <=< g)
{-
instance Arrow FibOp where
  arr = error "No. No arr."
  (***) =
-}
{-
class Monoidal k where
  (***) :: k a b -> k c d -> k (a,c) (b,d)
-}
tensor :: Num r => (b -> c -> a) -> W r b -> W r c -> W r a
tensor f (W xs) (W ys) =  W [ ( f x y , xc * yc ) | (x, xc) <- xs,  (y,yc) <- ys]

-- tensor :: Num r => (b -> c -> a) -> b -> c -> W r a
-- fg f g (TTT l r) = tensor TTT (f l) (g r)

{-
instance Monoidal FibOp where
  (***) :: FibOp a b -> FibOp c d -> FibOp (a,c) (b,d)
  (FibOp f) *** (FibOp g) = FibOp fg where
                               -- fg ::  (a,c) (b,d)
                               fg :: FibTree e (a,c) -> Q (FibTree e (b,d))
                               fg (TTT l r) = tensor TTT (f l) (g r)
                               -- fg (ITT l r) = FibOp $ tensor ITT (f l) (g r)
                               -- fg (TIT l r) = FibOp $ tensor TIT (f l) (g r)
-}
-- try to remove the issue of being in a instance. Which is kind of silly anyhow
{-
(****) :: FibOp' e a b -> FibOp' e c d -> FibOp' e (a,c) (b,d)
f **** g = fg where
                   -- fg ::  (a,c) (b,d)
                 fg :: FibTree e (a,c) -> Q (FibTree e (b,d))
                 fg (TTT l r) = tensor TTT (f l) (g r)
                 fg (ITT l r) = tensor ITT (f l) (g r)
-}
(***) :: (forall e. FibOp' e a b) -> (forall e'. FibOp' e' c d)  -> FibOp' e'' (a,c) (b,d)
(***) f g (TTT l r) = tensor TTT (f l) (g r)
(***) f g (ITT l r) = tensor ITT (f l) (g r)
(***) f g (TIT l r) = tensor TIT (f l) (g r)

first = lmap
second = rmap
-- braid
-- fmove

-- cup, cap? dot'

-- Is the problem that I'd trying to pattern match on kind of the output?
 -- densification

-- Quickcheck props

-- implement qubit abtraction

-- Did I fall into a way to make end, co-end work as einstein notation?6
-- The bounded meaning of forall and exists in the GADT context

main :: IO ()
main = mymain
--return ()


text:
{-# LANGUAGE DeriveGeneric #-}

{-|
Module      : ODMatrix.SmithDecomposition.SparseMat
Description : Sparse representation of a matrix.

-}
module ODMatrix.SmithDecomposition.SparseMat (
      SparseMat(..)
    , denseSparseMat
    , sparseDenseMat
    , fixIndexes
    , makeAssoc
  ) where

    import GHC.Generics (Generic)
    --import Data.Aeson (ToJSON, FromJSON)

    import Numeric.LinearAlgebra as L hiding (rows, cols)
    import Numeric.LinearAlgebra.Data as LD (toList)

    -- | Sparse representation of a matrix.
    data SparseMat = SparseMat {
        rows :: Int
      , cols :: Int
      , cells :: AssocMatrix
      } deriving (Generic, Show)

    --instance ToJSON SparseMat
    --instance FromJSON SparseMat

    -- | Transform a sparse matrix in a regular one.
    denseSparseMat :: SparseMat -> Matrix Double
    denseSparseMat m = assoc (rows m, cols m) 0 (fixIndexes $ cells m)

    -- | Transform a regular matrix to an sparse one.
    sparseDenseMat :: Matrix Double-> SparseMat
    sparseDenseMat m = SparseMat r c cs
      where (r, c) = size m
            vals = toList . flatten $ m
            ids = [(i,j) | i <- [0..r-1], j <- [0..c-1]]
            cs = filter (\(_,v) -> v /= 0) $ zip ids vals

    -- * Support functions

    -- | Adjust the indexes of the association matrix from 1-starting to 0-starting.
    -- This is needed by the Linear Algebra library.
    fixIndexes :: (Num a) => [((a,a),b)] -> [((a,a),b)]
    fixIndexes = map (\((a,b),v) -> ((a-1,b-1),v))

    -- | Given a constant value and a list of indexes, construct a association matrix only containing that value.
    makeAssoc :: Double -> [(Int,Int)] -> AssocMatrix
    makeAssoc v = fixIndexes . map (flip (,) v)

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

ext: hs

max_stars_repo_name: renecura/odmatrix

max_issues_repo_name: renecura/odmatrix

max_forks_repo_name: renecura/odmatrix


text:
{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses   #-}
{-# LANGUAGE ImpredicativeTypes, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Count-Min Sketch
-- | Creator: Xiao Ling
-- | Created: 12/17/2015
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module CountMinSketch where

import Control.Monad.Random.Class
import Control.Monad.Random
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Identity

import Data.Matrix
import Data.Conduit
import qualified Data.Conduit.List as Cl
import qualified Data.Matrix as M

import Core
import Utils
import Statistics

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

type W        = Int              -- * num cols
type D        = Int              -- * num rows
type P        = Int              -- * some prime   --> change this to Integer
type A        = Int              -- * coefficient a_j
type B        = Int              -- * coeeficient b_j
type Event    = Int              -- * some input event from
                                 -- * universe: {1, .., i, .. n}

type Sketch   = Matrix Int       -- * W x D matrix  of counters
type ABs      = Matrix Int       -- * 2 x D matrix  of a_j and b_j for j = 1..d
type Events   = Matrix Event     -- * 2 x D matrix of event i
type Hash     = Event -> ColIdx  -- * Maps event to d x 1 list of index to incr
type ColIdx   = [Int]            -- * indices of counters to increment

type Some m   = (MonadRandom m, MonadState (Sketch, Hash) m, MonadReader (D,W,P) m)

{-----------------------------------------------------------------------------
    Count Min Sketch
------------------------------------------------------------------------------}

{-----------------------------------------------------------------------------
    Helpers
------------------------------------------------------------------------------}

-- * initalize the sketch
sketch' :: D -> W -> Sketch
sketch' = zeros

-- * intialize random cofficients a_j, b_j, for d x 2 matrix:
-- * [ a_1 ... a_j ... a_d ]
-- * [ b_1 ... b_j ... b_d ]
ab' :: MonadRandom m => P -> D -> m ABs
ab' p d = (M.fromList 2 d . take (2*d)) <$> getRandomRs (1,p)

-- * h_j (i) = (a_j x i + b_j mod p) mod w
-- * sumRow-wise  [ a_1 ... a_j ... a_d ]  .*.  [ i_1 ...  i_d ]     mod p  mod w
-- *              [ b_1 ... b_j ... b_d ]       [ 1   ...  1   ]
-- * result:      [ idx_1, ..., idx_d   ]
-- *               where the indices start at 1
hash' :: P -> W -> D -> ABs -> Hash
hash' p w d ab = \i ->  go $ sumR (ab .*. is) ... (\v -> v `mod` p `mod` w)
    where go = fmap (+1) . toList
          is = is' d i

-- * map `i` <- event univese onto d x 2 matrix of form
-- * [ i_1 ... i_2 ]
-- * [ 1   ...  1  ]
is' :: D -> Event -> Events
is' d i = (i .+ (zeros 1 d)) <-> ones 1 d

-- * given hash function, construct a binary mask of size `w` x `d`
-- * to increment sketch
mask' :: W -> ColIdx -> Sketch
mask' w ks = let vec c w = zeros 1 (c-1) <|> single 1 <|> zeros 1 (w - c) in
    case ks of
        c:[] -> vec c w
        c:cs -> vec c w <-> mask' w cs

-- * given input i and some hash function `h`, and sketch `ms`,
-- * update `ms`
update' :: W -> D -> Hash -> Sketch -> Event -> Sketch
update' w d h ms i = mask' w (h i) .+. ms

maxi  = 30            :: Event
p     = 2^maxi - 1    :: P
(d,w) = (4,3)         :: (D,W)
i     = 12            :: Event

mab    = evalRandIO $ ab' p d
sketch = sketch' d w
mhashs = (hash' p w d)  <$> mab

update h i = update' w d h sketch i

mhash  = head <$> mhashs

--hs = hashF p w xs ys


text:
{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeInType           #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

module TensorOps.Learn.NeuralNet.FeedForward
  ( Network(..)
  , buildNet
  , runNetwork
  , trainNetwork
  , induceNetwork
  , nmap
  , (~*)
  , (*~)
  , liftNet
  , netParams
  , networkGradient
  , genNet
  , ffLayer
  ) where

import           Control.Category
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Kind
import           Data.Singletons
import           Data.Singletons.Prelude        (Sing(..))
import           Data.Type.Conjunction
import           Data.Type.Length
import           Data.Type.Product              as TCP
import           Data.Type.Product.Util         as TCP
import           Data.Type.Sing
import           Prelude hiding                 ((.), id)
import           Statistics.Distribution.Normal
import           System.Random.MWC
import           TensorOps.Learn.NeuralNet
import           TensorOps.NatKind
import           TensorOps.TOp                  as TO
import           TensorOps.Tensor               as TT
import           TensorOps.Types
import           Type.Class.Higher
import           Type.Class.Higher.Util
import           Type.Class.Known
import           Type.Class.Witness
import           Type.Family.List
import           Type.Family.List.Util

data Network :: ([k] -> Type) -> k -> k -> Type where
    N :: { _nsPs    :: !(Sing ps)
         , _nOp     :: !(TOp ('[i] ': ps) '[ '[o] ])
         , _nParams :: !(Prod t ps)
         } -> Network t i o

instance NFData1 t => NFData (Network t i o) where
    rnf = \case
      N _ o p -> o `seq` p `deepseq1` ()
    {-# INLINE rnf #-}

buildNet
    :: SingI ps
    => TOp ('[i] ': ps) '[ '[o] ]
    -> Prod t ps
    -> Network t i o
buildNet = N sing

netParams
    :: Network t i o
    -> (forall ps. SingI ps => Prod t ps -> r)
    -> r
netParams n f = case n of
    N o _ p -> f p \\ o

(~*~)
    :: Network t a b
    -> Network t b c
    -> Network t a c
N sPs1 o1 p1 ~*~ N sPs2 o2 p2 =
    N (sPs1 %:++ sPs2) (o1 *>> o2) (p1 `TCP.append'` p2)
        \\ singLength sPs1
infixr 4 ~*~
{-# INLINE (~*~) #-}

instance Category (Network t) where
    id = N SNil idOp Ø
    (.) = flip (~*~)

(~*) :: TOp '[ '[a] ] '[ '[b] ]
     -> Network t b c
     -> Network t a c
f ~* N sO o p = N sO (f *>> o) p
infixr 4 ~*
{-# INLINE (~*) #-}

(*~) :: Network t a b
     -> TOp '[ '[b] ] '[ '[c] ]
     -> Network t a c
N sO o p *~ f = N sO (o >>> f) p
infixl 5 *~
{-# INLINE (*~) #-}

liftNet
     :: TOp '[ '[i] ] '[ '[o] ]
     -> Network t i o
liftNet o = buildNet o Ø

nmap
     :: SingI o
     => (forall a. RealFloat a => a -> a)
     -> Network t i o
     -> Network t i o
nmap f n = n *~ TO.map f
{-# INLINE nmap #-}

runNetwork
    :: (RealFloat (ElemT t), Tensor t)
    => Network t i o
    -> t '[i]
    -> t '[o]
runNetwork (N _ o p) = head' . runTOp o . (:< p)
{-# INLINE runNetwork #-}

trainNetwork
    :: forall i o t. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> ElemT t
    -> t '[i]
    -> t '[o]
    -> Network t i o
    -> Network t i o
trainNetwork loss r x y = \case
    N s o p ->
      let p' = map1 (\(!(s1 :&: o1 :&: g1)) -> TT.zip stepFunc o1 g1 \\ s1)
             $ zipProd3 (singProd s) p (tail' $ netGrad loss x y s o p)
      in  N s o p'
  where
    stepFunc :: ElemT t -> ElemT t -> ElemT t
    stepFunc !o' !g' = o' - r * g'
    {-# INLINE stepFunc #-}
{-# INLINE trainNetwork #-}

induceNetwork
    :: forall i o t. (Tensor t, RealFloat (ElemT t), SingI i)
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> ElemT t
    -> t '[o]
    -> Network t i o
    -> t '[i]
    -> t '[i]
induceNetwork loss r y = \case
    N s o p -> \x -> TT.zip stepFunc x (head' $ netGrad loss x y s o p)
  where
    stepFunc :: ElemT t -> ElemT t -> ElemT t
    stepFunc o' g' = o' - r * g'
    {-# INLINE stepFunc #-}
{-# INLINE induceNetwork #-}

networkGradient
    :: forall i o t r. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> t '[i]
    -> t '[o]
    -> Network t i o
    -> (forall ps. SingI ps => Prod t ps -> r)
    -> r
networkGradient loss x y = \case
    N s o p -> \f -> f (tail' $ netGrad loss x y s o p) \\ s
{-# INLINE networkGradient #-}

netGrad
    :: forall i o ps t. (Tensor t, RealFloat (ElemT t))
    => TOp '[ '[o], '[o] ] '[ '[] ]
    -> t '[i]
    -> t '[o]
    -> Sing ps
    -> TOp ('[i] ': ps) '[ '[o] ]
    -> Prod t ps
    -> Prod t ('[i] ': ps)
netGrad loss x y s o p = (\\ appendSnoc lO (Proxy @'[o])) $
                         (\\ lO                         ) $
                         takeProd @'[ '[o] ] (LS lO)
                       $ gradTOp o' inp
  where
    lO  :: Length ps
    lO = singLength s
    o'  :: ((ps ++ '[ '[o] ]) ~ (ps >: '[o]), Known Length ps)
        => TOp ('[i] ': ps >: '[o]) '[ '[]]
    o' = o *>> loss
    inp  :: Prod t ('[i] ': ps >: '[o])
    inp = x :< p >: y
{-# INLINE netGrad #-}

ffLayer
    :: forall i o m t. (SingI i, SingI o, PrimMonad m, Tensor t)
    => Gen (PrimState m)
    -> m (Network t i o)
ffLayer g = (\w b -> buildNet ffLayer' (w :< b :< Ø))
          <$> genRand (normalDistr 0 0.5) g
          <*> genRand (normalDistr 0 0.5) g
  where
    ffLayer'
        :: TOp '[ '[i], '[o,i], '[o]] '[ '[o] ]
    ffLayer' = firstOp (TO.swap >>> TO.matVec)
           >>> TO.add
    {-# INLINE ffLayer' #-}
{-# INLINE ffLayer #-}

genNet
    :: forall k o i m (t :: [k] -> Type). (SingI o, SingI i, PrimMonad m, Tensor t)
    => [(Integer, Activation k)]
    -> Activation k
    -> Gen (PrimState m)
    -> m (Network t i o)
genNet xs0 f g = go sing xs0
  where
    go  :: forall (j :: k). ()
        => Sing j
        -> [(Integer, Activation k)]
        -> m (Network t j o)
    go sj = (\\ sj) $ \case
      []        -> (*~ getAct f) <$> ffLayer g
      (x,f'):xs -> withNatKind x $ \sl -> (\\ sl) $ do
        n <- go sl xs
        l <- ffLayer g
        return $ l *~ getAct f' ~*~ n
    {-# INLINE go #-}
{-# INLINE genNet #-}

max_stars_repo_name: mstksg/tensor-ops

max_issues_repo_name: mstksg/tensor-ops

max_forks_repo_name: mstksg/tensor-ops


text:
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE ViewPatterns        #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE FlexibleContexts    #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Data
-- Copyright   : David Johnson (c) 2019-2020
-- License     : BSD 3
-- Maintainer  : David Johnson <djohnson.m@gmail.com>
-- Stability   : Experimental
-- Portability : GHC
--
-- Functions for populating 'Array' with Data.
--
-- @
-- >>> constant @Double [2,2] 2.0
--  ArrayFire Array
-- [2 2 1 1]
--    2.0000     2.0000
--    2.0000     2.0000
-- @
--
--------------------------------------------------------------------------------
module ArrayFire.Data where

import Control.Exception
import Control.Monad
import Data.Complex
import Data.Int
import Data.Proxy
import Data.Word
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Marshal          hiding (void)
import Foreign.Storable
import System.IO.Unsafe
import Unsafe.Coerce

import ArrayFire.Exception
import ArrayFire.FFI
import ArrayFire.Internal.Data
import ArrayFire.Internal.Defines
import ArrayFire.Internal.Types
import ArrayFire.Arith

-- | Creates an 'Array' from a scalar value from given dimensions
--
-- >>> constant @Double [2,2] 2.0
--  ArrayFire Array
-- [2 2 1 1]
--    2.0000     2.0000
--    2.0000     2.0000
constant
  :: forall a . AFType a
  => [Int]
  -- ^ Dimensions
  -> a
  -- ^ Scalar value
  -> Array a
constant dims val =
  case dtyp of
    x | x == c64 ->
        cast $ constantComplex dims (unsafeCoerce val :: Complex Double)
      | x == c32 ->
        cast $ constantComplex dims (unsafeCoerce val :: Complex Float)
      | x == s64 ->
        cast $ constantLong dims (unsafeCoerce val :: Int)
      | x == u64 ->
        cast $ constantULong dims (unsafeCoerce val :: Word64)
      | x == s32 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Int32) :: Double)
      | x == s16 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Int16) :: Double)
      | x == u32 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Word32) :: Double)
      | x == u8 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Word8) :: Double)
      | x == u16 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Word16) :: Double)
      | x == f64 ->
        cast $ constant' dims (unsafeCoerce val :: Double)
      | x == b8  ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: CBool) :: Double)
      | x == f32 ->
        cast $ constant' dims (realToFrac (unsafeCoerce val :: Float))
      | otherwise -> error "constant: Invalid array fire type"
  where
    dtyp = afType (Proxy @ a)

    constant'
      :: [Int]
      -- ^ Dimensions
      -> Double
      -- ^ Scalar value
      -> Array Double
    constant' dims' val' =
      unsafePerformIO . mask_ $ do
        ptr <- alloca $ \ptrPtr -> do
          zeroOutArray ptrPtr
          withArray (fromIntegral <$> dims') $ \dimArray -> do
            throwAFError =<< af_constant ptrPtr val' n dimArray typ
            peek ptrPtr
        Array <$>
          newForeignPtr
            af_release_array_finalizer
              ptr
          where
            n = fromIntegral (length dims')
            typ = afType (Proxy @ Double)

    -- | Creates an 'Array (Complex Double)' from a scalar val'ue
    --
    -- @
    -- >>> constantComplex [2,2] (2.0 :+ 2.0)
    -- @
    --
    constantComplex
      :: forall arr . (Real arr, AFType (Complex arr))
      => [Int]
      -- ^ Dimensions
      -> Complex arr
      -- ^ Scalar val'ue
      -> Array (Complex arr)
    constantComplex dims' ((realToFrac -> x) :+ (realToFrac -> y)) = unsafePerformIO . mask_ $ do
      ptr <- alloca $ \ptrPtr -> do
        zeroOutArray ptrPtr
        withArray (fromIntegral <$> dims') $ \dimArray -> do
          throwAFError =<< af_constant_complex ptrPtr x y n dimArray typ
          peek ptrPtr
      Array <$>
        newForeignPtr
          af_release_array_finalizer
            ptr
          where
            n = fromIntegral (length dims')
            typ = afType (Proxy @ (Complex arr))

    -- | Creates an 'Array Int64' from a scalar val'ue
    --
    -- @
    -- >>> constantLong [2,2] 2.0
    -- @
    --
    constantLong
      :: [Int]
      -- ^ Dimensions
      -> Int
      -- ^ Scalar val'ue
      -> Array Int
    constantLong dims' val' = unsafePerformIO . mask_ $ do
      ptr <- alloca $ \ptrPtr -> do
        zeroOutArray ptrPtr
        withArray (fromIntegral <$> dims') $ \dimArray -> do
          throwAFError =<< af_constant_long ptrPtr (fromIntegral val') n dimArray
          peek ptrPtr
      Array <$>
        newForeignPtr
          af_release_array_finalizer
            ptr
          where
            n = fromIntegral (length dims')

    -- | Creates an 'Array Word64' from a scalar val'ue
    --
    -- @
    -- >>> constantULong [2,2] 2.0
    -- @
    --
    constantULong
      :: [Int]
      -> Word64
      -> Array Word64
    constantULong dims' val' = unsafePerformIO . mask_ $ do
      ptr <- alloca $ \ptrPtr -> do
        zeroOutArray ptrPtr
        withArray (fromIntegral <$> dims') $ \dimArray -> do
          throwAFError =<< af_constant_ulong ptrPtr (fromIntegral val') n dimArray
          peek ptrPtr
      Array <$>
        newForeignPtr
          af_release_array_finalizer
            ptr
          where
            n = fromIntegral (length dims')

-- | Creates a range of values in an Array
--
-- >>> range @Double [10] (-1)
-- ArrayFire Array
-- [10 1 1 1]
--     0.0000
--     1.0000
--     2.0000
--     3.0000
--     4.0000
--     5.0000
--     6.0000
--     7.0000
--     8.0000
--     9.0000
range
  :: forall a
   . AFType a
  => [Int]
  -> Int
  -> Array a
range dims (fromIntegral -> k) = unsafePerformIO $ do
  ptr <- alloca $ \ptrPtr -> mask_ $ do
    withArray (fromIntegral <$> dims) $ \dimArray -> do
      throwAFError =<< af_range ptrPtr n dimArray k typ
      peek ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        n = fromIntegral (length dims)
        typ = afType (Proxy @ a)

-- | Create an sequence [0, dims.elements() - 1] and modify to specified dimensions dims and then tile it according to tile_dims.
--
-- <http://arrayfire.org/docs/group__data__func__iota.htm>
--
-- >>> iota @Double [5,3] []
-- ArrayFire Array
-- [5 3 1 1]
--     0.0000     5.0000    10.0000
--     1.0000     6.0000    11.0000
--     2.0000     7.0000    12.0000
--     3.0000     8.0000    13.0000
--     4.0000     9.0000    14.0000
--
-- >>> iota @Double [5,3] [1,2]
-- ArrayFire Array
-- [5 6 1 1]
--     0.0000     5.0000    10.0000     0.0000     5.0000    10.0000
--     1.0000     6.0000    11.0000     1.0000     6.0000    11.0000
--     2.0000     7.0000    12.0000     2.0000     7.0000    12.0000
--     3.0000     8.0000    13.0000     3.0000     8.0000    13.0000
--     4.0000     9.0000    14.0000     4.0000     9.0000    14.0000
iota
  :: forall a . AFType a
  => [Int]
  -- ^ is the array containing sizes of the dimension
  -> [Int]
  -- ^ is array containing the number of repetitions of the unit dimensions
  -> Array a
  -- ^ is the generated array
iota dims tdims = unsafePerformIO $ do
  let dims' = take 4 (dims ++ repeat 1)
      tdims' =  take 4 (tdims ++ repeat 1)
  ptr <- alloca $ \ptrPtr -> mask_ $ do
    zeroOutArray ptrPtr
    withArray (fromIntegral <$> dims') $ \dimArray ->
      withArray (fromIntegral <$> tdims') $ \tdimArray -> do
        throwAFError =<< af_iota ptrPtr 4 dimArray 4 tdimArray typ
        peek ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        typ = afType (Proxy @ a)

-- | Creates the identity `Array` from given dimensions
--
-- >>> identity [2,2]
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     0.0000
--    0.0000     1.0000
identity
  :: forall a . AFType a
  => [Int]
  -- ^ Dimensions
  -> Array a
identity dims = unsafePerformIO . mask_ $ do
  let dims' = take 4 (dims ++ repeat 1)
  ptr <- alloca $ \ptrPtr -> mask_ $ do
    zeroOutArray ptrPtr
    withArray (fromIntegral <$> dims') $ \dimArray -> do
      throwAFError =<< af_identity ptrPtr n dimArray typ
      peek ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        n = fromIntegral (length dims)
        typ = afType (Proxy @ a)

-- | Create a diagonal matrix from input array when extract is set to false
--
-- >>> diagCreate (vector @Double 2 [1..]) 0
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     0.0000
--    0.0000     2.0000
diagCreate
  :: AFType (a :: *)
  => Array a
  -- ^	is the input array which is the diagonal
  -> Int
  -- ^ is the diagonal index
  -> Array a
diagCreate x (fromIntegral -> n) =
  x `op1` (\p a -> af_diag_create p a n)

-- | Create a diagonal matrix from input array when extract is set to false
--
-- >>> diagExtract (matrix @Double (2,2) [[1,2],[3,4]]) 0
-- ArrayFire Array
-- [2 1 1 1]
--     1.0000
--     4.0000
diagExtract
  :: AFType (a :: *)
  => Array a
  -> Int
  -> Array a
diagExtract x (fromIntegral -> n) =
  x `op1` (\p a -> af_diag_extract p a n)

-- | Join two Arrays together along a specified dimension
--
-- >>> join 0 (matrix @Double (2,2) [[1,2],[3,4]]) (matrix @Double (2,2) [[5,6],[7,8]])
-- ArrayFire Array
-- [4 2 1 1]
--     1.0000     3.0000
--     2.0000     4.0000
--     5.0000     7.0000
--     6.0000     8.0000
--
join
  :: Int
  -> Array (a :: *)
  -> Array a
  -> Array a
join (fromIntegral -> n) arr1 arr2 = op2 arr1 arr2 (\p a b -> af_join p n a b)

-- | Join many Arrays together along a specified dimension
--
-- *FIX ME*
--
-- >>> joinMany 0 [1,2,3]
-- ArrayFire Array
-- [3 1 1 1]
--    1.0000     2.0000     3.0000
--
joinMany
  :: Int
  -> [Array a]
  -> Array a
joinMany (fromIntegral -> n) arrays = unsafePerformIO . mask_ $ do
  fptrs <- forM arrays $ \(Array fptr) -> pure fptr
  newPtr <-
    alloca $ \fPtrsPtr -> do
      forM_ fptrs $ \fptr ->
        withForeignPtr fptr (poke fPtrsPtr)
      alloca $ \aPtr -> do
        zeroOutArray aPtr
        throwAFError =<< af_join_many aPtr n nArrays fPtrsPtr
        peek aPtr
  Array <$>
    newForeignPtr af_release_array_finalizer newPtr
  where
    nArrays = fromIntegral (length arrays)

-- | Tiles an Array according to specified dimensions
--
-- >>> tile @Double (scalar 22.0) [5,5]
-- ArrayFire Array
-- [5 5 1 1]
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
--
tile
  :: Array (a :: *)
  -> [Int]
  -> Array a
tile a (take 4 . (++repeat 1) -> [x,y,z,w]) =
  a `op1` (\p k -> af_tile p k (fromIntegral x) (fromIntegral y) (fromIntegral z) (fromIntegral w))
tile _ _ = error "impossible"

-- | Reorders an Array according to newly specified dimensions
--
-- *FIX ME*
--
-- >>> reorder @Double (scalar 22.0) [5,5]
-- ArrayFire Array
-- [5 5 1 1]
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
--
reorder
  :: Array (a :: *)
  -> [Int]
  -> Array a
reorder a (take 4 . (++ repeat 0) -> [x,y,z,w]) =
  a `op1` (\p k -> af_reorder p k (fromIntegral x) (fromIntegral y) (fromIntegral z) (fromIntegral w))
reorder _ _ = error "impossible"

-- | Shift elements in an Array along a specified dimension (elements will wrap).
--
-- >>> shift (vector @Double 4 [1..]) 2 0 0 0
-- ArrayFire Array
-- [4 1 1 1]
--     3.0000
--     4.0000
--     1.0000
--     2.0000
--
shift
  :: Array (a :: *)
  -> Int
  -> Int
  -> Int
  -> Int
  -> Array a
shift a (fromIntegral -> x) (fromIntegral -> y) (fromIntegral -> z) (fromIntegral -> w) =
  a `op1` (\p k -> af_shift p k x y z w)

-- | Modify dimensions of array
--
-- >>> moddims (vector @Double 3 [1..]) [1,3]
-- ArrayFire Array
-- [1 3 1 1]
--     1.0000     2.0000     3.0000
--
moddims
  :: forall a
   . Array (a :: *)
  -> [Int]
  -> Array a
moddims (Array fptr) dims =
  unsafePerformIO . mask_ . withForeignPtr fptr $ \ptr -> do
    newPtr <- alloca $ \aPtr -> do
      zeroOutArray aPtr
      withArray (fromIntegral <$> dims) $ \dimsPtr -> do
        throwAFError =<< af_moddims aPtr ptr n dimsPtr
        peek aPtr
    Array <$> newForeignPtr af_release_array_finalizer newPtr
  where
    n = fromIntegral (length dims)

-- | Flatten an Array into a single dimension
--
-- >>> flat (matrix @Double (2,2) [[1..],[1..]])
-- ArrayFire Array
-- [4 1 1 1]
--     1.0000
--     2.0000
--     1.0000
--     2.0000
--
-- >>> flat $ cube @Int (2,2,2) [[[1,1],[1,1]],[[1,1],[1,1]]]
-- ArrayFire Array
-- [8 1 1 1]
--          1
--          1
--          1
--          1
--          1
--          1
--          1
--          1
--
flat
  :: Array a
  -> Array a
flat = (`op1` af_flat)

-- | Flip the values of an Array along a specified dimension
--
-- >>> matrix @Double (2,2) [[2,2],[3,3]]
-- ArrayFire Array
-- [2 2 1 1]
--     2.0000     3.0000
--     2.0000     3.0000
--
-- >>> A.flip (matrix @Double (2,2) [[2,2],[3,3]]) 1
-- ArrayFire Array
-- [2 2 1 1]
--     3.0000     2.0000
--     3.0000     2.0000
--
flip
  :: Array a
  -> Int
  -> Array a
flip a (fromIntegral -> dim) =
  a `op1` (\p k -> af_flip p k dim)

-- | Create a lower triangular matrix from input array.
--
-- >>> lower (constant [2,2] 10 :: Array Double) True
-- ArrayFire Array
-- [2 2 1 1]
--     1.0000     0.0000
--    10.0000     1.0000
--
lower
  :: Array a
  -- ^ is the input matrix
  -> Bool
  -- ^ boolean parameter specifying if the diagonal elements should be 1
  -> Array a
lower a (fromIntegral . fromEnum -> b) =
  a `op1` (\p k -> af_lower p k b)

-- | Create an upper triangular matrix from input array.
--
-- >>> upper (constant [2,2] 10 :: Array Double) True
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     10.0000
--    0.0000     1.0000
--
upper
  :: Array a
  -> Bool
  -> Array a
upper a (fromIntegral . fromEnum -> b) =
  a `op1` (\p k -> af_upper p k b)

-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> arr2 = vector @Double 5 (repeat 2)
-- >>> select cond arr1 arr2
-- ArrayFire Array
-- [5 1 1 1]
--     1.0000
--     2.0000
--     1.0000
--     2.0000
--     1.0000
--
select
  :: Array CBool
  -- ^ is the conditional array
  -> Array a
  -- ^ is the array containing elements from the true part of the condition
  -> Array a
  -- ^	is the array containing elements from the false part of the condition
  -> Array a
  -- ^ is the output containing elements of a when cond is true else elements from b
select a b c = op3 a b c af_select

-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- <http://arrayfire.org/docs/group__data__func__select.htm#gab6886120d0bac4717276910e468bbe88>
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> x = 99
-- >>> selectScalarR cond arr1 x
-- ArrayFire Array
-- [5 1 1 1]
--     1.0000
--    99.0000
--     1.0000
--    99.0000
--     1.0000
--
selectScalarR
  :: Array CBool
  -- ^ is the conditional array
  -> Array a
  -- ^ is the array containing elements from the true part of the condition
  -> Double
  -- ^	is a scalar assigned to out when cond is false
  -> Array a
  -- ^ the output containing elements of a when cond is true else elements from b
selectScalarR a b c = op2 a b (\p w x -> af_select_scalar_r p w x c)

-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__data__func__select.htm#ga0ccdc05779f88cab5095bce987c2da9d)
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> x = 99
-- >>> selectScalarL cond x arr1
-- ArrayFire Array
-- [5 1 1 1]
--    99.0000
--     1.0000
--    99.0000
--     1.0000
--    99.0000
--
selectScalarL
  :: Array CBool
  -- ^ the conditional array
  -> Double
  -- ^ a scalar assigned to out when cond is true
  -> Array a
  -- ^ the array containing elements from the false part of the condition
  -> Array a
  -- ^ is the output containing elements of a when cond is true else elements from b
selectScalarL a n b = op2 a b (\p w x -> af_select_scalar_l p w n x)

-- af_err af_replace(af_array a, const af_array cond, const af_array b);
-- af_err af_replace_scalar(af_array a, const af_array cond, const double b);

max_stars_repo_name: arrayfire/arrayfire-haskell

max_issues_repo_name: arrayfire/arrayfire-haskell

max_forks_repo_name: arrayfire/arrayfire-haskell


text:
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_GHC -Wincomplete-patterns #-}

--{-# OPTIONS_GHC -W #-}

module Symphony.Symphony where

import AbsHashedLang
import Codec.Picture
import Control.Monad (when)
import Control.Monad.Except
import qualified Data.Array as Array
import Data.Char (toLower)
import Data.Complex
import Data.List (intercalate)
import Data.List.Extra (firstJust)
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe, mapMaybe)
import qualified Data.Set as Set
import Data.Tuple.HT (fst3)
import ErrM
import HashedExpression.Codegen
import qualified HashedExpression.Codegen.CSimple as CSimple
import HashedExpression.Internal
import HashedExpression.Internal.Expression (ExpressionMap, NodeID, Op (..))
import qualified HashedExpression.Internal.Expression as HE
import qualified HashedExpression.Internal.Node as HN
import HashedExpression.Internal.Utils
import qualified HashedExpression.Operation as HO
import HashedExpression.Prettify
import qualified HashedExpression.Problem as HS
import qualified HashedExpression.Value as HV
import LayoutHashedLang
import LexHashedLang
import ParHashedLang
import Symphony.Common
import Symphony.Exp
import Symphony.Solver
import Text.Regex.Posix

myLLexer :: String -> [Token]
myLLexer = resolveLayout True . myLexer

-------------------------------------------------------------------------------
-- Parse content of symphony file
-- Thanks to BNFC, we only need a simple call to `pProblem`
parse :: String -> Result Problem
parse fileContent =
  case pProblem . myLLexer $ fileContent of
    Ok problem -> return problem
    Bad errStr ->
      case (errStr =~ "line ([0-9]+), column ([0-9]+)" :: (String, String, String, [String])) of
        (_, _, _, [rs, cs]) ->
          -- because layout parsing add dummy ';'
          let r = min (read rs) (length r2c)
              c = min (read cs) (getNumColumn r)
           in throwError $ SyntaxError (r, c)
        _ -> throwError $ SyntaxError (1, 1)
  where
    r2c = Map.fromList . zip [1 ..] . map length . lines $ fileContent
    getNumColumn r = fromMaybe 0 $ Map.lookup r r2c

data ValidSymphony = ValidSymphony (ExpressionMap, NodeID) Vars Consts [HS.ConstraintStatement] AvailableSolver

-------------------------------------------------------------------------------
-- 1. Check if there is variables block
-- 2. Check if variable block is valid (no name clash)
-- 3. If there is a constant, check if it is valid (no name clash)
-- 4. Check if all expressions (in constraints and objective) is valid (should be scalar, all operation should be valid),
--        if yes construct Expression.
-- 5. Gen code
checkSemantic :: Problem -> Result ValidSymphony
checkSemantic problem = do
  let toVar name (shape, _) = varWithShape shape name
  let toParam name (shape, _) = paramWithShape shape name
  varDecls <- getVarDecls problem -- 1
  vars <- foldM checkVariableDecl Map.empty varDecls -- 2
  constDecls <- getConstDecls problem
  consts <- foldM (checkConstantDecl vars) Map.empty constDecls
  let initContext =
        Context
          { declarations = Map.mapWithKey toVar vars `Map.union` Map.mapWithKey toParam consts,
            vars = vars,
            consts = consts
          }
  letDecls <- getLetDecls problem
  finalContext <- foldM checkLetDecl initContext letDecls
  -- Complete processing context (variables, constants, declared values)
  -- Constraint
  constraintDecls <- getConstraintDecls problem
  css <- foldM (checkConstraintDecl finalContext) [] constraintDecls
  -- Objective
  parseObjectiveExp <- getMinimizeBlock problem
  objectiveExp <- constructExp finalContext (Just []) parseObjectiveExp
  -- Solver
  let shape = getShape objectiveExp
      nt = getNT objectiveExp
  solver <- getSolver problem
  when (shape /= [] || nt /= HE.R) $
    throwError $
      ErrorWithPosition
        ( "Objective should be a real scalar, here it is ("
            ++ toReadable shape
            ++ ", "
            ++ show nt
            ++ ")"
        )
        (getBeginningPosition parseObjectiveExp)
  liftIO $ putStrLn "Syntax & semantic is correct"
  return $ ValidSymphony objectiveExp vars consts css solver

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

-- | Generate code
generateCode :: String -> ValidSymphony -> Result ()
generateCode outputPath (ValidSymphony objectiveExp vars consts css solver) = do
  let problemGen =
        HS.constructProblem
          -- Objective
          (wrap objectiveExp)
          -- Constraints
          (HS.Constraint css)
  heProblem <- liftEitherString problemGen
--  case problemGen of
--    HS.ProblemValid heProblem -> do
  let valMap = Map.mapMaybeWithKey varVal vars `Map.union` Map.mapMaybeWithKey constVal consts
  res <- liftEitherString $ generateProblemCode CSimple.CSimpleConfig {output = CSimple.OutputHDF5} heProblem valMap
--  case generateProblemCode CSimple.CSimpleConfig {output = CSimple.OutputHDF5} heProblem valMap of
--    Invalid reason -> throwError $ GeneralError reason
--    Success res -> do
  liftIO $ putStrLn "Problem detail:"
  liftIO $ print $ heProblem
  liftIO $ res outputPath
  liftIO $ putStrLn "Download solver & adapter......"
  liftIO $ downloadSolver outputPath solver
--    HS.ProblemInvalid reason -> throwError $ GeneralError reason
  where
    varVal _ (_, Just val) = Just val
    varVal _ _ = Nothing
    constVal _ (_, val) = Just val

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

-- | Get variable declarations, there must be exact 1 variables block
getVarDecls :: Problem -> Result [VariableDecl]
getVarDecls (Problem blocks) =
  case filter isVariableBlock blocks of
    [] -> throwError $ GeneralError "No variale block"
    [BlockVariable declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 variables block"
  where
    isVariableBlock (BlockVariable declss) = True
    isVariableBlock _ = False

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

-- | Get constants declarations, there must be at most 1 constants block
getConstDecls :: Problem -> Result [ConstantDecl]
getConstDecls (Problem blocks) =
  case filter isConstantBlock blocks of
    [] -> return []
    [BlockConstant declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 constant block"
  where
    isConstantBlock (BlockConstant declss) = True
    isConstantBlock _ = False

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

-- | Get constraints declarations, there must be at most 1 constraints block
getConstraintDecls :: Problem -> Result [ConstraintDecl]
getConstraintDecls (Problem blocks) =
  case filter isConstraintBlock blocks of
    [] -> return []
    [BlockConstraint declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 constraint block"
  where
    isConstraintBlock (BlockConstraint declss) = True
    isConstraintBlock _ = False

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

-- | Get immediate declarations, at most 1 block
getLetDecls :: Problem -> Result [LetDecl]
getLetDecls (Problem blocks) =
  case filter isLetBlock blocks of
    [] -> return []
    [BlockLet declss] -> return $ concat declss
    _ -> throwError $ GeneralError "There are more than 1 let block"
  where
    isLetBlock (BlockLet declss) = True
    isLetBlock _ = False

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

-- | Get objective block, must be 1 block
getMinimizeBlock :: Problem -> Result Exp
getMinimizeBlock (Problem blocks) =
  case filter isMinimizeBlock blocks of
    [] -> throwError $ GeneralError "No minimize block"
    [BlockMinimize exp] -> return exp
    _ -> throwError $ GeneralError "There are more than 1 minimize block"
  where
    isMinimizeBlock (BlockMinimize declss) = True
    isMinimizeBlock _ = False

-- | Get the solver
getSolver :: Problem -> Result AvailableSolver
getSolver (Problem blocks) =
  case filter isSolverBlock blocks of
    [] -> do
      lift $ putStrLn "No solver specified, default to L-BFGS-B (https://github.com/stephenbeckr/L-BFGS-B-C)"
      return LBFGSB
    [BlockSolver (SolverName name)]
      | map toLower name `elem` ["lbfgs-b", "lbfgsb", "l-bfgs-b"] -> return LBFGSB
      | map toLower name `elem` ["ipopt"] -> return Ipopt
      | otherwise -> throwError $ GeneralError "Unknown solver, should be (ipopt | lbfgs-b)"
    _ -> throwError $ GeneralError "More than one solver specified"
  where
    isSolverBlock (BlockSolver name) = True
    isSolverBlock _ = False

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

-- | Check variable declarations
checkVariableDecl :: Vars -> VariableDecl -> Result Vars
checkVariableDecl accRes decl = do
  let checkName name pos =
        when (name `Map.member` accRes) $
          throwError $
            ErrorWithPosition ("Duplicate declaration of " ++ name) pos
  case decl of
    VariableNoInit (PIdent (pos, name)) shape -> do
      checkName name pos
      return $ Map.insert name (toHEShape shape, Nothing) accRes
    VariableWithInit (PIdent (pos, name)) shape val -> do
      checkName name pos
      checkVal (toHEShape shape) val
      let heShape = toHEShape shape
      heVal <- toHEVal heShape val
      return $ Map.insert name (heShape, Just heVal) accRes

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

-- | Check constant declarations
checkConstantDecl :: Vars -> Consts -> ConstantDecl -> Result Consts
checkConstantDecl vars accRes decl = do
  let (ConstantDecl (PIdent (pos, name)) shape val) = decl
  when (name `Map.member` accRes) $
    throwError $
      ErrorWithPosition ("Duplicate declaration of " ++ name) pos
  when (name `Map.member` vars) $
    throwError $
      ErrorWithPosition (name ++ " already defined as variables") pos
  let heShape = toHEShape shape
  checkVal heShape val
  heVal <- toHEVal heShape val
  return $ Map.insert name (toHEShape shape, heVal) accRes

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

-- | Check constraints declarations
checkConstraintDecl ::
  Context ->
  [HS.ConstraintStatement] ->
  ConstraintDecl ->
  Result [HS.ConstraintStatement]
checkConstraintDecl context@Context {..} acc decl = do
  (constructedExp, boundVal) <-
    case (isVariable exp, bound) of
      (Just varExp, ConstantBound (PIdent (pos, name)))
        | Just (shape, val) <- Map.lookup name consts -> do
          when (shape /= getShape varExp) $
            throwError $
              ErrorWithPosition
                "Shape mismatched: the bound doesn't have same shape as the variable"
                pos
          return (varExp, val)
        | otherwise ->
          throwError $ ErrorWithPosition (name ++ " not found") pos
      (Just varExp, NumberBound num) ->
        return (varExp, HV.VNum (numToDouble num))
      _ -> do
        scalarExp <- constructExp context Nothing exp
        when (getShape scalarExp /= []) $
          throwError $
            ErrorWithPosition
              ( "Higher-dimension (in)equality is not supported yet, here the expression has shape "
                  ++ toReadable (getShape scalarExp)
              )
              (getBeginningPosition exp)
        case bound of
          ConstantBound (PIdent (pos, name))
            | Just (shape, val) <- Map.lookup name consts -> do
              when (shape /= getShape scalarExp) $
                throwError $
                  ErrorWithPosition "The bound must be scalar" pos
              return (scalarExp, val)
            | otherwise ->
              throwError $
                ErrorWithPosition (name ++ " not found") pos
          NumberBound num ->
            return (scalarExp, HV.VNum (numToDouble num))
  let newEntry =
        case decl of
          ConstraintLower {} -> HS.Lower constructedExp boundVal
          ConstraintUpper {} -> HS.Upper constructedExp boundVal
          ConstraintEqual {} ->
            HS.Between constructedExp (boundVal, boundVal)
  return $ acc ++ [newEntry]
  where
    (exp, bound) =
      case decl of
        ConstraintLower exp bound -> (exp, bound)
        ConstraintUpper exp bound -> (exp, bound)
        ConstraintEqual exp bound -> (exp, bound)
    isVariable exp =
      case exp of
        EIdent (PIdent (_, name))
          | Just constructedExp <- Map.lookup name declarations,
            name `Map.member` vars ->
            Just constructedExp
        _ -> Nothing

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

-- | Check immediate declarations
checkLetDecl :: Context -> LetDecl -> Result Context
checkLetDecl context@Context {..} (LetDecl (PIdent (pos, name)) exp) = do
  when (name `Map.member` consts) $
    throwError $
      ErrorWithPosition (name ++ " already defined as a constant") pos
  when (name `Map.member` vars) $
    throwError $
      ErrorWithPosition (name ++ " already defined as a variable") pos
  when (name `Map.member` declarations) $
    throwError $
      ErrorWithPosition (name ++ " already taken") pos
  constructedExp <- constructExp context Nothing exp
  let newDeclarations = Map.insert name constructedExp declarations
  return $ context {declarations = newDeclarations}

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

-- | To HashedExpression's value
toHEVal :: HE.Shape -> Val -> Result HV.Val
toHEVal shape v =
  case v of
    ValFile filePath -> return $ HV.VFile $ HV.TXT filePath
    ValDataset filePath dataset ->
      return $ HV.VFile $ HV.HDF5 filePath dataset
    ValPattern (KWDataPattern pattern) ->
      case (pattern, shape) of
        ("FIRST_ROW_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate size2 1 ++ repeat 0
        ("LAST_ROW_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate (size2 * (size1 - 1)) 0 ++ repeat 1
        ("FIRST_COLUMN_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                1 : replicate (size2 - 1) 0
        ("LAST_COLUMN_1", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                replicate (size2 - 1) 0 ++ [1]
        ("FIRST_ROW_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate size2 0 ++ repeat 1
        ("LAST_ROW_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ replicate (size2 * (size1 - 1)) 1 ++ repeat 0
        ("FIRST_COLUMN_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                0 : replicate (size2 - 1) 1
        ("LAST_COLUMN_0", [size1, size2]) ->
          return
            . HV.V2D
            . Array.listArray ((0, 0), (size1 - 1, size2 - 1))
            $ concat $
              replicate size1 $
                replicate (size2 - 1) 1 ++ [0]
        _ ->
          throwError $
            GeneralError $
              "Pattern "
                ++ pattern
                ++ " is incompatible with the shape or not supported yet"
    ValRandom -> return $ HV.VNum 3
    ValLiteral num -> return $ HV.VNum (numToDouble num)
    ValImage imgPath -> do
      a <- liftIO $ readImage imgPath
      case a of
        Left err -> throwError $ GeneralError $ "Error reading image at " <> imgPath <> ":" <> err
        Right v -> do
          -- TODO : only support grayscale yet, and this is very slow
          let img = convertRGB8 v
          let col = imageWidth img
              row = imageHeight img
              toGrayscale :: Pixel8 -> Pixel8 -> Pixel8 -> Double
              toGrayscale r g b = (0.2126 * (fromIntegral r) + 0.7152 * (fromIntegral g) + 0.0722 * (fromIntegral b)) / 256
          if (shape /= [row, col])
            then throwError $ GeneralError $ "image size and variable shape don't match, image size is " ++ show row ++ "x" ++ show col
            else
              return $
                HV.V2D $
                  Array.listArray
                    ((0, 0), (row - 1, col - 1))
                    [ toGrayscale r g b
                      | i <- [0 .. row - 1],
                        j <- [0 .. col - 1],
                        let (PixelRGB8 r g b) = pixelAt img j i
                    ]

max_stars_repo_name: McMasterU/Symphony

max_issues_repo_name: McMasterU/Symphony

max_forks_repo_name: McMasterU/Symphony


text:
{-# LANGUAGE TypeFamilies, PatternGuards #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Statistics.Order
-- Copyright   :  (C) 2012 Edward Kmett,
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  Haskell 2011 + TypeFamilies
--
----------------------------------------------------------------------------

module Statistics.Order
  (
  -- * L-Estimator
    L(..)
  -- ** Applying an L-estimator
  , (@@), (@!)
  -- ** Analyzing an L-estimator
  , (@#)
  , breakdown
  -- ** Robust L-Estimators
  , trimean  -- Tukey's trimean
  , midhinge -- average of q1 and q3
  , iqr      -- interquartile range
  , iqm      -- interquartile mean
  , lscale   -- second L-moment
  -- ** L-Estimator Combinators
  , trimmed
  , winsorized, winsorised
  , jackknifed
  -- ** Trivial L-Estimators
  , mean
  , total
  , lmin, lmax
  , midrange
  -- ** Sample-size-dependent L-Estimators
  , nthSmallest
  , nthLargest
  -- ** Quantiles
  -- *** Common quantiles
  , quantile
  , median
  , tercile, t1, t2
  , quartile, q1, q2, q3
  , quintile, qu1, qu2, qu3, qu4
  , percentile
  , permille
  -- *** Harrell-Davis Quantile Estimator
  , hdquantile
  -- *** Compute a quantile using a specified quantile estimation strategy
  , quantileBy
  -- * Sample Quantile Estimators
  , Estimator
  , Estimate(..)
  , r1,r2,r3,r4,r5,r6,r7,r8,r9,r10
  ) where

import Data.Ratio
import Data.List (sort)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM
import Data.Vector (Vector, (!))
import qualified Data.Vector as V
import Data.VectorSpace
import Statistics.Distribution.Beta
import qualified Statistics.Distribution as D

-- | L-estimators are linear combinations of order statistics used by 'robust' statistics.
newtype L r = L { runL :: Int -> IntMap r }

-- TODO: Write code to test if the result of a given L-estimator will be always than or equal
-- to the result of another L-estimator at a given sample size.

-- | Calculate the result of applying an L-estimator after sorting list into order statistics
(@@) :: (Num r, Ord r) => L r -> [r] -> r
l @@ xs = l @! V.fromList (sort xs)

-- | Calculate the result of applying an L-estimator to a *pre-sorted* vector of order statistics
(@!) :: Num r => L r -> Vector r -> r
L f @! v = IM.foldrWithKey (\k x y -> (v ! k) * x + y) 0 $ f (V.length v)

-- | get a vector of the coefficients of an L estimator when applied to an input of a given length
(@#) :: Num r => L r -> Int -> [r]
L f @# n = map (\k -> IM.findWithDefault 0 k fn) [0..n-1] where fn = f n

-- | calculate the breakdown % of an L-estimator
breakdown :: (Num r, Eq r) => L r -> Int
breakdown (L f)
  | IM.null m = 50
  | otherwise = fst (IM.findMin m) `min` (100 - fst (IM.findMax m))
  where m = IM.filter (/= 0) $ f 101

instance Num r => AdditiveGroup (L r) where
  zeroV = L $ \_ -> IM.empty
  L fa ^+^ L fb = L $ \n -> IM.unionWith (+) (fa n) (fb n)
  negateV (L fa) = L $ fmap negate . fa

instance Num r => VectorSpace (L r) where
  type Scalar (L r) = r
  x *^ L y = L $ fmap (x *) . y

clamp :: Int -> Int -> Int
clamp n k
  | k <= 0 = 0
  | k >= n = n - 1
  | otherwise = k

-- | The average of all of the order statistics. Not robust.
--
-- > breakdown mean = 0%
mean :: Fractional r => L r
mean = L $ \n -> IM.fromList [ (i, 1 / fromIntegral n) | i <- [0..n-1]]

-- | The sum of all of the order statistics. Not robust.
--
-- > breakdown total = 0%
total :: Num r => L r
total = L $ \n -> IM.fromList [ (i, 1) | i <- [0..n-1]]

-- | Calculate a trimmed L-estimator. If the sample size isn't evenly divided, linear interpolation is used
-- as described in <http://en.wikipedia.org/wiki/Trimmed_mean#Interpolation>

-- Trimming can increase the robustness of a statistic by removing outliers.

trimmed :: Fractional r => Rational -> L r -> L r
trimmed p (L g) = L $ \n -> case properFraction (fromIntegral n * p) of
  (w, 0)               -> IM.fromDistinctAscList [ (k + w, v)                        | (k,v) <- IM.toAscList $ g (n - w*2)]
  (w, f) | w' <- w + 1 -> IM.fromListWith (+) $  [ (k + w, fromRational (1 - f) * v) | (k,v) <- IM.toList $ g (n - w*2)] ++
                                                 [ (k + w', fromRational f  * v)     | (k,v) <- IM.toList $ g (n - w'*2)]

-- | Calculates an interpolated winsorized L-estimator in a manner analogous to the trimmed estimator.
-- Unlike trimming, winsorizing replaces the extreme values.
winsorized, winsorised :: Fractional r => Rational -> L r -> L r
winsorised = winsorized
winsorized p (L g) = L $ \n -> case properFraction (fromIntegral n * p) of
  (w, 0)               -> IM.fromAscListWith (+) [ (w `max` min (n - 1 - w) k, v) | (k,v) <- IM.toAscList (g n) ]
  (w, f) | w' <- w + 1 -> IM.fromListWith (+) $ do
     (k,v) <- IM.toList (g n)
     [ (w  `max` min (n - 1 - w ) k, v * fromRational (1 - f)),
       (w' `max` min (n - 1 - w') k, v * fromRational f)]

-- | Jackknifes the statistic by removing each sample in turn and recalculating the L-estimator,
-- requires at least 2 samples!
jackknifed :: Fractional r => L r -> L r
jackknifed (L g) = L $ \n -> IM.fromAscListWith (+) $ do
  let n' = fromIntegral n
  (k,v) <- IM.toAscList (g (n - 1))
  let k' = fromIntegral k + 1
  [(k, (n' - k') * v / n'), (k + 1, k' * v / n')]

-- | The most robust L-estimator possible.
--
-- > breakdown median = 50
median :: Fractional r => L r
median = L go where
  go n
    | odd n        = IM.singleton (div (n - 1) 2) 1
    | k <- div n 2 = IM.fromList [(k-1, 0.5), (k, 0.5)]

-- | Sample quantile estimators
data Estimate r  = Estimate {-# UNPACK #-} !Rational (IntMap r)
type Estimator r = Rational -> Int -> Estimate r

-- | Compute a quantile using the given estimation strategy to interpolate when an exact quantile isn't available
quantileBy :: Num r => Estimator r -> Rational -> L r
quantileBy f p = L $ \n -> case f p n of
  Estimate h qp -> case properFraction h of
    (w, 0) -> IM.singleton (clamp n (w - 1)) 1
    _      -> qp

-- | Compute a quantile with traditional direct averaging
quantile :: Fractional r => Rational -> L r
quantile = quantileBy r2

tercile :: Fractional r => Rational -> L r
tercile n = quantile (n/3)

-- | terciles 1 and 2
--
-- > breakdown t1 = breakdown t2 = 33%
t1, t2 :: Fractional r => L r
t1 = quantile (1/3)
t2 = quantile (2/3)

quartile :: Fractional r => Rational -> L r
quartile n = quantile (n/4)

-- | quantiles, with breakdown points 25%, 50%, and 25% respectively
q1, q2, q3 :: Fractional r => L r
q1 = quantile 0.25
q2 = median
q3 = quantile 0.75

quintile :: Fractional r => Rational -> L r
quintile n = quantile (n/5)

-- | quintiles 1 through 4
qu1, qu2, qu3, qu4 :: Fractional r => L r
qu1 = quintile 1
qu2 = quintile 2
qu3 = quintile 3
qu4 = quintile 4

-- |
--
-- > breakdown (percentile n) = min n (100 - n)
percentile :: Fractional r => Rational -> L r
percentile n = quantile (n/100)

permille :: Fractional r => Rational -> L r
permille n = quantile (n/1000)

nthSmallest :: Num r => Int -> L r
nthSmallest k = L $ \n -> IM.singleton (clamp n k) 1

nthLargest :: Num r => Int -> L r
nthLargest k = L $ \n -> IM.singleton (clamp n (n - 1 - k)) 1

-- |
--
-- > midhinge = trimmed 0.25 midrange
-- > breakdown midhinge = 25%
midhinge :: Fractional r => L r
midhinge = (q1 ^+^ q3) ^/ 2

-- | Tukey's trimean
--
-- > breakdown trimean = 25
trimean :: Fractional r => L r
trimean = (q1 ^+^ 2 *^ q2 ^+^ q3) ^/ 4

-- | The maximum value in the sample
lmax :: Num r => L r
lmax = L $ \n -> IM.singleton (n-1) 1

-- | The minimum value in the sample
lmin :: Num r => L r
lmin = L $ \_ -> IM.singleton 0 1

-- |
-- > midrange = lmax - lmin
-- > breakdown midrange = 0%
midrange :: Fractional r => L r
midrange = L $ \n -> IM.fromList [(0,-1),(n-1,1)]

-- | interquartile range
--
-- > breakdown iqr = 25%
-- > iqr = trimmed 0.25 midrange
iqr :: Fractional r => L r
iqr = q3 ^-^ q1

-- | interquartile mean
--
-- > iqm = trimmed 0.25 mean
iqm :: Fractional r => L r
iqm = trimmed 0.25 mean

-- | Direct estimator for the second L-moment given a sample
lscale :: Fractional r => L r
lscale = L $ \n -> let
     r = fromIntegral n
     scale = 1 / (r * (r-1))
  in IM.fromList [ (i - 1, scale * (2 * fromIntegral i - 1 - r)) | i <- [1..n] ]

-- | The Harrell-Davis quantile estimate. Uses multiple order statistics to approximate the quantile
-- to reduce variance.
hdquantile :: Fractional r => Rational -> L r
hdquantile q = L $ \n ->
  let n' = fromIntegral n
      np1 = n' + 1
      q' = fromRational q
      d = betaDistr (q'*np1) (np1*(1-q')) in
  if q == 0 then IM.singleton 0 1
  else if q == 1 then IM.singleton (n - 1) 1
  else IM.fromAscList
    [ (i, realToFrac $ D.cumulative d ((fromIntegral i + 1) / n') -
                       D.cumulative d (fromIntegral i / n'))
    | i <- [0 .. n - 1]
    ]

-- More information on the individual estimators used below can be found in:
-- http://stat.ethz.ch/R-manual/R-devel/library/stats/html/quantile.html
-- and
-- http://en.wikipedia.org/wiki/Quantile#Estimating_the_quantiles_of_a_population

-- | Inverse of the empirical distribution function
r1 :: Num r => Estimator r
r1 p n = Estimate (np + 1%2) $ IM.singleton (clamp n (ceiling np - 1)) 1
  where np = fromIntegral n * p

-- | .. with averaging at discontinuities
r2 :: Fractional r => Estimator r
r2 p n =  Estimate (np + 1%2) $
  if p == 0      then IM.singleton 0       1
  else if p == 1 then IM.singleton (n - 1) 1
  else IM.fromList [(ceiling np - 1, 0.5), (floor np, 0.5)]
  where np = fromIntegral n * p

-- | The observation numbered closest to Np. NB: does not yield a proper median
r3 :: Num r => Estimator r
r3 p n = Estimate np $ IM.singleton (clamp n (round np - 1)) 1
  where np = fromIntegral n * p

-- continuous sample quantiles
continuousEstimator ::
  Fractional r =>
  (Rational -> (Rational, Rational)) -> -- take the number of samples, and return upper and lower bounds on 'p = k/n' for which this interpolation should be used
  (Rational -> Rational -> Rational) -> -- take p = k/q, and n the number of samples, and return the coefficient h which will be used for interpolation when h is not integral
  Estimator r
continuousEstimator bds f p n = Estimate h $
  if p < lo then IM.singleton 0 1
  else if p >= hi then IM.singleton (n - 1) 1
  else case properFraction h of
    (w,frac) | frac' <- fromRational frac -> IM.fromList [(w - 1, frac'), (w, 1 - frac')]
  where
    r = fromIntegral n
    h = f p r
    (lo, hi) = bds r

-- | Linear interpolation of the empirical distribution function. NB: does not yield a proper median.
r4 :: Fractional r => Estimator r
r4 = continuousEstimator (\n -> (1 / n, 1)) (*)

-- | .. with knots midway through the steps as used in hydrology. This is the simplest continuous estimator that yields a correct median
r5 :: Fractional r => Estimator r
r5 = continuousEstimator (\n -> let tn = 2 * n in (1 / tn, (tn - 1) / tn)) $ \p n -> p*n + 0.5

-- | Linear interpolation of the expectations of the order statistics for the uniform distribution on [0,1]
r6 :: Fractional r => Estimator r
r6 = continuousEstimator (\n -> (1 / (n + 1), n / (n + 1))) $ \p n -> p*(n+1)

-- | Linear interpolation of the modes for the order statistics for the uniform distribution on [0,1]
r7 :: Fractional r => Estimator r
r7 = continuousEstimator (\_ -> (0, 1)) $ \p n -> p*(n-1) + 1

-- | Linear interpolation of the approximate medans for order statistics.
r8 :: Fractional r => Estimator r
r8 = continuousEstimator (\n -> (2/3 / (n + 1/3), (n - 1/3)/(n + 1/3))) $ \p n -> p*(n + 1/3) + 1/3

-- | The resulting quantile estimates are approximately unbiased for the expected order statistics if x is normally distributed.
r9 :: Fractional r => Estimator r
r9 = continuousEstimator (\n -> (0.625 / (n + 0.25), (n - 0.375)/(n + 0.25))) $ \p n -> p*(n + 0.25) + 0.375

-- | When rounding h, this yields the order statistic with the least expected square deviation relative to p.
r10 :: Fractional r => Estimator r
r10 = continuousEstimator (\n -> (1.5 / (n + 2), (n + 0.5)/(n + 2))) $ \p n -> p*(n + 2) - 0.5

max_stars_repo_name: ekmett/order-statistics

max_issues_repo_name: ekmett/order-statistics

max_forks_repo_name: ekmett/order-statistics


text:
{-# LANGUAGE RecordWildCards #-}

module Data.BenchAna.Stats(stats) where

import qualified Data.Vector.Unboxed as V
import qualified Statistics.Sample as S
import qualified Statistics.Function as S

import Data.BenchAna.Types

stats :: BenchResults -> BenchStats
stats = fmap singleStats

singleStats :: V.Vector Int -> Stats
singleStats vec = Stats { .. }
    where statsWithOutliers = singleStats' vec'
          statsNoOutliers = singleStats' $ V.init $ V.tail $ S.sort vec'
          vec' = V.map fromIntegral vec

singleStats' :: V.Vector Double -> StatsPair
singleStats' vec = StatsPair { .. }
    where mean = S.mean vec
          stddev = S.stdDev vec

max_stars_repo_name: Maledictus/leechcraft

max_issues_repo_name: Maledictus/leechcraft

max_forks_repo_name: Maledictus/leechcraft