{-# LANGUAGE TemplateHaskell, FlexibleContexts, TypeFamilies, RankNTypes #-}


-- |
--
-- Module Data.Fieldwise provides 'Fieldwise' typeclass for operations
-- of fields of records treated as independent components.
module Data.Fieldwise
  ( Fieldwise(..)
  , deriveFieldwise
  )
where

import Language.Haskell.TH
import Data.Monoid
import Control.Applicative

-- | Fieldwise class provides some operations for record fields
-- treated as independent pieces of data. See individual methods.
class Fieldwise r where
  -- | UnliftedRecord points to original data structures.
  --
  -- For example:
  --
  -- > data T = T Int String
  -- > data T_f f = T (f Int) (f String)
  --
  -- Then @r = T_f@ and @UnliftedRecord r = T@

  type UnliftedRecord r
  -- | This function is similar to 'sequenceA', it composes fields
  -- using '<*>' operator.
  --
  -- Example:
  --
  -- > sequenceR (T_f (Just 1) (Just "abc")) = Just (T 1 "abc")
  sequenceR :: Applicative f => r f -> f (UnliftedRecord r)

  -- | This function can replace field wrapper.
  --
  -- Example:
  --
  -- > hoistR maybeToList (T_f (Just 1) (Just "abc")) = T_f [1] ["abc"]
  hoistR    :: (forall a . f a -> g a) -> r f -> r g

  -- | This function zips respective field values uzing zipper
  -- function,
  --
  -- Example:
  --
  -- > zipWithR mappend (T_f (Just 1) Nothing) (T_f Nothing (Just "c") =
  -- >                        T_f (Just 1) (Just "c")
  zipWithR  :: (forall a . f a -> g a -> h a) -> r f -> r g -> r h

  -- | This function uses a function to compose fields of a lifted
  -- record with selectors of original record.
  --
  -- Note that this is so complicated to support records with many
  -- constructors.
  appR      :: (forall a . f a -> (UnliftedRecord r -> a) -> g a) -> r f -> r g

  -- | Wrap each value in a record in 'pure'.
  --
  -- Example:
  --
  -- > liftR (T 123 "abc") = T_f (pure 123) (pure "abc")
  liftR     :: Applicative f => UnliftedRecord r -> r f

-- | Automatically derive lifted record type and 'Monoid' and
-- 'Fieldwise' instances for it.
--
-- First argument is the name of Haskell data type that should serve
-- as basis for derivation, second argument tell how to tranform names
-- in that type. Names need to be transformed if you want to derive
-- fieldwise in the same module as original data type.
--
-- Conceptually for a data type @T@ a derived data @T_f@ has type for
-- each field wrapped in a type constructor. @T@ is semantically equal
-- to @T_f Id@.
--
-- For example for data type:
--
-- > data Test1 = Test1 Int String
-- >            | Test2 { test2Char :: Char, test2IntList :: [Int], test2Func :: (Int -> Int) }
--
-- > $(deriveFieldwise ''Test1  (++ "_f"))
--
-- Will produce the following splice:
--
-- > data Test1_f f
-- >   = Test1_f (f Int) (f String) |
-- >     Test2_f {test2Char_f :: f Char,
-- >              test2IntList_f :: f [Int],
-- >              test2Func_f :: f (Int -> Int)}
-- >
-- > instance Alternative f => Monoid (Test1_f f) where
-- >   mempty = Test1_f empty empty
-- >   mappend (Test1_f l1 l2) (Test1_f r1 r2)
-- >     = Test1_f (l1 <|> r1) (l2 <|> r2)
-- >   mappend (Test2_f l1 l2 l3) (Test2_f r1 r2 r3)
-- >     = Test2_f (l1 <|> r1) (l2 <|> r2) (l3 <|> r3)
-- >
-- > instance Fieldwise Test1_f where
-- >   type instance UnliftedRecord Test1_f = Test1
-- >   sequenceR (Test1_f l1 l2) = (((pure Test1) <*> l1) <*> l2)
-- >   hoistR fg (Test1_f l1 l2) = Test1_f (fg l1) (fg l2)
-- >   hoistR fg (Test2_f l1 l2 l3)
-- >     = Test2_f (fg l1) (fg l2) (fg l3)
-- >   zipWithR fg (Test1_f l1 l2) (Test1_f r1 r2)
-- >     = Test1_f (fg l1 r1) (fg l2 r2)
-- >   zipWithR fg (Test2_f l1 l2 l3) (Test2_f r1 r2 r3)
-- >     = Test2_f (fg l1 r1) (fg l2 r2) (fg l3 r3)
-- >   appR fg (Test1_f l1 l2)
-- >     = Test1_f
-- >         (fg l1 (\ (Test1 q_ahwb _) -> q_ahwb))
-- >         (fg l2 (\ (Test1 _ q_ahwb) -> q_ahwb))
-- >   appR fg (Test2_f l1 l2 l3)
-- >     = Test2_f
-- >         (fg l1 test2Char)
-- >         (fg l2 test2IntList)
-- >         (fg l3 test2Func)
-- >   liftR (Test1 l1 l2) = Test1_f (pure l1) (pure l2)
-- >   liftR (Test2 l1 l2 l3) = Test2_f (pure l1) (pure l2) (pure l3)

deriveFieldwise :: Name -> (String -> String) -> Q [Dec]
deriveFieldwise t upName = do
  info <- reify t
  case info of
    TyConI (DataD ctx _name vars tcons@(firsttcon:_) _derivings) -> do
      let dname = mkName (upName (nameBase t))
      f <- newName "f"
      fg <- newName "fg"
      q <- newName "q"
      let upCon (NormalC nm stypes) = NormalC (mkName (upName (nameBase nm))) (map upStrictType stypes)
          upCon (RecC nm vstypes) = RecC (mkName (upName (nameBase nm))) (map upVarStrictType vstypes)
          upCon (InfixC stype1 nm stype2) = InfixC (upStrictType stype1) (mkName (upName (nameBase nm))) (upStrictType stype2)
          upCon (ForallC favars ctx2 con2) = ForallC favars ctx2 (upCon con2)
          upStrictType (strict, tp) = (strict, AppT (VarT f) tp)
          upVarStrictType (nm, strict, tp) = (mkName (upName (nameBase nm)), strict, AppT (VarT f) tp)
          getName (NormalC nm _) = nm
          getName (RecC nm _) = nm
          getName (InfixC _ nm _) = nm
          getName (ForallC _ _ c) = getName c
          getArity (NormalC _nm stypes) = length stypes
          getArity (RecC _nm vstypes) = length vstypes
          getArity (InfixC stype1 _nm stype2) = 2
          getArity (ForallC favars ctx2 con2) = getArity con2
          selectorNth i (RecC _nm vstypes) = VarE (fst3 (vstypes!!(i-1)))
          selectorNth i (NormalC nm stypes) = LamE [ConP nm [ if idx==i then VarP q else WildP | (_,idx) <- zip stypes [1..] ]]
                                                    (VarE q)
          selectorNth 1 (InfixC _ nm _) = LamE [InfixP (VarP q) nm WildP] (VarE q)
          selectorNth 2 (InfixC _ nm _) = LamE [InfixP WildP nm (VarP q)] (VarE q)
          selectorNth i (ForallC _ _ c) = selectorNth i c
          fst3 (a,_,_) = a
      return
        [ DataD ctx dname (PlainTV f : vars) (map upCon tcons) []
        , InstanceD ([ClassP ''Alternative [VarT f]])
                        (AppT (ConT ''Monoid) (ConT dname `AppT` VarT f))
          [
            FunD 'mempty [Clause [] (NormalB (foldl AppE (ConE (getName (upCon firsttcon)))
                    (take (getArity firsttcon) (repeat (VarE 'empty))))) []]

          , FunD 'mappend [Clause [ConP (getName (upCon tc)) [VarP (mkName ("l" ++ show i)) | i <- [1..getArity tc]],
                                   ConP (getName (upCon tc)) [VarP (mkName ("r" ++ show i)) | i <- [1..getArity tc]]]
                      (NormalB (foldl AppE (ConE (getName (upCon tc)))
                          [InfixE (Just (VarE (mkName ("l" ++ show i)))) (VarE '(<|>))  (Just (VarE (mkName ("r" ++ show i)))) |
                           i <- [1..getArity tc]])) []
                          | tc <- tcons
                          ]
          ]
        , InstanceD [] (ConT ''Fieldwise `AppT` ConT dname)
          [ TySynInstD ''UnliftedRecord [ConT dname] (ConT t)
          , FunD 'sequenceR [Clause [ConP (getName (upCon firsttcon)) [VarP (mkName ("l" ++ show i)) | i <- [1..getArity firsttcon]]]
                      (NormalB (foldl (\l r -> InfixE (Just l) (VarE '(<*>)) (Just r)) (VarE 'pure `AppE` ConE (getName firsttcon))
                          [VarE (mkName ("l" ++ show i)) | i <- [1..getArity firsttcon]])) []]
          , FunD 'hoistR [Clause [VarP fg, ConP (getName (upCon tc)) [VarP (mkName ("l" ++ show i)) | i <- [1..getArity tc]]]
                      (NormalB (foldl AppE (ConE (getName (upCon tc)))
                          [VarE fg `AppE` VarE (mkName ("l" ++ show i)) | i <- [1..getArity tc]])) []
                         | tc <- tcons ]
          , FunD 'zipWithR [Clause [VarP fg,
                                    ConP (getName (upCon tc)) [VarP (mkName ("l" ++ show i)) | i <- [1..getArity tc]],
                                    ConP (getName (upCon tc)) [VarP (mkName ("r" ++ show i)) | i <- [1..getArity tc]]]
                      (NormalB (foldl AppE (ConE (getName (upCon tc)))
                          [VarE fg `AppE` VarE (mkName ("l" ++ show i)) `AppE` VarE (mkName ("r" ++ show i)) |
                           i <- [1..getArity tc]])) []
                           | tc <- tcons ]
          , FunD 'appR [Clause [VarP fg,
                                    ConP (getName (upCon tc)) [VarP (mkName ("l" ++ show i)) | i <- [1..getArity tc]]]
                      (NormalB (foldl AppE (ConE (getName (upCon tc)))
                          [VarE fg `AppE` VarE (mkName ("l" ++ show i)) `AppE` selectorNth i tc |
                           i <- [1..getArity tc]])) []
                           | tc <- tcons ]
          , FunD 'liftR [Clause [ConP (getName tc) [VarP (mkName ("l" ++ show i)) | i <- [1..getArity tc]]]
                      (NormalB (foldl AppE (ConE (getName (upCon tc)))
                          [VarE 'pure `AppE` VarE (mkName ("l" ++ show i)) |
                           i <- [1..getArity tc]])) []
                           | tc <- tcons ]
          ]
        ]
    _ -> error $ "Not a data with single constructor declaration: " ++ show info