{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE ConstrainedClassMethods    #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE NamedFieldPuns             #-}
{-# LANGUAGE RecordWildCards            #-}

module Extrapolation (Extrapolator(..), LinearExtrapolator, linearDelay, distanceAlongLine) where
import           Data.Foldable      (maximumBy, minimumBy)
import           Data.Function      (on)
import           Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map           as M
import           Data.Time          (Day, UTCTime (UTCTime, utctDay),
                                     diffUTCTime, getCurrentTime,
                                     nominalDiffTimeToSeconds)
import qualified Data.Vector        as V
import           GHC.Float          (int2Double)
import           GHC.IO             (unsafePerformIO)

import           Conduit            (MonadIO (liftIO))
import           Data.List          (sortBy)
import           GTFS               (Depth (Deep), GTFS (..), Seconds (..),
                                     Shape (..), Station (stationName),
                                     Stop (..), Time, Trip (..), seconds2Double,
                                     stationGeopos, toSeconds)
import           Persist            (Running (..), TrainAnchor (..),
                                     TrainPing (..))
import           Server.Util        (utcToSeconds)

class Extrapolator a where
  -- | here's a position ping, guess things from that!
  extrapolateAnchorFromPing :: GTFS -> Running -> TrainPing -> TrainAnchor

  -- | extrapolate status at some time (i.e. "how much delay does the train have *now*?")
  extrapolateAtSeconds :: NonEmpty TrainAnchor -> Seconds -> Maybe TrainAnchor
  -- | extrapolate status at some places (i.e. "how much delay will it have at the next station?")
  extrapolateAtPosition :: NonEmpty TrainAnchor -> Double -> Maybe TrainAnchor

data LinearExtrapolator

instance Extrapolator LinearExtrapolator where
  -- TODO: this kind of sorting is bullshit; should look up depending on position,
  -- not time (else there's glitches)
  extrapolateAtSeconds history secondsNow =
    fmap (minimumBy (compare `on` difference))
    $ NE.nonEmpty $ NE.filter (\a -> trainAnchorWhen a < secondsNow) history
    where difference status = secondsNow - (trainAnchorWhen status)

  -- note that this sorts (descending) for time first as a tie-breaker
  -- (in case the train just stands still for a while, take the most recent update)
  extrapolateAtPosition history positionNow =
    fmap (minimumBy (compare `on` difference))
    $ NE.nonEmpty $ sortBy (flippedCompare `on` trainAnchorWhen)
    $ NE.filter (\a -> trainAnchorSequence a < positionNow) history
    where difference status = positionNow - (trainAnchorSequence status)
          flippedCompare a b = case compare a b of
            LT -> GT
            GT -> LT
            a  -> a

  extrapolateAnchorFromPing gtfs@GTFS{..} Running{..} ping@TrainPing{..} = TrainAnchor
    { trainAnchorCreated = trainPingTimestamp
    , trainAnchorTrip = runningTrip
    , trainAnchorDay = runningDay
    , trainAnchorWhen = utcToSeconds trainPingTimestamp runningDay
    , trainAnchorSequence
    , trainAnchorDelay
    , trainAnchorMsg = Nothing
    }
    where Just trip = M.lookup runningTrip trips
          (trainAnchorDelay, trainAnchorSequence) = linearDelay gtfs trip ping runningDay

linearDelay :: GTFS -> Trip Deep Deep -> TrainPing -> Day -> (Seconds, Double)
linearDelay GTFS{..} trip@Trip{..} TrainPing{..} runningDay = unsafePerformIO $ do
  print (observedPosition, observedProgress)
  print (stationName . stopStation $ lastStop, stationName . stopStation $ nextStop)
  print (distanceAlongLine line (stationGeopos $ stopStation lastStop) closestPoint
        , distanceAlongLine line (stationGeopos $ stopStation lastStop) (stationGeopos $ stopStation nextStop))
  pure (observedDelay, observedPosition)
  where closestPoint = minimumBy (compare `on` euclid (trainPingLat, trainPingLong)) line
        line = shapePoints tripShape
        lastStop = tripStops V.! (nextIndex - 1)
        nextStop = tripStops V.! nextIndex
        nextIndex = if idx' == 0
          then 1 else idx'
          where idx' = fst $ V.minimumBy (compare `on` snd)
                  $ V.filter (\(_,dist) -> dist > 0) $ V.indexed
                  $ fmap (distanceAlongLine line closestPoint . stationGeopos . stopStation) tripStops
        expectedTravelTime =
            toSeconds (stopArrival nextStop) tzseries runningDay
            - toSeconds (stopDeparture lastStop) tzseries runningDay
        expectedProgress = crop $
            seconds2Double (utcToSeconds trainPingTimestamp runningDay
                         - toSeconds (stopDeparture lastStop) tzseries runningDay)
            / seconds2Double expectedTravelTime
            where crop a
                    | a < 0 = 0
                    | a > 1 = 1
                    | otherwise = a
        observedProgress =
            distanceAlongLine line (stationGeopos $ stopStation lastStop) closestPoint
            / distanceAlongLine line (stationGeopos $ stopStation lastStop) (stationGeopos $ stopStation nextStop)
        observedPosition =
          (int2Double $ stopSequence lastStop) + observedProgress * (int2Double $ stopSequence nextStop - stopSequence lastStop)
        observedDelay = Seconds $ round $
          (expectedProgress - observedProgress) * int2Double (unSeconds expectedTravelTime)
          -- if the hypothetical on-time train is already at (or past) the next station,
          -- just add the time distance we're behind
          + if expectedProgress == 1
            then seconds2Double (utcToSeconds trainPingTimestamp runningDay - toSeconds (stopArrival nextStop) tzseries runningDay)
            else 0

distanceAlongLine :: V.Vector (Double, Double) -> (Double, Double) -> (Double, Double) -> Double
distanceAlongLine line p1 p2 = along2 - along1
  where along1 = along p1
        along2 = along p2
        along p@(x,y) =
          sumSegments
          $ V.take (index + 1) line
          where index = V.minIndexBy (compare `on` euclid p) line
        sumSegments :: V.Vector (Double, Double) -> Double
        sumSegments line = snd
          $ foldl (\(p,a) p' -> (p', a + euclid p p')) (V.head line,0) $ line



euclid :: Floating f => (f,f) -> (f,f) -> f
euclid (x1,y1) (x2,y2) = sqrt (x*x + y*y)
  where x = x1 - x2
        y = y1 - y2