import MyPrioQueue
import Control.Monad.State
import Control.Parallel.Strategies(using, parList, rparWith, rseq)
import System.Environment(getArgs)
import System.Exit(die)
import System.IO.Error(catchIOError, ioeGetFileName, isUserError)
import qualified Data.Map as M
import qualified Data.Set as S

data SearchState = SearchState
  { s_queue  :: MyPrioQueue
  , s_stats  :: M.Map (Int, Int) (Double, Double)
  , s_obs    :: S.Set (Int, Int)
  , seen_obs :: S.Set (Int, Int)
  , s_km     :: Double
  , s_start  :: (Int, Int)
  , s_start0 :: (Int, Int)
  , s_goal   :: (Int, Int)
  , s_last   :: (Int, Int)
  , s_height :: Int
  , s_width  :: Int
  }
  deriving (Show)

{- A function that takes two states as input and returns the Manhattan
   distance between them  -}
manDist :: (Int, Int) -> (Int, Int) -> Double
manDist (x1, y1) (x2, y2) = fromIntegral $ abs (x1 - x2) + abs (y1 - y2)

{- A function that takes two adjacent states and the current state monad as
   input and returns the edge cost between the two states -}
getEdge :: (Int, Int) -> (Int, Int) -> SearchState -> Double
getEdge s1 s2 ss 
    | b1 == True || b2 == True = 1/0
    | otherwise = 1
    where b1 = evalState (isObstacle s1) ss
          b2 = evalState (isObstacle s2) ss

{- A function that takes a state and the current state monad as input and
   returns the g-value of the state -}
getGVal :: (Int, Int) -> SearchState -> Double
getGVal s ss = g
    where Just (g, _) = M.lookup s (s_stats ss)

{- A function that takes a state as input and returns its successors -}
getAdj :: (Int, Int) -> State SearchState [(Int, Int)]
getAdj (x, y) = do params <- get
                   let height = s_height params
                   let width = s_width params
                   let adj = filter (\(a, b) -> a >= 0 && b >= 0 &&
                                                a < width && b < height)
                             [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]
                   return adj

{- A function that takes a state as input and scans its surrounding for
   obstacles. It returns a list containing any obstacles that were found. -}
scan :: (Int, Int) -> State SearchState [(Int, Int)]
scan s = do params <- get
            let known_obs = s_obs params
            let scanned_obs = seen_obs params
            adj <- getAdj s
            let adj_obs = filter (\s2 -> s2 `S.member` known_obs &&
                                  not (s2 `S.member` scanned_obs)) adj
            return adj_obs

{- A function that takes a list of newly found obstacles and updates the edge
   costs of adjacent states -}
scanUpdate :: [(Int, Int)] -> State SearchState ()
scanUpdate obs = do params <- get
                    let scanned_obs = seen_obs params
                    let new_set = foldl (\acc s2 -> S.insert s2 acc) scanned_obs
                                  obs
                    put (params {seen_obs = new_set})
                    mapM_ (\s2 -> mapM_ updateVertex
                          (s2 : (evalState (getAdj s2) params))) obs

{- A function that takes a state as input and returns whether it is an
   obstacle -}
isObstacle :: (Int, Int) -> State SearchState Bool
isObstacle s = do params <- get
                  let obs = seen_obs params
                  return (s `S.member` obs)

{- A function that takes a state as input, and returns its priority value -}
calculateKey :: (Int, Int) -> State SearchState (Double, Double)
calculateKey s = do params <- get
                    let Just (g, rhs) = M.lookup s (s_stats params)
                    return (min g rhs + manDist s (s_start params) +
                            s_km params, min g rhs)

{- A function that takes a state as input and updates its attributes -}
updateVertex :: (Int, Int) -> State SearchState ()
updateVertex s = do params <- get
                    let q = s_queue params
                    let m = s_stats params
                    let cost s' = getEdge s s' params + getGVal s' params
                    adj <- getAdj s
                    let new_rhs = minimum $ map cost adj
                    let goal = s_goal params
                    let Just (g, _) = M.lookup s m
                    let nm = if (s /= goal) then (M.insert s (g, new_rhs) m)
                        else m
                    put (params {s_stats = nm})
                    let nq1 = if (member s q) then (remove s q) else q
                    let Just (g', rhs') = M.lookup s nm
                    key <- calculateKey s
                    let nq2 = if (g' /= rhs') then (insert s key nq1) else nq1
                    put (params {s_queue = nq2, s_stats = nm})

{- A helper function for updateQueue -}
updateQueueHelper :: (Int, Int) -> (Double, Double) -> (Double, Double) ->
                     (Double, Double) -> State SearchState ()
updateQueueHelper s (g, rhs) new_key old_key
    | old_key < new_key = do params <- get
                             let q = s_queue params
                             put (params {s_queue = insert s new_key q})
    | g > rhs =           do params <- get
                             let m = s_stats params
                             put (params {s_stats = M.insert s (rhs, rhs) m})
                             adj <- getAdj s
                             mapM_ updateVertex adj
    | otherwise =         do params <- get
                             let m = s_stats params
                             put (params {s_stats = M.insert s (1/0, rhs) m})
                             adj <- getAdj s
                             mapM_ updateVertex (s : adj)

{- A function that updates the algorithm's priority queue -}
updateQueue :: State SearchState ()
updateQueue = do params <- get
                 let q = s_queue params
                 let start = s_start params
                 start_key <- calculateKey start
                 if (not (empty q)) then do
                    let (s, old_key) = findMin q
                    let Just (g, rhs) = M.lookup s (s_stats params)
                    let Just (s_g, s_rhs) = M.lookup start (s_stats params)
                    if (old_key < start_key || s_rhs /= s_g) then do
                       new_key <- calculateKey s
                       put (params {s_queue = remove s q})
                       updateQueueHelper s (g, rhs) new_key old_key
                       updateQueue
                    else return ()
                 else return ()

{- This ia the function that build the algorithm's current proposed path by
   following the cheapest transitions from the start to goal state. It returns
   the proposed path.  -}
computeShortestPath :: (Int, Int) -> [(Int, Int)] -> SearchState -> [(Int, Int)]
computeShortestPath s path ss
    | next == goal = path'
    | otherwise = computeShortestPath next path' ss
    where (_, next) = minimum $ map cost adj
          cost s' = (getEdge s s' ss + getGVal s' ss, s')
          adj = evalState (getAdj s) ss
          path' = path ++ [next]
          goal = s_goal ss

{- This is the main driving function for the D* Lite algorithm. It terminates
   when the goal state is reached, returning all of the proposed paths. -}
driver :: (Int, Int) -> SearchState -> [[(Int, Int)]] ->
          State SearchState [[(Int, Int)]]
driver s ss paths
    | s == s_goal ss = do return paths
    | otherwise = do params <- get
                     let cost s' = (getEdge s s' params + getGVal s' params, s')
                     adj <- getAdj s
                     let (_, new_start) = minimum $ map cost adj
                     put (params {s_start = new_start})
                     let l = s_last params
                     let km = s_km params
                     adj_obs <- scan s
                     updated_params <- get
                     if (adj_obs == [])
                        then (do driver new_start updated_params paths)
                     else do put (updated_params {s_km = km +
                                  manDist l new_start, s_last = new_start})
                             scanUpdate adj_obs
                             updateQueue
                             final_params <- get
                             let start0 = s_start0 final_params
                             let new_path = computeShortestPath start0
                                            [start0] final_params
                                            `using` rparWith rseq
                             if (new_path `elem` paths) then
                                driver s final_params paths
                             else driver s final_params (paths ++ [new_path])

{- This function is called by main and starts the driver function. It returns
   all of the algorithm's proposed paths. -}
getPaths :: State SearchState [[(Int, Int)]]
getPaths = do updateQueue
              params <- get
              let start = s_start params
              let first_path = computeShortestPath start [start] params
              let paths = evalState (driver start params [first_path]) params
                          `using` parList rseq
              return paths

main :: IO ()
main = do [height, weight, start, goal, obstacles] <- getArgs
          let h = read height::Int
          let w = read weight::Int
          let s = read start::(Int, Int)
          let g = read goal::(Int, Int)
          let obs = read obstacles::[(Int, Int)]

          let init_queue = insert g (manDist s g, 0)
                           (MyPrioQueue M.empty S.empty)
          let init_stats = M.insert g (1/0, 0) (M.fromList
                           [((x, y), (1/0, 1/0)) | x <- [0..(w - 1)],
                                                   y <- [0..(h - 1)]])
          let init_obs = S.fromList obs
          let ss = SearchState init_queue init_stats init_obs S.empty 0 s s g s
                               h w
          let paths = evalState getPaths ss `using` parList rseq
          mapM_ (putStrLn . show) paths
       `catchIOError` \e -> die $ case ioeGetFileName e of
          _ | isUserError e -> "Usage: ./dstar <height> <width> <start> <goal> \
                                \<list of obstacles>"
            | otherwise     -> show e
