```{-# OPTIONS_GHC -fno-warn-missing-methods #-}

module Data.DeriveGuess(DataName(..), tup1, guess) where

import Data.Generics
import Data.List
import Data.Char
import Data.Maybe

data DataName a = CtorZero
| CtorOne  a
| CtorTwo  a a
| CtorTwo' a a

ctorNames = ["CtorZero","CtorOne","CtorTwo","CtorTwo'"]

-- | Guess an instantiator from a sample instance.
guess :: (String, Q [Dec]) -> IO ()
guess (name,x) = runQ x >>= putStr . (++) line0. widthify . (++) line1 . guessStr . unQ
where
line0 = "make" ++ name ++ " :: Derivation\n" ++
"make" ++ name ++ " = derivation " ++ lname ++ "' \"" ++ name ++ "\"\n"
line1 = lname ++ "' dat = "
lname = toLower (head name) : tail name

-- | A fake constructor for the unary tuple.  Helps 'guess' to see
-- patterns in progressions of differently sized tuples.
tup1 = id

-- | Chop and mangle a String representing Haskell code so that it
-- fits in 80 columns, without regard for prettiness.
widthify :: String -> String
widthify xs = g 80 (f xs)
where
g n (x:xs) | n - length x <= 0 = "\n    " ++ g 76 ([x|x/=" "] ++ xs)
| otherwise = x ++ g (n - length x) xs
g n [] = "\n"

f (x:xs) | isSpace x = " " : f (dropWhile isSpace xs)
f x = case lex x of
[("","")] -> []
-- \\ must not occur at the end of a line (CPP restrictions)
[("\\",y)] -> let a:b = f y in ('\\':a) : b
[(x,y)] -> x : f y

-- | Process a tree produced by a quasiquote, stripping name
-- uniquifiers and changing applications and tuplings into a standard
-- form.
unQ :: [Dec] -> [Dec]
unQ x = everywhere (mkT g) \$ everywhere fAny \$ map normData x
where
fAny :: (Typeable a, Data a) => a -> a
fAny = mkT fE `extT` fP
fE (VarE x) = VarE (f x); fE x = x
fP (VarP x) = VarP (f x); fP x = x

-- | Remove_0 evil_1 ghc_2 name_3 uniquifiers_4
f :: Name -> Name
f name = if not ("_" `isSuffixOf` s) && match s then mkName \$ dropUnder s else name
where
s = show name
match = isPrefixOf "_" . dropWhile isDigit . reverse

-- | Turn infix applications into prefix, and normalise
-- tuples.
g :: Exp -> Exp
g (InfixE (Just x) y (Just z)) = AppE (AppE y x) z
g (AppE (VarE tup) x) | show tup == "tup1" = TupE [x]
g (ConE unit) | show unit == "()" = TupE []
g x = x

-- | Drop the first _ and everything after it; used to trim GHC name
-- uniques.
dropUnder = reverse . drop 1 . dropWhile (/= '_') . reverse

list x = "[" ++ concat (intersperse "," x) ++ "]"

unwordsb x = "(" ++ unwords x ++ ")"

fst3 (a,b,c) = a
snd3 (a,b,c) = b
thd3 (a,b,c) = c

snub x = nub \$ sort x

arityToCtors x = x : [3 | x == 2]
ctorToArity  x = if x == 3 then 2 else x

ctorArityEnv (Ctor i) = ctorToArity i

on op get a b = op (get a) (get b)

-- imagine the following environment table:
{-
[("CtorZero",0,0,[])
,("CtorOne" ,1,1,[1])
,("CtorTwo" ,2,2,[1,2])
,("CtorTwo'",2,3,[1,2])
]
-}

data Env = None | Ctor Int | Field Int
deriving (Show,Eq)

isField (Field _) = True; isField _ = False
isCtor  (Ctor  _) = True; isCtor  _ = False

fromField (Field i) = i
fromCtor  (Ctor  i) = i

fromEnv (Field i) = i
fromEnv (Ctor  i) = i

-- Show t and Ppr t for better error messages
-- Eq t required for hypothesis testing
class (Ppr t, Eq t, Show t) => Guess t where
-- invariant: all answers must be correct for this example
-- will never be given a different type of environment
guessEnv :: t -> [(Env, Env -> t, String)]

guessStr :: Guess t => t -> String
guessStr t = case [s | (None,_,s) <- guessEnv t] of
[] -> error \$ "\n\nNo hypothesis for:\n" ++ show t ++
"\n\nPretty version:\n" ++ show (ppr t)
(x:xs) -> x

checkGuess :: (Ppr t, Eq t, Show t) => t -> [(Env, Env -> t, String)] -> [(Env, Env -> t, String)]
checkGuess t xs = map f xs
where
f o@(env,gen,str) | t == gen env = o
| otherwise = error \$ unlines ["checkGuess failed:"
,"INPUT : " ++ show t
,"OUTPUT: " ++ show (gen env)
,"ENV   : " ++ show env
,"HYP   : " ++ str
]

guessEnvStr :: Guess t => t -> [(Env, Env -> t, String)]
guessEnvStr t = [(None, const t, guessStr t)]

guessPairStr :: (Guess a, Guess b) => String -> a -> b -> String
guessPairStr sjoin a b = sjoin ++ " " ++ guessStr a ++ " " ++ guessStr b

guessTripStr :: (Guess a, Guess b, Guess c) => String -> a -> b -> c -> String
guessTripStr sjoin a b c = unwords [sjoin, guessStr a, guessStr b, guessStr c]

joinEnvs :: [Env] -> Maybe Env
joinEnvs xs = if length ys > 1 then Nothing else Just \$ head (ys ++ [None])
where ys = filter (/= None) \$ nub xs

guessOneEnv :: Guess a => (a -> t) -> String -> a -> [(Env, Env -> t, String)]
guessOneEnv fjoin sjoin x1 =
[ (e1, \e -> fjoin (f1 e), unwordsb [sjoin,s1])
| (e1,f1,s1) <- guessEnv x1]

-- to join two elements either they are the same env, or one has None
guessPairEnv :: (Guess a, Guess b) => (a -> b -> t) -> String -> a -> b -> [(Env, Env -> t, String)]
guessPairEnv fjoin sjoin x1 x2 =
[ (env, \e -> fjoin (f1 e) (f2 e), unwordsb [sjoin,s1,s2])
| (e1,f1,s1) <- guessEnv x1
, (e2,f2,s2) <- guessEnv x2
, Just env <- [joinEnvs [e1,e2]]]

guessTripEnv :: (Guess a, Guess b, Guess c) => (a -> b -> c -> t) -> String -> a -> b -> c -> [(Env, Env -> t, String)]
guessTripEnv fjoin sjoin x1 x2 x3 =
[ (env, \e -> fjoin (f1 e) (f2 e) (f3 e), unwordsb [sjoin,s1,s2,s3])
| (e1,f1,s1) <- guessEnv x1
, (e2,f2,s2) <- guessEnv x2
, (e3,f3,s3) <- guessEnv x3
, Just env <- [joinEnvs [e1,e2,e3]]]

instance Guess a => Guess [a] where
guessEnv os = concatMap f \$ mapM guessEnv os
where
-- first try and induct based on the length of the list
f xs | all (== None) (map fst3 xs) &&
length xs == 2 &&
length vals == 1
= [(Ctor i, \e -> replicate (ctorArityEnv e) (head vals),
"(replicate (ctorArity ctor) " ++ thd3 (head xs) ++ ")")
| i <- [2,3]] ++
[(None, \e -> map (\$ e) gens, list strs)]
where
(envs,gens,strs) = unzip3 xs
vals = nub \$ zipWith (\$) gens envs

f xs | length es <= 1 = [(head (es ++ [None]), \e -> map (\$ e) gens, list strs)]
| otherwise = [(env,gen,"("++str++")")
| env <- newEnvs, (gen,str) <- nubBy ((==) `on` snd) \$ g xs]
where
(envs,gens,strs) = unzip3 xs
es = nub \$ filter (/= None) envs

ctors = snub [i | Ctor i <- envs]
fields = snub [i | Field i <- envs]
maxField = maximum fields

newEnvs = case ctors of
[] -> map Ctor \$ arityToCtors maxField
_ | null fields -> [None]
[x] | ctorToArity x == maxField -> [Ctor x]
_ -> []

ctorEnv = head newEnvs == None
varName = if ctorEnv then "(ctorInd,ctor)" else "field"

domain = if ctorEnv then [0..3] else [1..maxField]
getDomain (Ctor i) = take 2 [1..i]
getDomain None = [0..3]
getDomain _ = []
strDomain = if ctorEnv then "(zip [0..] (dataCtors dat))" else "[1..ctorArity ctor]"

construct = if ctorEnv then Ctor else Field

isNone x = x == None || (not ctorEnv && isCtor x)

g :: Eq t => [(Env, Env -> t, String)] -> [(Env -> [t], String)]
g [] = [(\e -> [], "[]")]
g ((none,gn,st):xs) | isNone none =
[(\e -> gn e : gen e, "[" ++ st ++ "]++" ++ str) | (gen,str) <- g xs]

g xs =  h id "id" xs ++ h reverse "reverse" xs

h :: Eq t => ([Int] -> [Int]) -> String -> [(Env, Env -> t, String)] -> [(Env -> [t], String)]
h fdir sdir xs
| map construct (fdir domain) `isPrefixOf` map fst3 xs
= [(\e -> map (fhyp . construct) (fdir \$ getDomain e) ++ gen e
,"(map (\\" ++ varName ++ " -> " ++ shyp ++ ") (" ++ sdir ++ " " ++ strDomain ++ "))++" ++ str)
| (fhyp,shyp) <- validHyp
, (gen,str) <- g rest]
where
(now,rest) = splitAt (length domain) xs

validHyp = filter (\hyp -> all (valid (fst hyp)) now) (map (\(a,b,c) -> (b,c)) now)
valid hyp (e,gen,_) = hyp e == gen e

h _ _ _ = []

guessType :: Type -> Type -> String
guessType principle x =
if hasPrinciple then "(map (\\tdat -> " ++ disp x ++ ") (dataVars dat))" else "[" ++ disp x ++ "]"
where
hasPrinciple = f x
where
f x | x == principle = True
f (AppT (ConT x) y) | show x == "DataName" = False
f (AppT a b) = f a || f b
f _ = False

disp x | x == principle = "tdat"
disp (AppT (ConT x) y) | show x == "DataName" = "(lK (dataName dat) (dataVars dat))"
disp (AppT a b) = "(AppT " ++ disp a ++ " " ++ disp b ++ ")"
disp (VarT a) = "(VarT (mkName " ++ show (show a) ++ "))"
disp (ConT a) = "(ConT (mkName " ++ show (show a) ++ "))"

instance Guess Dec where
guessEnv (InstanceD ctx typ inner) =
[ (None, \e -> InstanceD ctx typ (gen e), prefix ++ str)
| (None,gen,str) <- guessEnv inner]
where
principle = head (everything (++) ([] `mkQ` f) typ ++ [VarT \$ mkName "?"])
where
f (AppT (ConT x) y) | show x == "DataName" = [y]
f _ = []

prefix = "InstanceD " ++
"(concat (" ++ list (map (guessType principle) ctx) ++ ")) " ++
"(head " ++ guessType principle typ ++ ")"

guessEnv (FunD name claus) = guessPairEnv FunD "FunD" name claus
guessEnv (ValD pat bod whr) = guessTripEnv ValD "ValD" pat bod whr

guessEnv x = error \$ show ("Guess Dec",x)

instance Guess Name where
guessEnv name = if null guessCtor then guessRest else guessCtor
where
sname = show name
(pre,end) = (init sname, last sname)

guessCtor = [(Ctor i, \(Ctor e) -> mkName (pre ++ (ctorNames !! e))
,"(mkName (" ++ show pre ++ " ++ ctorName ctor))")
| (i,nam) <- zip [0..] ctorNames, nam `isSuffixOf` sname
, let pre = take (length sname - length nam) sname]

guessRest = guessLast ++ guessDefault

guessLast | isDigit end = [(e, \e -> mkName \$ pre ++ show (g e)
,"(mkName (" ++ show pre ++ " ++ show " ++ s ++ "))")
| (e,g,s) <- guessNum \$ read [end]]
| otherwise   = []

guessDefault = [(None,const name, "(mkName " ++ show sname ++ ")")
| not (isDigit end) || pre `notElem` ["x","y","z"]]

guessNum :: Int -> [(Env, Env -> Int, String)]
guessNum i = [(Field i, fromField, "field") | i `elem` [1,2]] ++
[(None, const 3, "(toInteger (length (dataCtors dat) - 1))") | i == 3] ++
[(None, const 4, "(toInteger (length (dataCtors dat)))") | i == 4] ++
[(Ctor i, fromCtor, "ctorInd") | i `elem` [0..3]] ++
[(Ctor i, ctorArityEnv, "(ctorArity ctor)") | i `elem` [0..2]] ++
[(Ctor 3, ctorArityEnv, "(ctorArity ctor)") | i == 2]

instance Guess Clause where
guessEnv (Clause pat bod whr) = guessTripEnv Clause "Clause" pat bod whr

instance Guess Stmt where
guessEnv (BindS x y) = guessPairEnv BindS "BindS" x y
guessEnv (NoBindS x) = guessOneEnv NoBindS "NoBindS" x
guessEnv x = error \$ show ("Guess Stmt",x)

instance Guess Pat where
guessEnv (VarP x) = guessOneEnv VarP "VarP" x
guessEnv (ConP x xs) = guessPairEnv ConP "ConP" x xs
guessEnv (WildP) = [(None, const WildP, "WildP")]
guessEnv (TildeP x) = guessOneEnv TildeP "TildeP" x
guessEnv (RecP x []) = guessOneEnv (flip RecP []) "(flip RecP [])" x
guessEnv (LitP x) = guessOneEnv LitP "LitP" x
guessEnv x = error \$ show ("Guess Pat",x)

instance Guess Body where
guessEnv (NormalB x) = guessOneEnv NormalB "NormalB" x
guessEnv x = error \$ show ("Guess Body",x)

instance Guess Exp where
guessEnv (VarE x) = guessOneEnv VarE "VarE" x
guessEnv (ConE x) = guessOneEnv ConE "ConE" x
guessEnv (LitE x) = guessOneEnv LitE "LitE" x
guessEnv (ListE x) = guessOneEnv ListE "ListE" x
guessEnv (LamE x y) = guessPairEnv LamE "LamE" x y
guessEnv (CompE x) = guessOneEnv CompE "CompE" x
guessEnv (CaseE x y) = guessPairEnv CaseE "CaseE" x y
guessEnv (TupE x) = guessOneEnv TupE "TupE" x
guessEnv (RecConE x []) = guessOneEnv (flip RecConE []) "(flip RecConE [])" x
guessEnv (CondE x y z) = guessTripEnv CondE "CondE" x y z
guessEnv (DoE x) = guessOneEnv DoE "DoE" x

guessEnv o@(AppE x y) = guessApply o ++ guessFold o ++ guessPairEnv AppE "AppE" x y

guessEnv x = error \$ show ("Guess Exp",x)

instance Guess Match where
guessEnv (Match a b c) = guessTripEnv Match "Match" a b c

instance Guess Lit where
guessEnv o@(IntegerL i) =
[ (env, \e -> IntegerL \$ toInteger \$ gen e, "(IntegerL " ++ str ++ ")")
| (env,gen,str) <- guessNum \$ fromInteger i] ++
[(None,const \$ IntegerL i,"(IntegerL " ++ show i ++ ")")]

guessEnv o@(StringL s) | s == "DataName" = [(None, const o, "(StringL (dataName dat))")]
| otherwise = [(None, const o, "(StringL " ++ show s ++ ")")]

guessEnv x = error \$ show ("Guess Lit",x)

-- for when an expression is just an application
guessApply :: Exp -> [(Env, Env -> Exp, String)]
guessApply o | length args <= 1 = []
| otherwise = guessPairEnv applyWith "applyWith" fn args
where
(fn,args) = list o

list (AppE x y) = let (fn,args) = list x in (fn, args ++ [y])
list x = (x, [])

-- for when an expression comes from folding
guessFold :: Exp -> [(Env, Env -> Exp, String)]
guessFold o@(AppE (AppE fn x) y) =
f (with foldl1With) "foldl1With" (list True  o) ++
f (with foldr1With) "foldr1With" (list False o)
where
with fold join [] = VarE \$ mkName "?"
with fold join xs = fold join xs

list b (AppE (AppE fn2 x) y) | fn == fn2 =
if b then x : list b y else y : list b x
list b x = [x]

f ffold sfold lst
| length lst <= 2 = []
| otherwise = guessPairEnv ffold sfold fn lst

guessFold _ = []

```