module MkZip where
	import Language.Haskell.TH
	import Language.Haskell.TH.Syntax

	{-
	 - genPE:
	 -
	 - Returns a double containing a list of patterns and expressions in the form: s1, s2, ... sn
	 - If the arguments were "n" 5 then the first element of the double would be a list of patterns
	 - [n1, n2, n3, n4, n5] and the second element would be the corresponding list of expressions.
	 -}

	genPE s n = 
		let
			ns = [ s++(show i) | i <- [1..n] ]
		in
			((map (varP . mkName) ns), (map (varE . mkName) ns))

	{-
	 - copies:
	 -
	 - Creates a list of n copies of something, for examples, copies 5 5 would return [5, 5, 5, 5, 5]
	 -}
	
	copies 0 cping = []
	copies n cping = cping:(copies (n - 1) cping)
	
	{-
	 - zipN:
	 - 
	 - $(zipN 3) will generate the following code (although the function name may be different):
	 -
	 - newZip_0 = 
	 -   \y1 y2 y3 ->
	 -     case (y1, y2, y3) of
	 -       ((x1:xs1), (x2:xs2), (x3:xs3)) -> (x1, x2, x3):(newZip_0 xs1 xs2 xs3)
	 -       (_, _, _) -> []
	 -}
	
	zipN n = [| let newZip = $(mkZip n [| newZip |]) in newZip |]

	{-
	 - unZipN:
	 - 
	 - $(unzipN 3) will generate the following code (although the function name may be different):
	 -
	 - newUnzip_0 = 
	 -   \lst -> 
	 -     case lst of
	 -       ((x1, x2, x3):ts) ->
	 -         let (xs1, xs2, xs3) = newUnzip_0 ts
	 -         in (x1:xs1, x2:xs2, x3:xs3)
	 -       [] -> ([], [], [])
	 -}

	unzipN n = [| let newUnzip = $(mkUnzip n [| newUnzip |]) in newUnzip |]

	{-
	 - mkZip:
	 -
	 - Does the actual generation of the function for zipN.
	 -}

	mkZip n name = lamE pYs (caseE (tupE eYs) [m1, m2])
		where
			(pXs, eXs) = genPE "x" n
			(pYs, eYs) = genPE "y" n
			(pXSs, eXSs) = genPE "xs" n
			pcons x xs = infixP x '(:) xs -- Creates the pattern of x:xs
			b = [| $(tupE eXs) : $(appsE(name : eXSs)) |] -- Creates the syntax tree for something like (x1, x2):(newZip_x xs1 xs2) depending on the size of the size function and the function name given.
			m1 = match (tupP (zipWith pcons pXs pXSs)) (normalB b) [] -- The match with the form ((x1:xs1), (x2, xs2)) -> (x1, x2):(newZip_x xs1 xs2)
			m2 = match (tupP (copies n wildP)) (normalB ([| [] |])) [] -- The match with the form (_, _) -> []

	{-
	 - mkUnzip:
	 -
	 - Does the actual generation of the function for unzipN.
	 -}

	mkUnzip n name = lamE [varP nList] (caseE (varE nList) [m1, m2])
		where
			nList = mkName "lst"
			nTs = mkName "ts"
			(pXs, eXs) = genPE "x" n
			(pXSs, eXSs) = genPE "xs" n
			econs x xs = [| $x : $xs |] -- Creates the expression of x:xs
			pcons x xs = infixP x '(:) xs -- Creates the pattern of x:xs
			b = appE name (varE nTs) -- Creates the expression: newUnzip_x ts
			eLet = letE ([valD (tupP pXSs) (normalB b) []]) (tupE (zipWith econs eXs eXSs)) -- A let with the form let (xs1, xs2) = newUnzip_x ts in (x1:xs1, x2:xs2)
			m1 = match (pcons (tupP pXs) (varP nTs)) (normalB eLet) [] -- Match with the form (x1, x2):ts -> (eLet)
			m2 = match (varP (mkName "[]")) (normalB(tupE (copies n [| [] |]))) [] -- Match with the form [] -> ([], [], [])