import Hardware.Chalk
import Control.Applicative

newtype TSignal a = TSignal (Signal (Ticked a))

instance Functor TSignal where
  fmap f (TSignal x) = TSignal (fmap (fmap f) x)

instance Applicative TSignal where
  pure x = TSignal (pure (pure x))
  (TSignal s) <*> (TSignal x) = TSignal (pure (<*>) <*> s <*> x)

data Ticked a = Ticked {tval :: a, cost :: Int}

instance Functor Ticked where
  fmap f (Ticked x cost) = Ticked (f x) cost

instance Applicative Ticked where
  pure x = Ticked x 0
  (<*>) (Ticked f c1) (Ticked x c2) = Ticked (f x) (c1 + c2)

instance Monad Ticked where
  return x = Ticked x 0
  c >>= f = join (fmap f c)

join :: Ticked (Ticked a) -> Ticked a            
join (Ticked (Ticked x c1) c2) = Ticked x (c1 + c2)

data Cmd = ADD | MUL deriving (Show, Eq, Read)

tick :: TSignal a -> TSignal a
tick (TSignal t) = TSignal (fmap inc t)
  where
  inc :: Ticked a -> Ticked a
  inc (Ticked x c) = Ticked x (c + 1)

pay :: Int -> TSignal a -> TSignal a
pay i (TSignal t) = TSignal (fmap (\x -> x {cost = cost x + i}) t)

costed :: Int -> a -> TSignal a
costed i x = pay i (pure x)

sjoin :: TSignal (Ticked a) -> TSignal a
sjoin (TSignal s) = TSignal (fmap join s)

alu :: TSignal Cmd -> TSignal (Int, Int) -> TSignal Int
alu cmds xys = pure cmdCase <*> cmds <*> adder xys <*> multiplier xys
  where
  cmdCase ADD x y = x
  cmdCase MUL x y = y
  adder xys  = pay adderCost (pure (uncurry (+)) <*> xys)
  multiplier xys = mux (sizeTest xys) (cheapMul xys) (dearMul xys)
  sizeTest xys = costed 1 (\(x,y) -> x < threshhold && y < threshhold) <*> xys
  cheapMul xys = costed 2 (uncurry (*)) <*> xys
  dearMul xys = costed 5 (uncurry (*)) <*> xys

adderCost = 3
sizeTest = 1
cheapMul = 3
dearMul = 6
threshhold = 32

mux :: TSignal Bool -> TSignal a -> TSignal a -> TSignal a
mux bs ts es = costed 1 cond <*> bs <*> ts <*> es
  where
  cond True x _ = x
  cond False _ x = x

smartMux :: Signal (Ticked Bool) -> Signal (Ticked a) 
  -> Signal (Ticked a) -> Signal (Ticked a)
smartMux bs ts es = pure f <*> bs <*> ts <*> es
  where
  f :: Ticked Bool -> Ticked a -> Ticked a -> Ticked a
  f (Ticked True c) t e = Ticked (tval t) (c + cost t)
  f (Ticked False c) t e = Ticked (tval e) (c + cost e)