{-# LANGUAGE TemplateHaskell, ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
module Data.Cursor.CLASE.Gen.Language(languageGen) where

import Control.Arrow
import Control.Monad
import Data.List
import Data.Map (Map)
import Data.Ord
import Data.Function
import Data.Maybe
import Data.Set (Set)
import Data.Cursor.CLASE.Gen.Util
import Data.Cursor.CLASE.Gen.PrintM
import System.FilePath
import Language.Haskell.TH
import qualified Data.Map as Map
import qualified Data.Set as Set
import System.IO

languageGen :: [String] -> Name -> [Name] -> Q [Dec]
languageGen moduleNames rootName acceptableNames = do
  nameMap <- buildMap acceptableNames
  
  fileOut <- runIO $ openFile (joinPath moduleNames <.> "hs") WriteMode

  let rootNameModule = nameModule rootName
  let moduleName = concat . intersperse "." $ moduleNames

  runIO . runPrint fileOut $ do
    printHeader moduleName rootNameModule (nameBase rootName)
    printLn ""

    printTypeRep nameMap
    printLn ""

    printReify rootName nameMap
    printLn ""

    let contextCtrs = buildContextCtrs nameMap

    printContext contextCtrs
    printLn ""

    printContextToMovement contextCtrs
    printLn ""

    printBuildOne contextCtrs
    printLn ""

    let leftRightNames = buildLeftRightNames nameMap
    printMovement contextCtrs leftRightNames
    printLn ""

    printInvertMovement contextCtrs leftRightNames
    printLn ""

    printMovementEq contextCtrs leftRightNames
    printLn ""

    printUnbuildOne contextCtrs leftRightNames
    printLn ""

    printReifyDirection contextCtrs
    printLn ""

    printDownMovements rootName contextCtrs
    printLn ""

    printGenericMoveLeft contextCtrs nameMap rootName
    printLn ""

    printGenericMoveRight contextCtrs nameMap rootName
    printLn ""

{-
    printToStringMovement contextCtrs leftRightNames
    printLn ""

    printContextReps contextCtrs
    printLn ""

    printMoveToRoot rootName
    printLn ""
-}

  runIO $ hFlush fileOut
  runIO $ hClose fileOut

  return []

printHeader :: String -> Maybe String -> String -> PrintM ()
printHeader modName mRootModName rootName = do
    printLn . unlines $ 
     [ "{-# LANGUAGE GADTs, MultiParamTypeClasses, TypeFamilies, TypeOperators, ScopedTypeVariables, ExistentialQuantification #-}"
     , "{-# OPTIONS_GHC -Wall -fno-warn-orphans -fno-warn-overlapping-patterns #-}"
     , "{- AUTOGENERATED (See Data.Cursor.CLASE.Gen.Language) -}"
     , "module " ++ modName ++ "("
     , "   ContextI(..)"
     , "  ,TypeRepI(..)"
     , "  ,MovementI(..)"
     , "  ,Context(..)"
     , "  ,Movement(..)"
     , "  ,TypeRep(..)"
     , ") where"
     , "import Data.Maybe"
     , "import Data.Cursor.CLASE.Util"
     , "import Control.Arrow"
     , maybe "" ("import " ++) mRootModName
     , "import Data.Cursor.CLASE.Language"
     , ""
     , "instance Language " ++ rootName ++ " where"
     , "  data Context " ++ rootName ++ " from to = CW (ContextI from to)"
     , "  data Movement " ++ rootName ++ " d from to = MW (MovementI d from to)"
     , "  data TypeRep " ++ rootName ++ " t = TW (TypeRepI t)"
     , ""
     , "  buildOne (CW x) = buildOneI x"
     , "  unbuildOne (MW m) a = fmap (first CW) (unbuildOneI m a)"
     , "  invertMovement (MW x) = MW (invertMovementI x)"
     , "  movementEq (MW x) (MW y) = fmap snd $ movementEqI x y"
     , "  reifyDirection (MW x) = reifyDirectionI x"
     , "  contextToMovement (CW x) = MW (contextToMovementI x)"
     , "  downMoves (TW t) = map (\\(ExistsR x) -> ExistsR (MW x)) (downMovesI t)"
     , "  moveLeft  (MW m) = fmap (\\(ExistsR x) -> ExistsR (MW x)) (moveLeftI m)"
     , "  moveRight (MW m) = fmap (\\(ExistsR x) -> ExistsR (MW x)) (moveRightI m)"
     ]

buildLeftRightNames :: Map Name DataType -> Set Name
buildLeftRightNames = Set.fromList . map dtChildName . filter isWrapsList . Map.elems

printTypeRep :: Map Name DataType -> PrintM ()
printTypeRep nmap = do
  printLn "data TypeRepI a where"
  mapM_ printTypeRepLine (Map.keys nmap)
  where
    printTypeRepLine :: Name -> PrintM ()
    printTypeRepLine name = printLn $ "  " ++ (nameBase name) ++ "T :: TypeRepI " ++ (nameBase name)


printReify :: Name -> Map Name DataType -> PrintM ()
printReify langName nmap = do
  mapM_ printReifyInstance (Map.keys nmap)
  where
    printReifyInstance :: Name -> PrintM ()
    printReifyInstance name = do
      printLn . unwords $ ["instance Reify", (nameBase langName),
                                            (nameBase name), "where"]
      printLn $ "  reify = const $ TW " ++ (nameBase name) ++ "T"
      printLn ""

printContext :: [ContextCtr] -> PrintM ()
printContext ctrs = do
  printLn "data ContextI a b where"
  mapM_ printContextCtr ctrs
  where
    printContextCtr :: ContextCtr -> PrintM ()
    printContextCtr cc = do
      printLn $ "  " ++ (ctxCtrName cc) ++ " :: " ++ tipes ++ " " ++ (typeFrom) ++ " " ++ (typeTo)
      where
        tipes = concat . intersperse " -> " $ (ctxArgs) ++ ["ContextI"]
        typeFrom = nameBase . ctxCtrTypeFrom $ cc
        typeTo = nameBase . ctxCtrTypeTo $ cc

        ctxListArgs
          | ctxCtrIsList cc = replicate 2 ("[" ++ typeFrom ++ "]")
          | otherwise = []
        ctxArgs = (ctxCtrArgsBefore cc) ++ ctxListArgs ++ (ctxCtrArgsAfter cc)

printContextToMovement :: [ContextCtr] -> PrintM ()
printContextToMovement ctrs = do
  printLn "contextToMovementI :: ContextI a b -> MovementI Up a b"
  mapM_ printContextToMovementLine ctrs
  where
    printContextToMovementLine :: ContextCtr -> PrintM ()
    printContextToMovementLine cc = do
      printLn $ "contextToMovementI " ++ pattern ++ " = " ++ movPattern
        where
          pattern = "(" ++ vars ++ ")"
          movPattern = "(MUp " ++ (downCtrName cc) ++ ")"
          vars = unwords $ [ctxCtrName cc] ++ replicate (numCCArgs cc) "_"

printBuildOne :: [ContextCtr] -> PrintM ()
printBuildOne ctrs = do
  printLn "buildOneI :: ContextI a b -> a -> b"
  mapM_ printBuildOneLine ctrs
  where
    printBuildOneLine :: ContextCtr -> PrintM ()
    printBuildOneLine cc = do
        printLn $ "buildOneI (" ++ pattern ++ ") h = " ++ (ctrPattern)
        where
          (pattern, ctrPattern) = buildPatterns cc

buildPatterns :: ContextCtr -> (String, String)
buildPatterns cc = (pattern, ctrPattern)
  where
    ctrPattern = unwords $ [nameBase . ctxCtrCtrTo $ cc] ++ argsPattern
    pattern = unwords $ [ctxCtrName cc] ++ varsPattern

    argsPattern = preVars ++ ctxArgs ++ postVars

    varsPattern = preVars ++ ctxVars ++ postVars
    (preVars, postVars) = second (take (length $ ctxCtrArgsAfter cc)) . 
                            splitAt (length $ ctxCtrArgsBefore cc) $ vars
        
    ctxVars
      | ctxCtrIsList cc = ["l", "r"]
      | otherwise = []

    ctxArgs
      | ctxCtrIsList cc = ["((reverse l) ++ [h] ++ r)"]
      | otherwise = ["h"]

    vars :: [String] 
    vars = map (('x':) . show) [(0::Integer)..]

buildApplyMovementPatterns :: ContextCtr -> (String, String)
buildApplyMovementPatterns cc = (ctxPattern, ctrPattern)
  where
    ctxPattern = unwords $ [ctxCtrName cc] ++ argsPattern
    ctrPattern = unwords $ [nameBase . ctxCtrCtrTo $ cc] ++ varsPattern

    argsPattern = preVars ++ ctxArgs ++ postVars
    varsPattern = preVars ++ ctxVars ++ postVars

    (preVars, postVars) = second (take (length $ ctxCtrArgsAfter cc)) . 
                            splitAt (length $ ctxCtrArgsBefore cc) $ vars

    ctxVars
      | ctxCtrIsList cc = ["(h:es)"]
      | otherwise = ["h"]

    ctxArgs
      | ctxCtrIsList cc = ["[] es"]
      | otherwise = []
    
    vars :: [String]
    vars = map (('x':) . show) [(0::Integer)..]

buildLeftRightPatterns :: ContextCtr -> String -> String -> (String, String)
buildLeftRightPatterns cc lop rop = (ctxPatternBefore, ctxPatternAfter)
  where
    ctxPatternBefore = unwords $ [ctxCtrName cc] ++ beforePattern
    ctxPatternAfter  = unwords $ [ctxCtrName cc] ++ afterPattern

    beforePattern  = preVars ++ bArgs ++ postVars
    afterPattern   = preVars ++ aArgs ++ postVars

    (preVars, postVars) = second (take (length $ ctxCtrArgsAfter cc)) . 
                            splitAt (length $ ctxCtrArgsBefore cc) $ vars

    bArgs
      | ctxCtrIsList cc = ["l", "r"]
      | otherwise = error "CtxCtr is not list!"

    aArgs
      | ctxCtrIsList cc = ["(" ++ lop ++ "l" ++")","(" ++ rop ++ "r" ++ ")"]
      | otherwise = error "CtxCtr is not list!"
    
    vars :: [String]
    vars = map (('x':) . show) [(0::Integer)..]

printMovement :: [ContextCtr]  -> Set Name -> PrintM ()
printMovement downRoutes _ {- TODO leftRights -} = do
  printLn "data MovementI d a b where"
  printLn "  MUp :: MovementI Down b a -> MovementI Up a b"
  mapM_ printDownMovement downRoutes
  --mapM_ printLeftRight (Set.toList leftRights)
  where
    printDownMovement :: ContextCtr -> PrintM ()
    printDownMovement cc = do
      printLn $ "  " ++ (downCtrName cc) ++ " :: MovementI Down " ++ fromS ++ " " ++ toS
      where
        fromS = nameBase . ctxCtrTypeTo $ cc
        toS   = nameBase . ctxCtrTypeFrom $ cc

{-
    printLeftRight :: Name -> PrintM ()
    printLeftRight name = do
      printLn $ "  M" ++ nameS ++ "Left" ++ rest
      printLn $ "  M" ++ nameS ++ "Right" ++ rest
      where
        nameS = nameBase name 
        rest = " :: Movement " ++ nameS ++ " " ++ nameS
-}

printMovementEq :: [ContextCtr] -> Set Name -> PrintM ()
printMovementEq downRoutes leftRights = do
  printLn "movementEqI :: MovementI d x y -> MovementI d a b -> Maybe (TyEq x a, TyEq y b)"
  printLn "movementEqI (MUp a) (MUp b) = fmap (\\(x,y) -> (y,x)) $ movementEqI a b"
  mapM_ printDownMovementEq downRoutes
  mapM_ printLeftRightEq (Set.toList leftRights)
  printLn "movementEqI _ _ = Nothing"
  where
    printDownMovementEq :: ContextCtr -> PrintM ()
    printDownMovementEq cc = do
      printLn $ "movementEqI " ++ (downCtrName cc) ++ " " ++ (downCtrName cc) ++ " = Just (Eq, Eq)"

    printLeftRightEq :: Name -> PrintM ()
    printLeftRightEq name = do
      printLn $ "movementEqI " ++ (nameS ++ "Left") ++ " " ++ (nameS ++ "Left") ++ " = Just (Eq, Eq)"
      printLn $ "movementEqI " ++ (nameS ++ "Right") ++ " " ++ (nameS ++ "Right") ++ " = Just (Eq, Eq)"
      where
        nameS = "M" ++ nameBase name

printInvertMovement :: [ContextCtr] -> Set Name -> PrintM ()
printInvertMovement ctxCtrs leftRights = do
  printLn "invertMovementI :: MovementI d a b -> MovementI (Invert d) b a"
  printLn "invertMovementI (MUp dwn) = dwn"
  mapM_ printInvertDownRoute downPatterns
  mapM_ printInvertLeftRights (Set.toList leftRights)
  where
    downPatterns = map getDownCtr $ ctxCtrs

    getDownCtr :: ContextCtr -> String
    getDownCtr ctxCtr = sDownCtr
      where
        sDownCtr = liftM3 (\x y z -> "M" ++ x ++ "To" ++ y ++ z) 
                             (nameBase . ctxCtrCtrTo) 
                             (nameBase . ctxCtrTypeFrom) 
                             (maybe "" show . ctxCtrOffset) ctxCtr 

    printInvertDownRoute :: String -> PrintM ()
    printInvertDownRoute sDownCtr = do
      printLn $ "invertMovementI " ++ sDownCtr ++ " = MUp (" ++ sDownCtr ++ ")"

    printInvertLeftRights :: Name -> PrintM ()
    printInvertLeftRights name = do
      printLn $ "invertMovementI " ++ (nameS ++ "Left = IMLeftRight " ++ nameS ++ "Right")
      printLn $ "invertMovementI " ++ (nameS ++ "Right = IMLeftRight " ++ nameS ++ "Left")
      where
        nameS = "M" ++ nameBase name

printUnbuildOne :: [ContextCtr] -> Set Name -> PrintM ()
printUnbuildOne contextCtrs leftRights = do
  printLn . unlines $ 
    [ "unbuildOneI :: MovementI Down a b -> a -> Maybe (ContextI b a, b)",
      "unbuildOneI mov here = case mov of"]
  downMoveCases
  leftRightMoveCases
  printLn $ "  _ -> Nothing"
  where
    downMoveCases = mapM_ downMoveCase contextCtrs

    downMoveCase cc = do
      printLn $ "  " ++ (downCtrName cc) ++ " -> case here of"
      printLn $ "    (" ++ ctrPattern ++ ") -> " ++ "Just $ (" ++ ctrctxPattern ++ ", h)"
      printLn $ "    _ -> Nothing"
      where
        (ctrctxPattern, ctrPattern) = buildApplyMovementPatterns cc

    leftRightMoveCases = mapM_ leftRightMoveCase (Set.toList leftRights)

    leftRightMoveCase name = do
      printLn $ "      M" ++ nameBase name ++ "Left -> case step of"
      mapM_ leftMoveCase ctxCtrs
      printLn $ "        _ -> Nothing"
      printLn $ "      M" ++ nameBase name ++ "Right -> case step of"
      mapM_ rightMoveCase ctxCtrs
      printLn $ "        _ -> Nothing"
      where
        ctxCtrs = filter (liftM2 (&&) ((== name) . ctxCtrTypeFrom)
                                      (ctxCtrIsList)) contextCtrs

        leftMoveCase = moveCase "l" "tail " "it:"
        rightMoveCase = moveCase "r" "it:" "tail " 

        moveCase lr y z ctxCtr = do
          printLn $ "        (" ++ ctxCtrPatternBefore ++ ") -> if' (null " ++ lr ++ ") " ++ 
                      "Nothing " ++
                      "(Just $ ((head " ++ lr ++ "), " ++ ctxCtrPatternAfter ++ ")"
          where
            (ctxCtrPatternBefore, ctxCtrPatternAfter) = buildLeftRightPatterns ctxCtr y z


printReifyDirection :: [ContextCtr] -> PrintM ()
printReifyDirection ctrs = do
  printLn $ "reifyDirectionI :: MovementI d a b -> DirectionT d"
  printLn $ "reifyDirectionI d = case d of"
  printLn $ "  (MUp _) -> UpT"
  mapM_ reifyDirectionDown ctrs
  where
    reifyDirectionDown :: ContextCtr -> PrintM ()
    reifyDirectionDown cc = printLn $ "  " ++ (downCtrName cc) ++ " -> DownT"

printGenericMoveLR :: String -> (ContextCtr -> [Child] -> Maybe (Child, Maybe Int)) -> [ContextCtr] -> Map Name DataType -> Name -> PrintM ()
printGenericMoveLR leftOrRight beforeOrAfterFn ctxcts dataMap rootName = do
  printLn $ "move" ++ leftOrRight ++ "I :: MovementI Down a x -> Maybe (ExistsR " ++ (nameBase rootName) ++  " (MovementI Down a))"
  printLn $ "move" ++ leftOrRight ++ "I mov = case mov of"
  mapM_ genericMove ctxcts
  printLn "    _ -> Nothing"
  where
    genericMove :: ContextCtr -> PrintM ()
    genericMove cc
      | isJust nextKid = do
          printLn $ "    " ++ (downCtrName cc) ++ " -> Just $ ExistsR " ++ notListMovement
      | otherwise = return ()
      where
        currName = ctxCtrCtrTo cc
        notListMovement = "M" ++ nameBase currName ++ "To" ++ nxtName
        --TODO listMovement = "M" ++ nameBase lstNxtName ++ leftOrRight
        --TODO lstNxtName = ctxCtrTypeFrom cc

        (Just ctr) = liftM (head . filter ( (== (ctxCtrCtrTo cc)) . ctrName ) . dtCtrs) . 
                      Map.lookup (ctxCtrTypeTo cc) $ dataMap
        nextKid@(~(Just (child, moffset))) = beforeOrAfterFn cc (ctrKids ctr)
        nxtName = (nameBase . childType $ child) ++ (maybe "" show moffset)
      

printGenericMoveRight :: [ContextCtr] -> Map Name DataType -> Name -> PrintM ()
printGenericMoveRight = printGenericMoveLR "Right" grabNextChild
  where
    grabNextChild :: ContextCtr -> [Child] -> Maybe (Child, Maybe Int)
    grabNextChild cc kids = liftM (id &&& calcOffset . childType) nextKid
      where
        calcOffset :: Name -> Maybe Int
        calcOffset nam
          | length allOtherKidsWithTheSameName > 1 = Just (length otherKidsWithTheSameNameBefore)
          | otherwise = Nothing
          where
            otherKidsWithTheSameNameBefore = filter (== nam) $ beforeKids ++ [currKid] 
            allOtherKidsWithTheSameName    = otherKidsWithTheSameNameBefore ++ (filter (== nam) afterKids)
        nextKid = listToMaybe . filter isNavigable . drop 1 . drop (length beforeKids) $ kids
        beforeKids = ctxCtrTypesBefore cc
        afterKids  = ctxCtrTypesAfter cc
        currKid    = ctxCtrTypeFrom cc

printGenericMoveLeft :: [ContextCtr] -> Map Name DataType -> Name -> PrintM ()
printGenericMoveLeft = printGenericMoveLR "Left" grabNextChild
  where
    grabNextChild :: ContextCtr -> [Child] -> Maybe (Child, Maybe Int)
    grabNextChild cc kids = liftM (id &&& calcOffset . childType) nextKid
      where
        calcOffset :: Name -> Maybe Int
        calcOffset nam
          | length allOtherKidsWithTheSameName > 1 = Just (pred . length $ otherKidsWithTheSameNameBefore)
          | otherwise = Nothing
          where
            otherKidsWithTheSameNameBefore = filter (== nam) $ beforeKids

            allOtherKidsWithTheSameName    = otherKidsWithTheSameNameBefore ++ (filter (== nam) [currKid] ++ afterKids)

        nextKid    = listToMaybe . filter isNavigable . reverse . take (length beforeKids) $ kids
        beforeKids = ctxCtrTypesBefore cc
        afterKids  = ctxCtrTypesAfter cc
        currKid    = ctxCtrTypeFrom cc


printDownMovements :: Name -> [ContextCtr] -> PrintM ()
printDownMovements langName ctxCtrs = do
  printLn $ "downMovesI :: TypeRepI a -> [ExistsR " ++ lNameS ++ " (MovementI Down a)]"
  printLn $ "downMovesI tr = case tr of "
  mapM_ (uncurry printTrCase) ctrsByTypeTo
  where
    lNameS = nameBase langName

    ctrsByTypeTo :: [(Name, [ContextCtr])]
    ctrsByTypeTo = map ((ctxCtrTypeTo . head) &&& id) . 
                   groupBy ((==) `on` ctxCtrTypeTo) . 
                   sortBy (comparing ctxCtrTypeTo) $ ctxCtrs

    printTrCase :: Name -> [ContextCtr] -> PrintM ()
    printTrCase trCase localCtrs = do
      printLn $ "  " ++ (nameBase trCase) ++ "T -> [" ++ dctrs ++ "]"
      where
        dctrs = concat . intersperse ", " . map (\x -> "(ExistsR " ++ x ++ ")") $ downCtrs
        downCtrs = map downCtrName localCtrs