{-# LANGUAGE FlexibleContexts , BangPatterns #-}
module Perceptron.Vector  
    ( I(..)
    , Global
    , Local(..)
    , Weights
    , WeightsST
    , toSV
    , for_
    , plus_
    , minus_
    , plus
    , minus
    , scale
    , dot 
    , dot'
    )
where

import Data.Array.ST
import Data.Array.Unboxed
import Control.Monad.ST
import Data.STRef
import Control.Monad
import qualified Data.Map as Map
import Data.List (foldl',sort)
import qualified Data.Vector.Unboxed as V
import Data.Binary

data I = I {-# UNPACK #-} !Int {-# UNPACK #-} !Int deriving (Eq,Ord,Ix,Show)
instance Binary I where 
    put (I i j) = put (i,j)
    get = uncurry I `fmap` get

type Global = Map.Map I Float
data Local  = Local {-# UNPACK #-} !Int !(V.Vector Int)
type WeightsST s = STUArray s I Float
type Weights = UArray I Float



for_ xs f = mapM_ f xs

plus_ :: WeightsST s -> Global -> ST s ()
plus_ w v = do
  for_ (Map.toList v) $ \(i,vi) -> do
              wi <- readArray w i 
              writeArray w i (wi + vi)
minus_ w v = plus_ w (v `scale` (-1))

scale :: Global -> Float -> Global
scale v n = Map.map (*n) v

plus :: Global -> Global -> Global
plus u v = Map.unionWith (+) u v
minus :: Global -> Global -> Global
minus u v = u `plus` (v `scale` (-1))


dot :: Weights -> Local -> Float
{-# INLINE dot #-}
dot w (Local !y x) = V.foldl' (\ !z !i -> z + w ! I y i) 0 x
-- For some reason explicit loop doesn't help here
-- dot !w (Local y x) = go 0 0
--     where !len = V.length x
--           go !z !j | j == len = z
--           go !z !j = go (z + w ! I y (x V.! j)) (j+1)


dot' :: (Float,Weights,Weights) -> Local -> Float
{-# INLINE dot' #-}
-- dot' (!c,!params,!params_a) (Local y x) = V.foldl' (\ !z !j -> 
--                                                     let i   = I y j
--                                                         e   = params   ! i 
--                                                         e_a = params_a ! i
--                                                  in z + (e - (e_a / c)))
--                                             0
--                                             x


dot' (!c,!params,!params_a) (Local y x) = go 0 0
    where !len = V.length x
          go !z !j | j == len = z
          go !z !j = 
              let i   = I y (x V.! j)
                  e   = params   ! i
                  e_a = params_a ! i
              in  go (z + (e - (e_a / c))) (j+1)

toSV :: (V.Unbox Int) => Local -> Global
toSV (Local y v) = Map.fromList [ (I y i,1) | i <- V.toList v ]