-- Copyright 2022 University of Freiburg
-- Janek Spaderna <janek.spaderna@pluto.uni-freiburg.de>
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module Ex03 where

import Control.Exception (evaluate)
import Control.Monad (void)
import GHC.Generics (Generic)
import Test.QuickCheck

-------------------------------------------------------------------------------
-- Exercise 1.1

filter' :: (a -> Bool) -> [a] -> [a]
filter' f = foldr' g []
  where
    -- With the ScopedTypeVariables extension, and an explicit `forall` in the
    -- type signature for `filter` we could write the signature for g:
    --
    -- g :: a -> [a] -> [a]
    --
    -- Without either the extension or the explicit `forall` GHC will complain
    -- that the two `a`s are different.
    g a as = if f a then a : as else as

remdups :: Eq a => [a] -> [a]
remdups = foldr' f []
  where
    -- f :: a -> [a] -> [a]
    f a [] = [a]
    f a (a' : as)
      | a == a' = a' : as
      | otherwise = a : a' : as

avg :: [Double] -> Double
avg [] = error "avg of zero elements"
avg ds = dsSum / dsCount
  where
    (dsSum, dsCount) = foldr' f (0, 0) ds

    f :: Double -> (Double, Double) -> (Double, Double)
    f d (runningSum, runningCount) = (runningSum + d, runningCount + 1)

-------------------------------------------------------------------------------
-- Exercise 1.2

foldr' :: (a -> b -> b) -> b -> [a] -> b
foldr' _ nil [] = nil
foldr' f nil (a : as) = a `f` foldr' f nil as

foldl' :: (b -> a -> b) -> b -> [a] -> b
foldl' _ nil [] = nil
foldl' f nil (a : as) = foldl' f (nil `f` a) as

foldl'2 :: forall a b. (b -> a -> b) -> b -> [a] -> b
foldl'2 f nil as = foldr' f' nil' as nil
  where
    f' :: a -> (b -> b) -> (b -> b)
    f' a k b = k (f b a)

    nil' :: b -> b
    nil' = id

prop_foldr_is_id :: [Integer] -> Bool
prop_foldr_is_id is = foldr' (:) [] is == is

prop_foldl_is_reverse :: [Integer] -> Bool
prop_foldl_is_reverse is = foldl' (flip (:)) [] is == reverse is

prop_foldl2_is_reverse :: [Integer] -> Bool
prop_foldl2_is_reverse is = foldl'2 (flip (:)) [] is == reverse is

-------------------------------------------------------------------------------
-- Exercise 3
-------------------------------------------------------------------------------
-- BoolTerm definition

data BoolTerm
  = T
  | F
  | Var Char
  | Not BoolTerm
  | Conj BoolTerm BoolTerm
  | Disj BoolTerm BoolTerm
  deriving (Eq, Show, Generic)

type Position = [Integer]

-------------------------------------------------------------------------------
-- Implement the functions here.

pos :: BoolTerm -> [Position]
pos T = [[]]
pos F = [[]]
pos (Var _) = [[]]
pos (Not t) = [] : map (1 :) (pos t)
pos (Conj t u) = [] : map (1 :) (pos t) ++ map (2 :) (pos u)
pos (Disj t u) = [] : map (1 :) (pos t) ++ map (2 :) (pos u)

(|.) :: BoolTerm -> Position -> BoolTerm
t |. [] = t
Not t |. (1 : p) = t |. p
Conj t _ |. (1 : p) = t |. p
Conj _ u |. (2 : p) = u |. p
Disj t _ |. (1 : p) = t |. p
Disj _ u |. (2 : p) = u |. p
t |. p = error ("invalid position " ++ show p ++ " for term " ++ show t)

replace :: BoolTerm -> BoolTerm -> Position -> BoolTerm
replace _ r [] = r
replace (Not t) r (1 : p) = Not (replace t r p)
replace (Conj t u) r (1 : p) = Conj (replace t r p) u
replace (Conj t u) r (2 : p) = Conj t (replace u r p)
replace (Disj t u) r (1 : p) = Disj (replace t r p) u
replace (Disj t u) r (2 : p) = Disj t (replace u r p)
replace t _ p = error ("invalid position " ++ show p ++ " for term " ++ show t)

-------------------------------------------------------------------------------
-- Properties

-- | All terms should contain ε in their set of positions.
prop_pos_empty :: BoolTerm -> Bool
prop_pos_empty t = [] `elem` pos t

-- | Checks that all positions returned from `pos` resolve to valid subterms.
prop_pos_valid :: BoolTerm -> Property
prop_pos_valid t = forAll (elements (pos t)) $ \p ->
  ioProperty $ void $ evaluate (t |. p)

-- | 1. t_pq = (t|p)|q
prop_at_split :: BoolTerm -> Property
prop_at_split t =
  forAll (elements (pos t)) $ \pq ->
    forAll (choose (0, length pq)) $ \n ->
      let (p, q) = splitAt n pq
       in t |. p |. q == t |. pq

-- | 2. (s[t]p)|pq = t|q
prop_replace_split :: BoolTerm -> BoolTerm -> Property
prop_replace_split s t =
  forAll (elements (pos s)) $ \p ->
    forAll (elements (pos t)) $ \q ->
      replace s t p |. (p ++ q) == t |. q

-- | 3. (s[t]p)[r]pq = s[t[r]q]p
prop_replace_twice :: BoolTerm -> BoolTerm -> BoolTerm -> Property
prop_replace_twice s t r =
  forAll (elements (pos s)) $ \p ->
    forAll (elements (pos t)) $ \q ->
      replace (replace s t p) r (p ++ q) == replace s (replace t r q) p

-------------------------------------------------------------------------------
-- Allow quantification over `BoolTerm` values.

instance Arbitrary BoolTerm where
  arbitrary = sized go
    where
      go :: Int -> Gen BoolTerm
      go 0 = oneof [pure T, pure F, Var <$> choose ('a', 'z')]
      go n =
        oneof
          [ Not <$> go (n - 1),
            binary (n - 1) Conj,
            binary (n - 1) Disj
          ]

      binary :: Int -> (BoolTerm -> BoolTerm -> BoolTerm) -> Gen BoolTerm
      binary n f = do
        !a <- invE 0.25 <$> choose (0, 1)
        !b <- invE 0.25 <$> choose (0, 1)
        let rescaled x = go $ round $ x * fromIntegral n / (a + b)
        f <$> rescaled a <*> rescaled b

      invE :: Double -> Double -> Double
      invE lambda u = (-log (1 - u)) / lambda

  shrink = genericShrink

return []

checkAll :: IO Bool
checkAll = $quickCheckAll