{-# LANGUAGE FlexibleContexts #-}
module Main where

-- import Debug.Trace

import Control.Monad
import Data.Array (Array)
import Data.Coerce
import Data.Maybe
import Data.Semigroup
import System.IO

import qualified Data.Array.IArray as IArray
import qualified Data.Array.ST.Safe as STArray
import qualified Data.ByteString.Char8 as BS
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as List

main :: IO ()
main = do
  (g1Vc, _, g1Edges) <- parseData "./data.txt"
  let g1 = buildAdjacencyListByHead g1Vc g1Edges
  let g1Result = bellmanFord g1 g1Vc 1
  print g1Result

parseData :: FilePath -> IO (Int, Int, [(Int, Int, Double)])
parseData fp = withFile fp ReadMode $ \h -> do
  (vertexCount, edgeCount) <- getMetaData h
  edges <- replicateM edgeCount (getEdge h)
  pure (vertexCount, edgeCount, edges)
  where
    getMetaData h = parseMetaData <$> BS.hGetLine h
    parseMetaData line = fromJust $ do
      (x, line') <- BS.readInt line
      (y, line'') <- BS.readInt (dropSpace line')
      if BS.null (dropSpace line'')
        then pure (x, y)
        else error "Invalid data format"

    getEdge h = parseEdge <$> BS.hGetLine h
    parseEdge line = fromJust $ do
      (edgeTail, line') <- BS.readInt line
      (edgeHead, line'') <- BS.readInt (dropSpace line')
      (weight, line''') <- BS.readInteger (dropSpace line'')
      if BS.null (dropSpace line''')
        then pure (edgeTail, edgeHead, fromInteger weight)
        else error "Invalid data format"

    dropSpace = BS.dropWhile (== ' ')

buildAdjacencyListByHead :: (Fractional a, Ord a) => Int -> [(Int, Int, a)] -> Array Int [(Int, a)]
buildAdjacencyListByHead vertexCount =
  IArray.listArray (1, vertexCount)
  . IntMap.elems
  . List.foldl' (\acc (f, t, w) -> IntMap.adjust ((f, w) :) t acc) initialMap
  where
    initialMap = IntMap.fromAscList (List.zip [1 .. vertexCount] (List.repeat []))
    
bellmanFord :: Array Int [(Int, Double)] -> Int -> Int -> Either String (Array Int Double)
bellmanFord g vC s =
  if hasNegativeCycle then
    Left "Bellman-Ford algorithm halted: negative cycle found"
  else
    Right resultArray
  where
    calculated = STArray.runSTUArray $ do
      array <- STArray.newArray (1, vC) (1 / 0)
      STArray.writeArray array s 0
      go array 1 1
      where
        go array stage v
          | v > vC =
            if stage < vC
            then go array (stage + 1) 1
            else pure array
          | otherwise = do
            forM_ (g IArray.! v) $ coreCmp array v
            go array stage (v + 1)

        coreCmp array v (f, w) = do
          vW <- STArray.readArray array v
          fW <- STArray.readArray array f
          STArray.writeArray array v (min (fW + w) vW)
        {-# INLINE coreCmp #-}

    resultList = IArray.elems calculated
    resultArray = IArray.listArray (1, vC) resultList

    getResultOf = (resultArray IArray.!)
    {-# INLINE getResultOf #-}

    hasNegativeCycle =
      coerce
      . foldMap (\(t, efs) -> foldMap (\(f, w) -> coerce (getResultOf f + w < getResultOf t) :: Any) efs)
      $ IArray.assocs g