{-# LANGUAGE
  DeriveGeneric,
  DerivingVia #-}

import Data.Bifunctor
import GHC.Generics
import Generic.Functor
import System.Exit (exitFailure)

-- Testing DeriveFunctor (gfmap) and gsolomap

data Empty a
  deriving Generic
  deriving Functor via (DeriveFunctor Empty)

data Unit a = Unit
  deriving (Eq, Show, Generic)
  deriving Functor via (DeriveFunctor Unit)

data Result a r = Error a | Ok r
  deriving (Eq, Show, Generic)
  deriving Functor via (DeriveFunctor (Result a))
  
mapError :: (a -> b) -> Result a r -> Result b r
mapError = gsolomap

mapOk :: (r -> s) -> Result a r -> Result a s
mapOk = gsolomap

mapBoth :: (a -> b) -> Result a a -> Result b b
mapBoth = gsolomap

data Writer w a = Writer w a
  deriving (Eq, Show, Generic)
  deriving Functor via (DeriveFunctor (Writer w))

mapW :: (w -> w') -> Writer w a -> Writer w' a
mapW = gsolomap

data Square a b = Square a a b b
  deriving (Eq, Show, Generic)
  deriving Functor via (DeriveFunctor (Square a))

mapFirst :: (a -> a') -> Square a b -> Square a' b
mapFirst = gsolomap

data Twice a = Twice (Either a a)
  deriving (Eq, Show, Generic)
  deriving Functor via (DeriveFunctor Twice)

-- Testing solomap

map1, map1' :: (a -> b) -> Either e (Maybe [(e, a)]) -> Either e (Maybe [(e, b)])
map1 = solomap
map1' = fmap . fmap . fmap . fmap  -- equivalent definition, just making sure it typechecks

map2, map2' :: (a -> b) -> (e -> Either [a] r) -> (e -> Either [b] r)
map2 = solomap
map2' f = fmap (bimap (fmap f) id)

type F a = ([a], Either a ())

map3, map3' :: (a -> b) -> F a -> F b
map3 = solomap
map3' f = bimap (fmap f) (bimap f id)

type G t a = (t, Maybe [Either Bool a])

map4 :: (a -> b) -> G t a -> G t b
map4 = solomap
map4' = fmap . fmap . fmap . fmap

-- Deriving Bifunctor

data Tree a b = Node a (Tree a b) (Tree a b) | Leaf b
  deriving (Eq, Show, Generic)
  deriving Functor via (DeriveFunctor (Tree a))
  deriving Bifunctor via (DeriveBifunctor Tree)

data CofreeF f a b = a :< f b
  deriving (Eq, Show, Generic)
  deriving Bifunctor via (DeriveBifunctor (CofreeF f))

-- Multimap

data Three a b c = One a | Two b | Three c
  deriving (Eq, Show, Generic)

mapThree :: (a -> a') -> (b -> b') -> (c -> c') -> Three a b c -> Three a' b' c'
mapThree f g h = gmultimap (f :+ g :+ h :+ ())

type F5 a b c = Either a (b, c)

map5 :: (a -> a') -> (b -> b') -> (c -> c') -> F5 a b c -> F5 a' b' c'
map5 f g h = multimap (f :+ g :+ h :+ ())
map5' f g h = bimap f (bimap g h)

-- Run at least once

twice :: Int -> Int
twice = (* 2)

main :: IO ()
main = do
  Unit @= fmap twice Unit
  Ok 8 @= fmap twice (Ok 4 :: Result () Int)
  Error 8 @= mapError twice (Error 4 :: Result Int ())
  Writer () 8 @= fmap twice (Writer () 4)
  Writer 8 () @= mapW twice (Writer 4 ())
  Square () () 8 10 @= fmap twice (Square () () 4 5)
  Square 8 10 () () @= mapFirst twice (Square 4 5 () ())
  [Twice (Left 8), Twice (Right 10)] @= (fmap . fmap) twice [Twice (Left 4), Twice (Right 5)]
  Node 8 (Leaf 10) (Leaf 12) @= bimap twice twice (Node 4 (Leaf 5) (Leaf 6))
  (8 :< Just 10) @= bimap twice twice (4 :< Just 5)
  [One 8, Two False, Three 1] @= fmap (mapThree twice not length) [One 4, Two True, Three [()]]

  let t1 = Right (Just [((), 4)])
  map1 twice t1 @= map1' twice t1

  let t2 x = Left [x] :: Either [Int] ()
  map2 twice t2 4 @= map2' twice t2 4

  let t3 = ([4], Left 5)
  map3 twice t3 @= map3' twice t3

  let t4 = ((), Just [Right 4])
  map4 twice t4 @= map4' twice t4

  let t5 = [Left 4, Right (False, [()])]
  fmap (map5 twice not length) t5 @= fmap (map5 twice not length) t5

-- Assert equality
(@=) :: (Eq a, Show a) => a -> a -> IO ()
(@=) x y | x == y = pure ()
         | otherwise = do
  putStrLn "Not equal:"
  print x
  print y
  exitFailure