{-# LANGUAGE Trustworthy, TemplateHaskell, FlexibleContexts, ScopedTypeVariables #-}

-- | Data.THMu: automated generation of instances for data conversion.
module Data.THMu (thMuInstance, thMuDataDec, out'', inn'') where

import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Lib
import Control.Monad hiding (Functor, fmap)
import Data.Char
import Generics.Pointless.Functors
import Generics.Pointless.Combinators
import Generics.Pointless.RecursionPatterns
import Prelude hiding (Functor, fmap)

lower :: String -> String
lower = map toLower

makeNice :: Con -> Con
makeNice (RecC nm types) = NormalC nm (map (\(_, s, t) -> (s, t)) types)
makeNice (InfixC ty1 nm ty2) = NormalC nm [ty1, ty2]
makeNice c@(NormalC _ _) = c

makeNice2 :: Dec -> Dec
makeNice2 (NewtypeD ctx nm tys k con ctx2) = DataD ctx nm tys k [con] ctx2
makeNice2 d@(DataD _ _ _ _ _ _) = d

ctorHead :: Type -> Name
ctorHead (AppT x _) = ctorHead x
ctorHead (ConT x) = x
ctorHead _ = mkName ""

formFunctorType :: Name -> Type -> TypeQ
formFunctorType tName t = if ctorHead t == tName then
		[t| Id |]
	else
		case t of
			AppT t2 t3 -> [t| (:@:) $(return t2) $(formFunctorType tName t3) |]
			_ -> [t| Const $(return t) |]

thMuTypeCase :: Name -> Con -> TypeQ
thMuTypeCase tName con = if null types then
		[t| Const () |]
	else
		foldr1 (\t t2 -> [t| (:*:) $(t) $(t2) |]) (map (formFunctorType tName . snd) types) where
	NormalC nm types = makeNice con

thMuType :: Name -> [Con] -> TypeQ
thMuType tName cons = foldr1 (\t t2 -> [t| (:+:) $(t) $(t2) |])
	(map (thMuTypeCase tName) cons)

thEithers :: Int -> Int -> ExpQ -> ExpQ
thEithers n total exp = foldr ($) exp' (replicate n (\x -> [| Right $(x) |])) where
	exp' = if n == total - 1 then
			exp
		else
			[| Left $(exp) |]

thPairs :: [ExpQ] -> ExpQ
thPairs [] = [| () |]
thPairs ls = foldr1 (\e1 e2 -> [| (,) $(e1) $(e2) |]) ls

name :: ExpQ -> Q Name
name e = liftM (\(VarE nm) -> nm) e

nameT :: TypeQ -> Q Name
nameT t = liftM (\(ConT nm) -> nm) t

thRep :: [Con] -> DecsQ
thRep cons = liftM (:[]) $ name [| out |] >>= \outN -> funD outN
	(zipWith (\con conI -> let

		NormalC nm types = makeNice con

		vars = zipWith (\n _ -> mkName ("x" ++ show n)) [0..] types

		pattern = conP nm (map varP vars)

		body = liftM NormalB (thEithers conI (length cons) (thPairs (map varE vars))) in

		liftM2 (\x1 x2 -> Clause [x1] x2 [])
			pattern
			body)
		cons
		[0..])

thEithersPat :: Int -> Int -> PatQ -> PatQ
thEithersPat n total exp = foldr ($) exp' (replicate n (\x -> name [| Right |] >>= \n -> conP n [x])) where
	exp' = if n == total - 1 then
			exp
		else
			 name [| Left |] >>= \n -> conP n [exp]

thPairsPat :: [PatQ] -> PatQ
thPairsPat [] = [p| () |]
thPairsPat ls = foldr1 (\e1 e2 -> tupP [e1,e2]) ls

thUnrep :: [Con] -> DecsQ
thUnrep cons = liftM (:[]) $ name [| inn |] >>= \inN -> funD inN
	(zipWith (\con conI -> let

		NormalC nm types = makeNice con

		vars = zipWith (\n _ -> mkName ("x" ++ show n)) [0..] types

		pattern = thEithersPat conI (length cons) (thPairsPat (map varP vars))

		body = liftM NormalB (foldl (\x -> appE x . varE) (conE nm) vars) in

		liftM2 (\x1 x2 -> Clause [x1] x2 [])
			pattern
			body)
		cons
		[0..])

tyvarBinderToType :: TyVarBndr -> TypeQ
tyvarBinderToType (PlainTV nm) = varT nm
tyvarBinderToType (KindedTV nm kind) = sigT (varT nm) kind

liftM6 f m m2 m3 m4 m5 m6 = do
	x1 <- m
	liftM5 (f x1) m2 m3 m4 m5 m6

-- | Defines a PF instance, and a Mu instance.
thMuInstance :: Dec -> DecsQ
thMuInstance dec = let DataD _ nm tyvars _ con _ = makeNice2 dec in
	liftM6 (\x1 x2 x3 x4 x5 x6 -> [TySynInstD x1 (TySynEqn [ConT nm] x3),InstanceD Nothing [] (AppT x2 x4) (x5 ++ x6)])
		(nameT [t| PF |])
		[t| Mu |]
		(thMuType nm con)
		(foldl (\x -> appT x . tyvarBinderToType) (conT nm) tyvars)
		(thRep con)
		(thUnrep con)

-- | Accepts a data type declaration in a splice, and defines the data type
--   along with calling 'thMuInstance'.
thMuDataDec :: DecsQ -> DecsQ
thMuDataDec q = q >>= \[dec] -> liftM (dec:) (thMuInstance dec)

-- | Conveniences to convert to/from sums of products.
out'' :: forall t. (Mu t, Functor (PF t)) => t -> Fix (PF t)
out'' = nu (_L :: Ann t)

inn'' :: forall t. (Mu t, Functor (PF t)) => Fix (PF t) -> t
inn'' = nu (_L :: Ann t)