module Weighted
       where

import Control.Applicative
import Control.Monad
import Data.Foldable
import Data.Monoid

newtype Weight = Weight Float deriving (Show,Eq,Ord)

instance Monoid Weight where
  mempty  = Weight 1
  Weight w `mappend` Weight w'  =  Weight (w * w')

inv :: Weight -> Weight
inv (Weight w) = (Weight (1.0 - w))
                                   
data Weighted a = Weighted Weight a

unweighted :: Weighted a -> a
unweighted (Weighted _ x) = x

weight :: Weighted a -> Weight
weight (Weighted w _) = w 

instance Functor Weighted where
    fmap f (Weighted w x) = Weighted w (f x)

instance Applicative Weighted where
    pure = Weighted mempty
    Weighted w f <*> Weighted w' x = Weighted (w <> w') (f x)

instance Monad Weighted where
    return = pure
    (Weighted w x) >>= f = let (Weighted w' y) = f x
                           in Weighted (w<>w') y


        
    
instance (Show a) => Show (Weighted a) where
    show (Weighted (Weight w) x) = show x ++ "#" ++ show w

infix 9 #
x#w = Weighted (Weight w) x

instance Eq (Weighted a) where
    Weighted w1 _ == Weighted w2 _ =  w1 == w2

instance Ord (Weighted a) where
    Weighted w1 _ <= Weighted w2 _ =  w1 <= w2



newtype WeightedT m a = WeightedT { runWeightedT :: m (Weighted a) }

instance Functor m => Functor (WeightedT m) where
    f `fmap` x = WeightedT (liftA f `fmap` runWeightedT x)

instance Applicative m => Applicative (WeightedT m) where
    pure = WeightedT . pure . pure
    f <*> x = WeightedT (liftA2 (<*>) (runWeightedT f) (runWeightedT x))

instance Alternative m => Alternative (WeightedT m) where
    empty = WeightedT empty
    x <|> y = WeightedT ( runWeightedT x <|> runWeightedT y )

instance (MonadPlus m, Foldable m) => Foldable (WeightedT m) where
    -- foldr f z c = foldr f z (map unweighted $ toList $ runWeightedT c)
    foldr f z c = foldr f z (fmap unweighted $ runWeightedT c)
    -- foldr f z c = foldr (f . unweighted) z (runWeightedT c)

instance Monad m => Monad (WeightedT m) where
    return = WeightedT . return . return
    x >>= f = WeightedT $ do
                Weighted w a <- runWeightedT x
                Weighted w' b <- runWeightedT (f a)
                return (Weighted (w<>w') b)

instance (Alternative m, Monad m) => MonadPlus (WeightedT m)
                
-- instance (MonadPlus m, Traversable m) => Traversable (WeightedT m) where
--     traverse f xs = return $ traverse (liftM f) (runWeightedT xs)

    -- traverse f xs | null xs   = pure mzero
    --               | otherwise = foldr (\x v -> mplus <$> f x <*> v) (pure mzero) xs


    -- (WeightedT x)= WeightedT $ sequenceA x

    

type WeightedList = WeightedT []