{-# LANGUAGE ApplicativeDo                        #-}
{-# LANGUAGE DataKinds                            #-}
{-# LANGUAGE DeriveFunctor                        #-}
{-# LANGUAGE DeriveGeneric                        #-}
{-# LANGUAGE FlexibleContexts                     #-}
{-# LANGUAGE FlexibleInstances                    #-}
{-# LANGUAGE FunctionalDependencies               #-}
{-# LANGUAGE GADTs                                #-}
{-# LANGUAGE KindSignatures                       #-}
{-# LANGUAGE LambdaCase                           #-}
{-# LANGUAGE MultiParamTypeClasses                #-}
{-# LANGUAGE OverloadedStrings                    #-}
{-# LANGUAGE PartialTypeSignatures                #-}
{-# LANGUAGE PolyKinds                            #-}
{-# LANGUAGE RankNTypes                           #-}
{-# LANGUAGE RecordWildCards                      #-}
{-# LANGUAGE ScopedTypeVariables                  #-}
{-# LANGUAGE StandaloneDeriving                   #-}
{-# LANGUAGE TemplateHaskell                      #-}
{-# LANGUAGE TupleSections                        #-}
{-# LANGUAGE TypeApplications                     #-}
{-# LANGUAGE TypeFamilies                         #-}
{-# LANGUAGE TypeOperators                        #-}
{-# LANGUAGE TypeSynonymInstances                 #-}
{-# LANGUAGE UndecidableInstances                 #-}
{-# LANGUAGE ViewPatterns                         #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}

module Egg.Research (
    BonusAmount(..), AsBonusAmount(..)
  , BonusType(..), AsBonusType(..)
  , Bonuses(..), _Bonuses
  , Research(..), HasResearch(..)
  , ResearchTier(..), rtUnlock, rtTechs, SomeResearchTier
  , ResearchData(..), rdCommon, rdEpic, SomeResearchData
  , ResearchStatus(..), rsCommon, rsEpic, SomeResearchStatus
  , ResearchIx(..), _RICommon, _RIEpic
  , ResearchError(..), _REMaxedOut, _RELocked
  , scaleAmount
  , hasEffect
  , bonusEffect
  , bonusEffectFor
  , bonusing
  , bonusingFor
  , emptyBonuses
  , bonuses
  , maxLevel
  , emptyResearchStatus
  , researchBonuses
  , totalBonuses
  , researchCount
  , legalTiers
  , foldResearch
  , purchaseResearch
  , researchIxNum
  , researchIxData
  , researchIxStatus
  , researchIxStatusLegal
  , withSomeResearch
  , researchIxesCommon
  , researchIxesEpic
  , legalResearchIxesCommon
  , legalResearchIxesEpic
  , legalResearchesCommon
  , legalResearchesEpic
  -- * Why?
  , rdrsCommon
  , rdrsEpic
  , researchIx
  ) where

import           Control.Applicative hiding      (some)
import           Control.Lens hiding             ((.=), (:<), Index)
import           Data.Aeson.Encoding
import           Data.Aeson.Types
import           Data.Bifunctor
import           Data.Dependent.Sum
import           Data.Finite
import           Data.Foldable
import           Data.Kind
import           Data.Maybe
import           Data.Monoid                     (First(..))
import           Data.Singletons
import           Data.Singletons.Prelude hiding  (Flip)
import           Data.Singletons.TypeLits
import           Data.Type.Combinator
import           Data.Type.Combinator.Singletons
import           Data.Type.Combinator.Util
import           Data.Type.Conjunction
import           Data.Type.Index
import           Data.Type.Product               as TCP
import           Data.Type.Sum
import           Data.Type.Vector
import           Data.Vector.Sized.Util
import           Egg.Commodity
import           GHC.Generics                    (Generic)
import           Numeric.Lens
import           Numeric.Natural
import           Statistics.LinearRegression
import           Type.Class.Higher
import           Type.Class.Witness
import           Type.Family.Tuple
import qualified Data.Map                        as M
import qualified Data.Text                       as T
import qualified Data.Vector                     as V
import qualified Data.Vector.Sized               as SV

data BonusAmount
    -- | Bonus adds an absolute amount.  Scales by multiplication, and
    -- compounds by addition.
    = BAIncrement Double
    -- | Bonus adds a percentage point multiplier.  Scales by
    -- multiplication (three 10% buffs is a 15% buff) and compounds by
    -- appropriate multiplication.
    | BAPercent Double
    -- | Bonus multiplies by a given ratio.  Scales by exponentiation
    -- (three 2x buffs is a 8x buff) and compounds by multiplication.
    | BAMultiplier Double
  deriving (Show, Eq, Ord, Generic)

makeClassyPrisms ''BonusAmount

data BonusType =
        BTBuildCosts
      | BTDroneRewards
      | BTEggValue
      | BTFarmValue
      | BTFleetSize
      | BTHabCapacity
      | BTHatcheryCapacity
      | BTHatcheryRate
      | BTHoldingHatch
      | BTHoverVehicleCapacity
      | BTInternalHatchery
      | BTInternalHatcheryCalm
      | BTInternalHatcherySharing
      | BTLayingRate
      | BTLongWarpTime
      | BTMaxRunningBonus
      | BTPrestigeEggs
      | BTResearchCosts
      | BTRunningBonus
      | BTSiloQuality
      | BTSoulEggBonus
      | BTVehicleCapacity
      | BTVehicleCosts
      | BTVideoDoublerTime
      | BTVehicleSpeed
  deriving (Show, Eq, Ord, Generic)

makeClassyPrisms ''BonusType

newtype Bonuses = Bonuses { _bMap :: M.Map BonusType [BonusAmount] }
    deriving (Show, Eq, Ord)

makePrisms ''Bonuses
makeWrapped ''Bonuses

data Research a =
    Research { _rName        :: T.Text
             , _rDescription :: T.Text
             , _rBaseBonuses :: Bonuses
             , _rCosts       :: V.Vector (Maybe a)
             }
  deriving (Show, Eq, Ord, Generic, Functor)

makeClassy ''Research

data ResearchTier :: Nat -> Type where
    ResearchTier
        :: { _rtUnlock :: Natural
           , _rtTechs  :: SV.Vector n (Research Bock)
           }
        -> ResearchTier n
  deriving (Show, Eq, Ord, Generic)

makeLenses ''ResearchTier

type SomeResearchTier = DSum Sing ResearchTier

data ResearchData :: [Nat] -> Nat -> Type where
    ResearchData
        :: { _rdCommon :: Prod ResearchTier tiers
           , _rdEpic   :: SV.Vector epic (Research GoldenEgg)
           }
        -> ResearchData tiers epic
  deriving (Show, Eq, Ord)

makeLenses ''ResearchData

type SomeResearchData = DSum Sing (Uncur ResearchData)

data ResearchStatus :: [Nat] -> Nat -> Type where
    ResearchStatus
        :: { _rsCommon :: Prod (Flip SV.Vector Natural) tiers
           , _rsEpic   :: SV.Vector epic Natural
           }
        -> ResearchStatus tiers epic
  deriving (Show, Eq, Ord, Generic)

makeLenses ''ResearchStatus

type SomeResearchStatus = DSum Sing (Uncur ResearchStatus)

-- | A safe index for a given research item, usable with 'ResearchData' and
-- 'ResearchStatus'.
data ResearchIx :: [Nat] -> Nat -> Type -> Type where
    RICommon :: Sum Finite tiers -> ResearchIx tiers epic Bock
    RIEpic   :: Finite epic      -> ResearchIx tiers epic GoldenEgg

deriving instance Show (Sum Finite tiers) => Show (ResearchIx tiers epic a)

_RICommon :: Iso (ResearchIx t1 epic Bock) (ResearchIx t2 epic Bock)
                 (Sum Finite t1          ) (Sum Finite t2          )
_RICommon = iso (\case RICommon i -> i) RICommon

_RIEpic :: Iso (ResearchIx tiers e1 GoldenEgg) (ResearchIx tiers e2 GoldenEgg)
               (Finite e1                    ) (Finite e2                    )
_RIEpic = iso (\case RIEpic i -> i) RIEpic

bonusAmountParseOptions :: Options
bonusAmountParseOptions = defaultOptions
    { sumEncoding = TaggedObject
                      { tagFieldName      = "type"
                      , contentsFieldName = "value"
                      }
    , constructorTagModifier = camelTo2 '-' . drop 2
    }

instance FromJSON BonusAmount where
    parseJSON  = genericParseJSON  bonusAmountParseOptions
instance ToJSON BonusAmount where
    toJSON     = genericToJSON     bonusAmountParseOptions
    toEncoding = genericToEncoding bonusAmountParseOptions

bonusTypeParseOptions :: Options
bonusTypeParseOptions = defaultOptions
    { sumEncoding = UntaggedValue
    , constructorTagModifier = camelTo2 '-' . drop 2
    }

instance FromJSON BonusType where
    parseJSON  = genericParseJSON  bonusTypeParseOptions
instance ToJSON BonusType where
    toJSON     = genericToJSON     bonusTypeParseOptions
    toEncoding = genericToEncoding bonusTypeParseOptions
instance FromJSONKey BonusType where
    fromJSONKey = FromJSONKeyTextParser (parseJSON . String)
instance ToJSONKey BonusType where
    toJSONKey = toJSONKeyText $
        T.pack . constructorTagModifier bonusTypeParseOptions . show

instance Semigroup Bonuses where
    (<>) = mappend

instance Monoid Bonuses where
    mempty = Bonuses M.empty
    mappend (Bonuses x) (Bonuses y) =
      Bonuses (M.unionWith (++) x y)

instance (FromJSON a, RealFrac a) => FromJSON (Research a) where
    parseJSON = withObject "Research" $ \v ->
        Research <$> v .: "name"
                 <*> v .: "description"
                 <*> (mkBase =<< v .: "bonuses")
                 <*> (mkCosts <$> (v .:? "costs") <*> (v .:? "levels"))
      where
        mkBase :: M.Map BonusType Object -> Parser Bonuses
        mkBase = fmap (Bonuses . fmap (:[])) . traverse (.: "base-amount")
        mkCosts :: Maybe [a] -> Maybe Natural -> V.Vector (Maybe a)
        mkCosts Nothing        Nothing  = V.empty
        mkCosts Nothing        (Just l) = V.replicate (fromIntegral l) Nothing
        mkCosts (Just [] )     (Just l) = V.replicate (fromIntegral l) Nothing
        mkCosts (Just v  )     Nothing  = Just <$> V.fromList v
        mkCosts (Just [x])     (Just l) = Just x `V.cons` V.replicate (fromIntegral l - 1) Nothing
        mkCosts (Just v@(_:_)) (Just l) = Just <$> V.fromList v V.++ extras
          where
            lv = length v
            extras = V.generate (fromIntegral l - lv) $ \i -> realToFrac . exp $
              α + β * fromIntegral (i + lv)
            (α, β) = linearRegression (V.fromList $ zipWith const [0..] v)
                                      (V.fromList (log . realToFrac <$> v))

instance ToJSON a => ToJSON (Research a) where
    toJSON Research{..} = object $
        [ "name"        .= _rName
        , "description" .= _rDescription
        , "bonuses"     .= object [ "base-amount" .= _bMap _rBaseBonuses ]
        ] ++ case sequence _rCosts of
            Just v  ->   [ "costs"  .= v ]
            Nothing -> case sequence (V.takeWhile isJust _rCosts) of
              Nothing -> [ "levels" .= V.length _rCosts ]
              Just v  -> [ "costs"  .= v
                         , "levels" .= V.length _rCosts
                         ]
    toEncoding Research{..} = pairs . mconcat $
        [ "name"        .= _rName
        , "description" .= _rDescription
        , pair "bonuses" (pairs ("base-amount" .= _bMap _rBaseBonuses))
        ] ++ case sequence _rCosts of
            Just v  ->   [ "costs"  .= v ]
            Nothing -> case sequence (V.takeWhile isJust _rCosts) of
              Nothing -> [ "levels" .= V.length _rCosts ]
              Just v  -> [ "costs"  .= v
                         , "levels" .= V.length _rCosts
                         ]

researchTierParseOptions :: Options
researchTierParseOptions = defaultOptions
    { fieldLabelModifier = camelTo2 '-' . drop 3
    }

instance KnownNat n => FromJSON (ResearchTier n) where
    parseJSON  = genericParseJSON  researchTierParseOptions
instance KnownNat n => ToJSON (ResearchTier n) where
    toJSON     = genericToJSON     researchTierParseOptions
    toEncoding = genericToEncoding researchTierParseOptions
instance FromJSON SomeResearchTier where
    parseJSON = withObject "ResearchTier" $ \v -> do
      u <- v .: "unlock"
      t <- v .: "techs"
      SV.withSized t $ \tV ->
        return $ sing :=> ResearchTier u tV
instance ToJSON SomeResearchTier where
    toJSON = \case
        SNat :=> r -> toJSON r
    toEncoding = \case
        SNat :=> r -> toEncoding r
instance FromJSON SomeResearchData where
    parseJSON = withObject "ResearchData" $ \v -> do
        res   <- v .: "common"
        epics <- (fmap . fmap) (round @Double) <$> v .: "epic"
        withV res $ \resV ->
          some (withProd (dsumSome . getI) resV) $ \(_ :&: (unzipP->(resS :&: resP))) ->
            SV.withSized epics   $ \sizedV ->
              return (STuple2 (fromTC @(Prod Sing) resS) SNat :=> Uncur (ResearchData resP sizedV))
instance ToJSON SomeResearchData where
    toJSON = \case
        _ :=> Uncur r -> toJSON r
    toEncoding = \case
        _ :=> Uncur r -> toEncoding r
instance (SingI tiers, KnownNat epic) => FromJSON (ResearchData tiers epic) where
    parseJSON = withObject "ResearchData" $ \v -> do
        res   <- v .: "common"
        resV  <- go sing res
        epics <- (fmap . fmap) (round @Double) <$> v .: "epic"
        epicsV <- case SV.toSized epics of
          Nothing -> fail "Bad number of items in list."
          Just eV -> return eV
        return $ ResearchData resV epicsV
      where
        go :: Sing ts -> [Value] -> Parser (Prod ResearchTier ts)
        go = \case
          SNil -> \case
            []  -> return Ø
            _:_ -> fail "Too many items in list"
          SNat `SCons` ss -> \case
            []   -> fail "Too few items in list"
            x:xs -> (:<) <$> parseJSON x <*> go ss xs

instance ToJSON (ResearchData tiers epic) where
    toJSON ResearchData{..} = object
        [ "common" .= TCP.toList (\ResearchTier{..} ->
                          object [ "unlock" .= _rtUnlock
                                 , "techs"  .= SV.fromSized _rtTechs
                                 ]
                        ) _rdCommon
        , "epic"   .= SV.fromSized _rdEpic
        ]
    toEncoding ResearchData{..} = pairs . mconcat $
        [ "common" .= TCP.toList (\ResearchTier{..} ->
                          object [ "unlock" .= _rtUnlock
                                 , "techs"  .= SV.fromSized _rtTechs
                                 ]
                        ) _rdCommon
        , "epic"   .= SV.fromSized _rdEpic
        ]

instance (SingI tiers, KnownNat epic) => FromJSON (ResearchStatus tiers epic) where
    parseJSON = withObject "ResearchStatus" $ \v -> do
        res   <- v .: "common"
        resV  <- go sing res
        epics <- v .: "epic"
        epicsV <- case SV.toSized epics of
          Nothing -> fail "Bad number of items in list."
          Just eV -> return eV
        return $ ResearchStatus resV epicsV
      where
        go :: Sing ts -> [Value] -> Parser (Prod (Flip SV.Vector Natural) ts)
        go = \case
          SNil -> \case
            []  -> return Ø
            _:_ -> fail "Too many items in list"
          SNat `SCons` ss -> \case
            []   -> fail "Too few items in list"
            x:xs -> (:<) <$> parseJSON x <*> go ss xs
instance ToJSON (ResearchStatus tiers epic) where
    toJSON ResearchStatus{..} = object
        [ "common" .= TCP.toList (SV.fromSized . getFlip) _rsCommon
        , "epic"   .= SV.fromSized _rsEpic
        ]
    toEncoding ResearchStatus{..} = pairs . mconcat $
        [ "common" .= TCP.toList (SV.fromSized . getFlip) _rsCommon
        , "epic"   .= SV.fromSized _rsEpic
        ]
instance FromJSON SomeResearchStatus where
    parseJSON = withObject "ResearchStatus" $ \v -> do
        res   <- v .: "common"
        epics <- v .: "epic"
        withV res $ \resV ->
          some (withProd (go . getI) resV) $ \(_ :&: (unzipP->(resS :&: resP))) ->
            SV.withSized epics   $ \sizedV ->
              return (STuple2 (fromTC resS) SNat :=> Uncur (ResearchStatus resP sizedV))
      where
        go :: V.Vector Natural -> Some (Sing :&: Flip SV.Vector Natural)
        go v = SV.withSized v $ \u -> Some (SNat :&: Flip u)
instance ToJSON SomeResearchStatus where
    toJSON = \case
        _ :=> Uncur r -> toJSON r
    toEncoding = \case
        _ :=> Uncur r -> toEncoding r

-- | Scales a bonus amount by a "level".
scaleAmount :: Natural -> BonusAmount -> BonusAmount
scaleAmount n = \case
    BAIncrement  i -> BAIncrement  (fromIntegral n * i)
    BAPercent    p -> BAPercent    (fromIntegral n * p)
    BAMultiplier r -> BAMultiplier (r ^ n)

-- | Tells whether or not a bonus amount creates any actual effect.
hasEffect :: BonusAmount -> Bool
hasEffect = \case
    BAIncrement  i -> i /= 0
    BAPercent    p -> p /= 0
    BAMultiplier r -> r /= 1

-- | A "smart constructor" that clears bonuses with no effect.
bonuses :: M.Map BonusType [BonusAmount] -> Bonuses
bonuses = Bonuses . M.filter (not . null) . fmap (filter hasEffect)

-- | Apply a list of bonuses (from left to right) to a value.
bonusEffect :: Fractional a => [BonusAmount] -> a -> a
bonusEffect = flip . foldl' $ \e -> \case
    BAIncrement  i -> e + realToFrac i
    BAPercent    p -> e * (1 + realToFrac p / 100)
    BAMultiplier r -> e * realToFrac r

-- | Iso on a value under a list of bonuses (from left to right).
bonusing :: (Fractional a, Eq a) => [BonusAmount] -> Iso' a a
bonusing = flip foldl' id $ \e -> \case
             BAIncrement i  -> e . adding (realToFrac i)
             BAPercent   p  -> e . multiplying (1 + realToFrac p / 100)
             BAMultiplier r -> e . multiplying (realToFrac r)

-- | Iso on a value under the bonuses of a given bonus type.
bonusingFor
    :: (Fractional a, Eq a)
    => Bonuses
    -> BonusType
    -> Iso' a a
bonusingFor bs bt = case M.lookup bt (_bMap bs) of
                      Nothing -> id
                      Just bl -> bonusing bl

-- | Apply bonuses for a given type on a value.
bonusEffectFor :: Fractional a => Bonuses -> BonusType -> a -> a
bonusEffectFor bs bt = maybe id bonusEffect $ M.lookup bt (_bMap bs)

-- | No bonuses
emptyBonuses :: Bonuses
emptyBonuses = Bonuses M.empty

-- | Maximum level for a given research.
maxLevel :: Research a -> Natural
maxLevel = fromIntegral . V.length . _rCosts

-- | Empty status (no research).
emptyResearchStatus :: ResearchData tiers epic -> ResearchStatus tiers epic
emptyResearchStatus ResearchData{..} =
    ResearchStatus (map1 clear _rdCommon) (0 <$ _rdEpic)
  where
    clear :: ResearchTier a -> Flip SV.Vector Natural a
    clear = Flip . set mapped 0 . view rtTechs

-- | Zips together all research data and research statuses into an
-- accumulator.
foldResearch
    :: Monoid b
    => (Either (Research Bock) (Research GoldenEgg) -> Natural -> b)
    -> ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> b
foldResearch f ResearchData{..} ResearchStatus{..} = mconcat
    [ foldMap1 (\case d :&: Flip s -> fold $ SV.zipWith (f . Left) (_rtTechs d) s
               )
        (zipP (_rdCommon :&: _rsCommon))
    , fold $ SV.zipWith (f . Right) _rdEpic _rsEpic
    ]

-- | Bonuses from a given 'Research' at a given level.  Assumes no "maximum
-- level".
researchBonuses :: Research a -> Natural -> Bonuses
researchBonuses _ 0 = Bonuses M.empty
researchBonuses r l = _rBaseBonuses r & _Bonuses . mapped . mapped %~ scaleAmount l

-- | Total bonuses from a given 'ResearchStatus'.
totalBonuses :: ResearchData tiers epic -> ResearchStatus tiers epic -> Bonuses
totalBonuses = foldResearch (either researchBonuses researchBonuses)

-- | How many common techs have been researched?
--
-- Used for opening tiers.
researchCount :: ResearchStatus tiers epic -> Natural
researchCount = sumOf $ rsCommon . liftTraversal (_Flip . folded)

legalTiers
    :: forall tiers epic. ()
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> Prod (C Bool) tiers
legalTiers rd rs = rd ^. rdCommon . to (map1 go)
  where
    tot = researchCount rs
    go  :: ResearchTier a
        -> C Bool a
    go = view $ rtUnlock . to (>= tot) . _Unwrapped

data ResearchError = REMaxedOut
                   | RELocked
  deriving (Show, Eq, Ord)

makePrisms ''ResearchError

-- | Purchase research at a given index, incrementing the counter in the
-- 'ResearchStatus'.  Returns 'Nothing' if research is already maxed out.
purchaseResearch
    :: forall tiers epic a. (KnownNat epic, SingI tiers)
    => ResearchData tiers epic
    -> ResearchIx tiers epic a
    -> ResearchStatus tiers epic
    -> Either ResearchError (a, ResearchStatus tiers epic)
purchaseResearch rd i rs0 = pp . getComp . researchIxStatusLegal rd i (Comp . go) $ rs0
  where
    bs = totalBonuses rd rs0
    pp :: Maybe (First a, ResearchStatus tiers epic)
        -> Either ResearchError (a, ResearchStatus tiers epic)
    pp Nothing                     = Left REMaxedOut
    pp (Just (First Nothing , _ )) = Left RELocked
    pp (Just (First (Just c), rs)) = case i of
        RICommon _ -> Right (c ^. bonusingFor bs BTResearchCosts, rs)
        RIEpic   _ -> Right (c, rs)
    go :: Natural -> Maybe (First a, Natural)
    go currLevel = first (First . Just) <$>
      rd ^? researchIxData i
          . rCosts
          . ix (fromIntegral currLevel)
          . to ((, currLevel + 1) . fromMaybe (0 \\ researchIxNum i))

-- | Get a 'Num' instance from a 'ResearchIx'.
researchIxNum :: ResearchIx tiers epic a -> Wit (Num a)
researchIxNum = \case
    RICommon _ -> Wit
    RIEpic   _ -> Wit

-- | A lens into a 'ResearchData', given the appropriate index.
researchIxData
    :: (KnownNat epic, SingI tiers)
    => ResearchIx tiers epic a
    -> Lens' (ResearchData tiers epic) (Research a)
researchIxData = \case
    RICommon i -> \f ->
      let g :: forall a. _ a -> _ (_ a)
          g x@(slot :&: (_ :&: SNat)) = (_2 . _1 . rtTechs . ixSV slot) f x
      in  rdCommon $ \rs -> map1 fanFst . fanSnd
            <$> sumProd g (i :&: zipP (rs :&: toTC sing))
    RIEpic i   -> rdEpic . ixSV i

-- | A lens into a 'ResearchStatus', given the appropriate index.
researchIxStatus
    :: (KnownNat epic, SingI tiers)
    => ResearchIx tiers epic a
    -> Lens' (ResearchStatus tiers epic) Natural
researchIxStatus = \case
    RICommon i -> \f ->
      let g :: forall a. _ a -> _ (_ a)
          g x@(slot :&: (_ :&: SNat)) = (_2 . _1 . _Flip . ixSV slot) f x
      in  rsCommon $ \rs -> map1 fanFst . fanSnd
            <$> sumProd g (i :&: zipP (rs :&: toTC sing))
    RIEpic i   -> rsEpic . ixSV i

-- | Lens into both the research data and research status common slots
-- together.
rdrsCommon
    :: Lens' ((Uncur ResearchData :&: Uncur ResearchStatus) (tiers # epic))
             (Prod (ResearchTier :&: Flip SV.Vector Natural) tiers)
rdrsCommon f (Uncur rd :&: Uncur rs) =
    f (zipP (_rdCommon rd :&: _rsCommon rs)) <&> \(unzipP->(rdc :&: rsc)) ->
      Uncur (rd & rdCommon .~ rdc) :&: Uncur (rs & rsCommon .~ rsc)

-- | Lens into both the research data and research status epic slots
-- together.
rdrsEpic
    :: KnownNat epic
    => Lens' ((Uncur ResearchData :&: Uncur ResearchStatus) (tiers # epic))
             (SV.Vector epic (Research GoldenEgg, Natural))
rdrsEpic f (Uncur rd :&: Uncur rs) =
    f (liftA2 (,) (_rdEpic rd) (_rsEpic rs)) <&> \rdrs ->
      Uncur (rd & rdEpic .~ fmap fst rdrs) :&: Uncur (rs & rsEpic .~ fmap snd rdrs)

-- | Lens into both a ResearchData and ResearchStatus together, from
-- a given index.
researchIx
    :: (SingI tiers, KnownNat epic)
    => ResearchIx tiers epic a
    -> Lens' ((Uncur ResearchData :&: Uncur ResearchStatus) (tiers # epic))
             (Research a, Natural)
researchIx = \case
    RICommon i -> \f ->
      let g :: forall a. _ a
            -> _ (_ a)
          g (slot :&: ((c :&: Flip e) :&: SNat)) =
              f (SV.index (_rtTechs c) slot, SV.index e slot) <&> \(d, s) ->
                let c' = c & rtTechs . ixSV slot .~ d
                    e' = e & ixSV slot .~ s
                in  slot :&: ((c' :&: Flip e') :&: SNat)
      in  rdrsCommon $ \rdrs ->
            map1 fanFst . fanSnd <$> sumProd g (i :&: zipP (rdrs :&: toTC sing))
    RIEpic i -> rdrsEpic . ixSV i

withSomeResearch
    :: DSum Sing (Uncur res)
    -> (forall tiers epic. (KnownNat epic, SingI tiers) => res tiers epic -> r)
    -> r
withSomeResearch = \case
    STuple2 sTs SNat :=> Uncur r -> \f -> withSingI sTs $ f r

-- | Traversal into a given index of a 'ResearchStatus' if the item is in
-- a legal tier.
--
-- Only a legal traversal if the mapping function doesn't change the legal
-- status.
researchIxStatusLegal
    :: (SingI tiers, KnownNat epic)
    => ResearchData tiers epic
    -> ResearchIx tiers epic a
    -> Traversal' (ResearchStatus tiers epic) Natural
researchIxStatusLegal rd = \case
    RICommon i -> \f rs0 -> getUncur . fanSnd <$>
      let totCount = researchCount rs0
          g :: forall a. _ a -> _ (_ a)
          g x@(slot :&: ((rt :&: _) :&: SNat))
            | _rtUnlock rt <= totCount = (_2 . _1 . _2 . _Flip . ixSV slot) f x
            | otherwise                = pure x
      in  (Uncur rd :&: Uncur rs0) & rdrsCommon %%~ \rtrs ->
            map1 fanFst . fanSnd <$> sumProd g (i :&: zipP (rtrs :&: toTC sing))
    RIEpic i   -> rsEpic . ixSV i

-- | All 'ResearchIx' into common researches.
researchIxesCommon
    :: SingI tiers
    => Prod (Flip SV.Vector (ResearchIx tiers epic Bock)) tiers
researchIxesCommon = go sing
  where
    go :: Sing ts -> Prod (Flip SV.Vector (ResearchIx ts epic Bock)) ts
    go = \case
      SNil         -> Ø
      SNat `SCons` ss ->
        let rest = map1 (over (_Flip . mapped . _RICommon) InR) (go ss)
        in  Flip (SV.generate (RICommon . InL)) :< rest

-- | All 'ResearchIx' into epic researches.
researchIxesEpic
    :: KnownNat epic
    => SV.Vector epic (ResearchIx tiers epic GoldenEgg)
researchIxesEpic = SV.generate RIEpic

-- | All legal 'ResearchIx' for common research.
legalResearchIxesCommon
    :: forall tiers epic. SingI tiers
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> Prod (Flip SV.Vector (Either ResearchError (ResearchIx tiers epic Bock))) tiers
legalResearchIxesCommon rd rs =
    imap1 (\i -> Flip . go i) (zipP (zipP (_rdCommon rd :&: _rsCommon rs) :&: toTC sing))
  where
    totCount = researchCount rs
    go  :: forall t. ()
        => Index tiers t
        -> ((ResearchTier :&: Flip SV.Vector Natural) :&: Sing) t
        -> SV.Vector t (Either ResearchError (ResearchIx tiers epic Bock))
    go i ((rt :&: Flip c) :&: SNat)
      | rt ^. rtUnlock > totCount = pure (Left RELocked)
      | otherwise                 =
          let mkIx
                  :: Finite t
                  -> Research Bock
                  -> Natural
                  -> Either ResearchError (ResearchIx tiers epic Bock)
              mkIx j r n
                | n < maxLevel r = Right . RICommon . someSum $ Some (i :&: j)
                | otherwise      = Left REMaxedOut
          in  SV.izipWith mkIx (rt ^. rtTechs) c

-- | All legal 'ResearchIx' for epic research.
--
-- 'Nothing' implies research is maxed-out.
legalResearchIxesEpic
    :: forall tiers epic. KnownNat epic
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> (SV.Vector epic :.: Maybe) (ResearchIx tiers epic GoldenEgg)
legalResearchIxesEpic rd rs = Comp $ SV.izipWith go (_rdEpic rd) (_rsEpic rs)
  where
    go  :: Finite epic
        -> Research GoldenEgg
        -> Natural
        -> Maybe (ResearchIx tiers epic GoldenEgg)
    go i r n
      | n < maxLevel r = Just . RIEpic $ i
      | otherwise      = Nothing

-- | All legal common researches.
legalResearchesCommon
    :: forall tiers epic. SingI tiers
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> Prod (Flip SV.Vector (Either ResearchError Bock)) tiers
legalResearchesCommon rd rs =
    map1 (Flip . go) (zipP (zipP (_rdCommon rd :&: _rsCommon rs) :&: toTC sing))
  where
    bs       = totalBonuses rd rs
    totCount = researchCount rs
    go  :: forall t. ()
        => ((ResearchTier :&: Flip SV.Vector Natural) :&: Sing) t
        -> SV.Vector t (Either ResearchError Bock)
    go ((rt :&: Flip c) :&: SNat)
      | rt ^. rtUnlock > totCount = pure (Left RELocked)
      | otherwise                 = do
          r         <- rt ^. rtTechs
          currLevel <- c
          pure $ case r ^? rCosts . ix (fromIntegral currLevel) of
            Nothing -> Left REMaxedOut
            Just b  -> Right $ fromMaybe 0 b ^. bonusingFor bs BTResearchCosts

-- | All legal epic researches.
--
-- 'Nothing' implies research is maxed-out.
legalResearchesEpic
    :: forall tiers epic. KnownNat epic
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> (SV.Vector epic :.: Maybe) GoldenEgg
legalResearchesEpic rd rs = Comp $ do
    r         <- _rdEpic rd
    currLevel <- _rsEpic rs
    pure $ fromMaybe 0 <$> r ^? rCosts . ix (fromIntegral currLevel)


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE TemplateHaskell   #-}
{-# LANGUAGE TupleSections     #-}

module Topical.Text.Tokenizer.TextTiling
    ( SmoothingParams(..)

    , textTilingTokenizer
    , suppressSmallBlocks
    , smooth

    , windows
    , partition
    , triples
    , frequencies
    , blockComparison
    , vocabularyIntroduction

    , Sequence(..)
    , seqNo
    , seqSpan
    , seqItems
    , seqFreqs

    , SeqScore
    ) where

import           Control.Arrow
import           Control.Error
import           Control.Lens
-- import           Data.Foldable       hiding (concat)
import           Data.Hashable
import qualified Data.HashMap.Strict as M
import qualified Data.HashSet        as S
import qualified Data.List           as L
import qualified Data.List.Split     as Split
import           Data.Monoid
import           Data.Ord
import qualified Data.Text           as T
import           Data.Traversable
-- import qualified Data.Vector         as V
-- import           Statistics.Sample

-- import           Topical.Text.Types

type Frequency         = Int
-- type TokenSeqFrequency = (SeqNumber, Frequency)

-- | "The morpholocially-analyzed token is stored in a table along with
-- a record of the token-windowSeq number it occurred in, and the number of
-- times it appeared in the token-sequence."
-- type TokenTable a      = M.HashMap a [TokenSeqFrequency]

type SeqNumber = Int

-- | A record is also kept of the locations of the paragraph breaks within
-- the text. Stop words contribute to the computation of the size of the
-- token windowSeq, but not to the computation of the similarity between
-- blocks of text.
type SeqSpan = (Int, Int)

type FrequencyTable a = M.HashMap a Frequency

data Sequence a b = Sequence
                  { _seqNo    :: !SeqNumber
                  , _seqSpan  :: !SeqSpan
                  , _seqItems :: ![a]
                  , _seqFreqs :: !(FrequencyTable b)
                  } deriving (Eq)
makeLenses ''Sequence

instance Show (Sequence T.Text b) where
    show Sequence{..} =
        mconcat [ "Sequence #"
                , show _seqNo
                , ": "
                , T.unpack (T.unwords _seqItems)
                ]

type TokenSequence a = Sequence a a
type BlockSequence a = Sequence (TokenSequence a) a

type SeqScore a b = (TokenSequence a, b)

data SmoothingParams  = Smoothing Int    -- ^ Number of rounds of smoothing
                                  Int    -- ^ Smoothing window size

type ScoringFunction a =
    TokenSeqSize -> BlockSize -> [TokenSequence a] -> [SeqScore a Double]
type CountFunction a b = (Hashable b, Eq b) => [a] -> FrequencyTable b

type TokenSeqSize = Int
type BlockSize    = Int

-- | The basic processing flow:
-- boundary identification . lexical score determination . tokenization
textTilingTokenizer :: (Hashable a, Eq a, Show a)
                    => TokenSeqSize
                    -- ^ The token windowSeq size parameter (/w/). 20 is often
                    -- a good default value.
                    -> BlockSize
                    -- ^ The number of token sequences to group together.
                    -- This is the average paragraph length (in token
                    -- sequences) (/blocksize/). 6 often works well.
                    -> ScoringFunction a
                    -- ^ This either uses block comparison to look at
                    -- shared vocabulary or vocabulary introduction to base
                    -- it on the number of new words are found in the
                    -- second window. Either way, larger output values
                    -- should indicate a more likely change of topic.
                    -> SmoothingParams
                    -> [a]
                    -- ^ The input sequence. If this is text, this should
                    -- be tokenized, case-folded, and stop-words filtered
                    -- out. Also, affixes and irregular forms normalized
                    -- should be removed so that it's only the
                    -- morphological base.
                    -> [BlockSequence a]
textTilingTokenizer w k scoring (Smoothing sWin sIters) =
      cutOffs
    . smooth sWin sIters
    . snd
    . mapAccumConcat boundary (0, 0, [])
    . scoring w k
    . partitionSeq w frequencies

cutOffs :: [SeqScore a (Double, Double)] -> [BlockSequence a]
cutOffs = undefined

-- The first Double is the raw score. The second is the depth score
-- actually used to make the boundary determination.
-- type SeqNode a = SeqScore a (Double, Double)

-- | This takes a list of @SeqScore a (Double, Double)@, sorted by
-- likelihood of beginning a new section, and returns a tree of descending
-- likelihoods, but in sequence order.
-- hangTree :: Show a => [SeqNode a] -> Tree (SeqNode a)
-- hangTree = unfoldTree hang
--     where
--         hang :: Show a
--                 => [SeqNode a]
--                     -> (SeqNode a, (Maybe [SeqNode a], Maybe [SeqNode a]))
--         hang = undefined
--              . maximumBy (comparing (snd . snd . fst))
--              . snd
--              . L.mapAccumL packageNode []
--              . filter (not . L.null)
--              . L.tails
--
--         packageNode :: Show a
--                        => [SeqNode a]
--                            -> [SeqNode a]
--                            -> ([SeqNode a], (SeqNode a, [[SeqNode a]]))
--         packageNode ls (s:rs) = (s:ls, (s, [reverse ls, rs]))
--         packageNode ls []     = (ls,   (undefined, []))
--         -- It should never reach here, so I'll just plant a bomb.
--         -- KA-BOOM!

smooth :: Int -> Int -> [SeqScore a (Double, Double)]
       -> [SeqScore a (Double, Double)]
smooth _      0     sss = sss
smooth window iters sss =
    smooth window (pred iters) $ moveAvgBy window (_2 . _2) sss

moveAvgBy :: Int       -- ^ The number of items to offset from each side
                       -- of the center. For instance, for a window of 3,
                       -- use a value of 1.
          -> Lens' a Double -> [a] -> [a]
moveAvgBy n l xs = zipWith (set l) (moveavg n $ map (^. l) xs) xs

moveavg :: Int         -- ^ The number of items to offset from each side
                       -- of the center. For instance, for a window of 3,
                       -- use a value of 1.
        -> [Double] -> [Double]
moveavg n = go []
    where
        go _    []     = []
        go pref (x:xs) = avg (x : take n pref ++ take n xs) : go (x:pref) xs

avg :: [Double] -> Double
avg = uncurry (/) . L.foldl' accum (0, 0)
    where
        accum p x = ((+ x) *** succ) p

suppressSmallBlocks :: Int
                    -> [SeqScore a (Double, Double)]
                    -> [SeqScore a (Double, Double)]
suppressSmallBlocks minP = sortNo . removeTiny minP . sortScore
    where
        removeTiny :: Int
                   -> [SeqScore a0 (Double, Double)]
                   -> [SeqScore a0 (Double, Double)]
        removeTiny _ []     = []
        removeTiny w (x:xs) =
            x : sortScore (map (modifyClose w $ x ^. _1 . seqNo) xs)
        closeTo :: Int -> Int -> SeqScore a1 (Double, Double) -> Bool
        closeTo w seq1No seq2 = abs (seq1No - (seq2 ^. _1 . seqNo)) <= w
        modifyClose :: Int
                    -> Int
                    -> SeqScore a2 (Double, Double)
                    -> SeqScore a2 (Double, Double)
        modifyClose w seq1No seq2 = if closeTo w seq1No seq2
                                        then seq2 & _2 . _2 .~ 0
                                        else seq2

sortScore :: [SeqScore a (Double, Double)] -> [SeqScore a (Double, Double)]
sortScore = L.sortBy (comparing (Down . snd . snd))

sortNo :: [SeqScore a b] -> [SeqScore a b]
sortNo = L.sortBy (comparing (_seqNo . fst))

type BoundaryId a = (Double, Double, [Double -> SeqScore a (Double, Double)])

boundary :: BoundaryId a
         -> SeqScore a Double
         -> (BoundaryId a, [SeqScore a (Double, Double)])
boundary (leftS, lastS, pending) (ts, s)
    | lastS <= s = ((leftS', s, nextf:pending), [])
    | otherwise  = ((s, s, []), reverse (map ($ lastS) (nextf:pending)))
    where
        leftS'  = max leftS s
        nextf r = (ts, (s, (leftS' - s) + (r - s)))

-- | This converts a two-item list into a tuple pair.
toPair :: [a] -> Maybe (a, a)
toPair [a, b] = Just (a, b)
toPair _      = Nothing

-- toTriples :: [a] -> Maybe (a, a, a)
-- toTriples [a, b, c] = Just (a, b, c)
-- toTriples _         = Nothing

mapAccumConcat :: (a -> b -> (a, [c])) -> a -> [b] -> (a, [c])
mapAccumConcat f s = fmap concat . mapAccumL f s

-- TODO: I think this should be a comonad, but right now I just want to get
-- it written.
-- mapAccumContext :: (s -> [x] -> x -> [x] -> (s, [a])) -> s -> [x] -> (s, [a])
-- mapAccumContext f state = go state []
--     where
--         go s _  []     = (s, [])
--         go s ls (x:xs) = let (s',  as)  = f s ls x xs
--                              (s'', ass) = go s' (x:ls) xs
--                          in  (s'', as ++ ass)

-- | The block comparison scoring function. This looks at words in common
-- between two windows.
-- TODO: I think we're loosing `bsize` token sequences off the front.
blockComparison :: (Hashable a, Eq a) => ScoringFunction a
blockComparison _ _ []      = []
blockComparison _ bsize tss =
      uncurry appHead
    . (floatHead &&& comparePairs)
    $ windowSeq bsize mergeFrequencies tss
    where
        floatHead seqs = (,0.0) <$> seqs ^? i0 . seqItems . i0
        comparePairs   = map (uncurry blockCompare) . mapMaybe toPair . windows 2
        appHead mh xs  = maybe xs (:xs) mh
        i0             = traversed . index 0

blockCompare :: (Hashable a, Eq a)
             => BlockSequence a -> BlockSequence a -> SeqScore a Double
blockCompare (Sequence _ _ _ b1) (Sequence _ _ (ts:_) b2) =
    (ts,) . final . foldMap step $ terms b1 `S.union` terms b2
    where
        topf w1 w2     = w1 * w2
        bottoml w1 _w2 = w1 * w1
        bottomr _w1 w2 = w2 * w2
        getFreqs t     = (M.lookupDefault 0 t b1, M.lookupDefault 0 t b2)
        step           = (   Sum . uncurry topf
                         &&& Sum . uncurry bottoml
                         &&& Sum . uncurry bottomr
                         ) . getFreqs
        final (Sum t, (Sum bl, Sum br)) =
            fromIntegral t / sqrt (fromIntegral $ bl * br)
blockCompare _ _ = undefined

-- | This returns the terms in a frequency table.
terms :: (Hashable a, Eq a) => FrequencyTable a -> S.HashSet a
terms = S.fromList . M.keys

-- | The vocabulary introduction scoring function. This looks at how many
-- words are introduced in the second window.
-- TODO: What happens to the first item from the first pair in `collapse`?
-- TODO: Are we looking a token sequence off the front?
vocabularyIntroduction :: (Hashable a, Eq a) => ScoringFunction a
vocabularyIntroduction tsSize _ = snd
                                . mapAccumL accum S.empty
                                . map (uncurry collapse)
                                . mapMaybe toPair
                                . windows 2
                                . map (id &&& (terms . _seqFreqs))
    where
        blockSize :: Double
        blockSize = fromIntegral tsSize * 2.0
        collapse (_, t1) (ts, t2) = (ts, t1 `S.union` t2)
        accum s (ts, block) =
            let unseen = S.size $ block `S.difference` s
                s'     = s `S.union` block
            in  (s', (ts, fromIntegral unseen / blockSize))

-- | This extends @windows@ to create @Sequence@ instances.
windowSeq :: (Eq b, Hashable b)
          => Int -> CountFunction a b -> [a] -> [Sequence a b]
windowSeq k = toSeq (windows k)

-- | This extends @partition@ to create @Sequence@ data.
partitionSeq :: (Eq b, Hashable b)
             => Int -> CountFunction a b -> [a] -> [Sequence a b]
partitionSeq k = toSeq (partition k)

toSeq :: (Eq b, Hashable b)
      => ([(Int, a)] -> [[(Int, a)]])
      -> CountFunction a b -> [a] -> [Sequence a b]
toSeq breaker counter =
    mapMaybe (uncurry toSeq') . zip [0..] . breaker . zip [0..]
    where
        toSeq' n spanned =
            let items = map snd spanned
            in  Sequence n <$> getSpan spanned
                           <*> pure items
                           <*> pure (counter items)
        getSpan []          = Nothing
        getSpan [(s, _)]    = Just (s, s)
        getSpan ((s, _):xs) = (s,) . fst <$> lastZ xs

-- | This takes an input windowSeq and divides it into overlapping
-- subsequences of a given size.
--
-- >>> windows 2 [1, 2, 3, 4, 5]
-- [[1, 2], [2, 3], [3, 4], [4, 5], [5]]
windows :: Int      -- ^ The size of the sliding window.
        -> [a]      -- ^ The input sequence.
        -> [[a]]    -- ^ The output windowSeq of windows.
windows k = filter (not . L.null) . map (take k) . L.tails

-- | This implementation is taken from Data.List.Split.Lens, before it was
-- removed.
chunking :: Int -- ^@n@
         -> Getting (Endo [a]) s a -> Fold s [a]
chunking s l f = coerce . traverse f . Split.chunksOf s . toListOf l
{-# INLINE chunking #-}

-- | This takes an input windowSeq and divides it into non-overlapping
-- partitions of a given size.
--
-- >> partition 2 [1, 2, 3, 4, 5]
-- [[1, 2], [3, 4], [5]]
partition :: Int    -- ^ The size of the non-overlapping partitions.
          -> [a]    -- ^ The input windowSeq to partition.
          -> [[a]]  -- ^ The non-overlapping partitions.
partition k xs = xs ^.. chunking k each

-- | This takes an input windowSeq and divides it into overlapping chains of
-- three.
--
-- >>> triples [1, 2, 3, 4, 5]
-- [(1, 2, 3), (2, 3, 4), (3, 4, 5)]
triples :: [a] -> [(a, a, a)]
triples = mapMaybe triples' . L.tails
    where
        triples' (a:b:c:_) = Just (a, b, c)
        triples' _         = Nothing

frequencies :: CountFunction a a
frequencies xs = M.fromListWith (+) . zip xs $ L.repeat 1

mergeFrequencies :: CountFunction (Sequence a b) b
mergeFrequencies = foldMap _seqFreqs

-- cosSimilarity :: (Hashable a, Eq a)
--               => FrequencyTable a -> FrequencyTable a -> Double
-- cosSimilarity b1 b2 =
--     finish . foldMap step $ terms b1 `S.union` terms b2
--     where step k = let w1 = M.lookupDefault 0 k b1
--                        w2 = M.lookupDefault 0 k b2
--                     in (Sum (w1 * w2), Sum (w1 * w1), Sum (w2 * w2))
--           finish (Sum a, Sum b, Sum c) =
--               fromIntegral a / sqrt (fromIntegral b * fromIntegral c)

-- | This calculates all boundaries as being places where the depth score
-- is greater than the function below (mean minus half the standard
-- deviation).
--
-- If I need better performance, I can use
-- http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online_algorithm

-- boundaries :: [Double] -- ^ The cutoff scores between each paragraph block.
--            -> [Int]    -- ^ A list of the paragraph indexes that are boundaries.
-- boundaries depths = undefined $ mean depths' - (fastStdDev depths' / 2.0)
--     where depths' = V.fromList depths

-- Stuff from Main for this:
-- main :: IO ()
-- main = do
--     (stopfile:files) <- getArgs
--
--     stoplist <- S.fromList . tokenize <$> TIO.readFile stopfile
--
--     forM_ files $ \filename -> do
--         putStrLn filename
--
--         tokens <-  filter (not . (`S.member` stoplist)) . tokenize
--                <$> TIO.readFile filename
--         putStrLn $ "Token count = " ++ show (length tokens)
--
--         let tiles = map _seqItems
--                   $ textTilingTokenizer 20 6 blockComparison (Smoothing 1 2) tokens
--         putStrLn $ "Tile count  = " ++ show (length tiles)
--
--         forM_ tiles $ TIO.putStr
--                     . (<> "\n\n")
--                     . T.intercalate "\n\n"
--                     . map ((<> ".") . T.intercalate " " . _seqItems)
--
-- {-
--  -         chart tiles
--  -
--  -         forM_ tiles $ \(Sequence{..}, (raw, score)) ->
--  -             putStrLn $ L.intercalate "," [ show _seqNo
--  -                                          , show (fst _seqSpan)
--  -                                          , show (snd _seqSpan)
--  -                                          , show raw
--  -                                          , show score
--  -                                          , unwords (map T.unpack _seqItems)
--  -                                          ]
--  -}
--
-- chart :: [SeqScore T.Text (Double, Double)] -> IO ()
-- chart seqs = toFile def "tiles.png" $ do
--     layout_title .= "TextTiling scores"
--     -- layout_y_axis . laxis_override .= axisGridHide
--     -- plot $ line "raw"   [map (fst . _seqSpan *** fst) seqs]
--     plot $ line "score" [map (fst . _seqSpan *** snd) seqs]
--
-- {-
--  - toParagraphs :: Int
--  -              -> Tree (SeqScore T.Text (Double, Double))
--  -              -> [[SeqScore T.Text (Double, Double)]]
--  - toParagraphs 0 tree = [L.sortBy (comparing (_seqNo . fst)) $ flatten tree]
--  - toParagraphs n (Node root forest) = [sortp pre, root : sortp post]
--  -     where
--  -         rootNo      = root ^. _1 . seqNo
--  -         (pre, post) = L.break ((< rootNo) . _seqNo . fst . head . head)
--  -                     $ map (toParagraphs (n - 1)) forest
--  -         sortp       = L.sortBy (comparing (_seqNo . fst . _))
--  -}
--
-- showp :: [SeqScore T.Text (Double, Double)] -> String
-- showp = L.intercalate ". " . map (unwords . map T.unpack . _seqItems . fst)
--
-- tshow :: Show a => a -> T.Text
-- tshow = T.pack . show
--
-- word :: Parser T.Text
-- word = T.pack <$> many1 (satisfy isAlphaNum)
--
-- tokenize :: T.Text -> [T.Text]
-- tokenize = map T.toLower . parserTokenizer word


{-# LANGUAGE GADTs #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FunctionalDependencies #-}

module Dynamical.Sim.Internal where

import Control.Applicative
import Control.DeepSeq (force, NFData(..))
import Control.Monad
import Control.Monad.Fix (MonadFix(mfix))
import Control.Monad.State (StateT, State, get, put, runState, modify, execState, evalState)
import Control.Monad.Trans (lift)
import Control.Monad.Writer (Writer, tell, WriterT, runWriterT)
import Control.Parallel (pseq)
import qualified Control.Parallel.Strategies as Par
import Data.Complex (Complex((:+)))
import Data.Default.Class (Default(def))
import Data.Fixed (Fixed, mod')
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as Set
import Data.List
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
import Data.Monoid (Monoid(), Any(Any), (<>), mempty, mappend)
import Data.Proxy (Proxy(Proxy))
import Data.Scientific (Scientific)
import Data.StableMemo (memo)
import Data.Time (getCurrentTime, diffUTCTime, UTCTime)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Debug.Trace as Dbg
import GHC.Generics (Generic)
import qualified GHC.Prim as Prim
import GHC.TypeLits (Symbol, symbolVal, KnownSymbol)
import qualified Graphics.Rendering.Chart.Backend.Cairo as Chart
import qualified Graphics.Rendering.Chart.Easy as Chart
import Linear (V1(..), V2(..), angle, V3(..), V4(..), norm, _x, _y,_z,_w, qd,Additive(..),signorm,(*^),(^/), perp)
import System.IO.Unsafe (unsafePerformIO, unsafeInterleaveIO)
import System.Mem.StableName (hashStableName,eqStableName,StableName,makeStableName)
import Text.Printf (printf, PrintfArg)
import Unsafe.Coerce (unsafeCoerce)

-------------------------------------------------------------------
-- Signal
-------------------------------------------------------------------

-- | A `Signal t a` represents a time varying value of type `a`, where
-- time is measured using a `t`.
data Signal t a where
    SPure   :: a -> Signal t a
    SAp     :: Signal t (a -> b) -> Signal t a -> Signal t b
    SAdd    :: Num a => Signal t a -> Signal t a -> Signal t a
    SMul    :: Num a => Signal t a -> Signal t a -> Signal t a
    SDiv    :: Fractional a => Signal t a -> Signal t a -> Signal t a
    SMap    :: (a -> b) -> Signal t a -> Signal t b
    SInt    :: Splat t a => Int -> (a -> b) -> Signal t b
    SFn     :: Int -> (t -> a) -> Signal t a
    SSwitch :: Signal t a -> Event t (Sim t (Signal t a)) -> Signal t a
    SShare  :: Signal t a -> Signal t a

instance Monoid a => Monoid (Signal t a) where
    mempty = pure mempty
    mappend = liftA2 mappend

instance Functor (Signal t) where
    fmap f (SMap g s) = SMap (f . g) s
    fmap f (SInt i g) = SInt i (f . g)
    fmap f (SPure a) = SPure (f a)
    -- fmap f (SAp g a) = SAp (fmap (f.) g) a
    fmap f (SFn i g) = SFn i (f . g)
    -- fmap f (SSwitch s e) = SSwitch (fmap f s) (fmap (fmap (fmap f)) e)
    fmap f s = SMap f s

instance Applicative (Signal t) where
    pure = SPure
    (SPure f) <*> (SPure a) = pure (f a)
    (SPure f) <*> s = fmap f s
    a <*> b = SAp a b

dbgSig :: Signal t a -> String
dbgSig (SPure _) = "SPure"
dbgSig (SAp a b) = "SAp (" ++ dbgSig a ++ ") (" ++ dbgSig b ++ ")"
dbgSig (SMap _ s) = "SMap (" ++ dbgSig s ++ ")"
dbgSig (SInt i _) = "SInt " ++ show i
dbgSig (SFn i _) = "SFn " ++ show i
dbgSig (SSwitch s e) = "SSwitch (" ++ dbgSig s ++ ")"
dbgSig (SShare s) = "SShare (" ++ dbgSig s ++ ")"
dbgSig (SAdd a b) = "SAdd (" ++ dbgSig a ++ ") (" ++ dbgSig b ++ ")"
dbgSig (SDiv a b) = "SDiv (" ++ dbgSig a ++ ") (" ++ dbgSig b ++ ")"
dbgSig (SMul a b) = "SMul (" ++ dbgSig a ++ ") (" ++ dbgSig b ++ ")"

instance Num a => Num (Signal t a) where
    fromInteger = pure . fromInteger
    (SPure a) + (SPure b) = pure $ a + b
    a + b = SAdd a b
    (SPure a) * (SPure b) = pure $ a * b
    a * b = SMul a b
    abs = fmap abs
    signum = fmap signum
    negate = fmap negate

instance Fractional a => Fractional (Signal t a) where
    (SPure a) / (SPure b) = pure $ a / b
    a / b = SDiv a b
    recip = fmap recip
    fromRational = pure . fromRational

instance Floating a => Floating (Signal t a) where
    pi = pure pi
    exp = fmap exp
    log = fmap log
    sqrt = fmap sqrt
    (**) = liftA2 (**)
    logBase = liftA2 (**)
    sin = fmap sin
    cos = fmap cos
    tan = fmap tan
    asin = fmap asin
    acos = fmap acos
    atan = fmap atan
    sinh = fmap sinh
    cosh = fmap cosh
    tanh = fmap tanh
    asinh = fmap asinh
    acosh = fmap acosh
    atanh = fmap atanh

-------------------------------------------------------------------
-- Event
-------------------------------------------------------------------

-- | An `Event t a` represents a series of instantaneous events.
data Event t a where
    ERoot  :: (Num a, Ord a) => Signal t a -> Event t ()
    ETag   :: Event t a -> Signal t b -> Event t (a,b)
    EMap   :: (a -> b) -> Event t a -> Event t b

evalEvent :: Time t => Network t o -> Event t a -> a
evalEvent n (ERoot s) = ()
evalEvent n (EMap f e) = f $ evalEvent n e
evalEvent n (ETag e s) = (evalEvent n e, evalSignal n s)

instance Functor (Event t) where
    fmap = EMap

-------------------------------------------------------------------
-- Sim
-------------------------------------------------------------------

-- | The `Sim t` monad is used for constructing certain `Signal` values.
newtype Sim t a = Sim {unSim :: forall o. State (Network t o) a}

instance Num a => Num (Sim t a) where
    fromInteger = pure . fromInteger
    (+) = liftA2 (+)
    (*) = liftA2 (*)
    abs = fmap abs
    signum = fmap signum
    negate = fmap negate

instance Functor (Sim t) where
    fmap f (Sim s) = Sim (fmap f s)

instance Applicative (Sim t) where
    pure a = Sim (pure a)
    (Sim f) <*> (Sim a) = Sim (f <*> a)

instance Monad (Sim t) where
    (Sim a) >>= f = Sim $ do
        r <- a
        let (Sim k) = f r
        k

instance MonadFix (Sim t) where
    mfix f = Sim $ mfix (unSim . f)

-------------------------------------------------------------------
-- Network
-------------------------------------------------------------------

-- | This type family is used for determining the best way to store
-- a vector of values. It should evaluate to something that is an
-- instance of @Data.Vector.Generic.Vector@.
type family NetStoreVec a :: * -> *
type instance NetStoreVec Double = UV.Vector
type instance NetStoreVec Float = UV.Vector
type instance NetStoreVec Word = UV.Vector
type instance NetStoreVec Int = UV.Vector
type instance NetStoreVec (Fixed a) = V.Vector
type instance NetStoreVec Scientific = V.Vector

type NetStore t = NetStoreVec t t
type NetStoreMut t = G.Mutable (NetStoreVec t) t

type Time t = (NFData (NetStoreVec t t), Show t, G.Vector (NetStoreVec t) t, Show (NetStore t))

-- | A `Network` represents a compiled simulation.
data Network t o = Network
    { netIntState :: !(NetStore t)
    , netIntDeriv :: !(IntMap (Signal t (NetStore t)))
    , netFnTime   :: !(NetStore t)
    , netRoot     :: (Signal t o)  -- TODO: Make this a container of signals?
    } deriving (Generic)

-- TODO: Rewrite

-- | Cleans up a `Network` after a switching event, removing signals
-- that are no longer needed.
gc :: forall v t o. Time t => Network t o -> Network t o
gc n@Network{..} =
    let
        (intAll, intReachable, fnReachable) = findNames netRoot

        findNames :: Signal t a -> (IntSet, IntSet, IntSet)
        findNames s = execState (go s)  mempty

        go :: forall a. Signal t a -> State (IntSet, IntSet, IntSet) ()
        go (SInt ix (f :: x -> a)) = do
            (i,i',f) <- get
            when (not $ Set.member ix i) $ do
                let
                    pt = Proxy :: Proxy t
                    pa = Proxy :: Proxy x
                    len = splen pt pa - 1
                    names = Set.fromList [ix..ix + len]
                put (i `Set.union` names, Set.insert ix i', f)
                go (netIntDeriv IntMap.! ix)
        go (SAp a b) = go a >> go b
        go (SMap _ s) = go s
        go (SAdd a b) = go a >> go b
        go (SMul a b) = go a >> go b
        go (SDiv a b) = go a >> go b
        go (SPure _) = return ()
        go (SFn ix _) = modify $ mappend (mempty, mempty, Set.singleton ix)
        go (SSwitch s e) = go s >> goE e
        go (SShare s) = go s

        goE :: Event t a -> State (IntSet, IntSet, IntSet) ()
        goE (ERoot s) = go s
        goE (EMap _ e) = goE e
        goE (ETag e s) = goE e >> go s

        intMap  = UV.fromList (Set.toAscList intReachable)
        intMapAll = UV.fromList (Set.toAscList intAll)
        intMap' = IntMap.fromList $ zip (Set.toAscList intReachable) [0..]

        fnMap   = UV.fromList (Set.toAscList fnReachable)
        fnMap' = IntMap.fromList $ zip (Set.toAscList fnReachable) [0..]

        rename :: Signal t a -> Signal t a
        rename (SAp a b) = SAp (rename a) (rename b)
        rename (SMap f a) = SMap f (rename a)
        rename s@(SPure _) = s
        rename (SInt ix f) = SInt (intMap' IntMap.! ix) f
        rename (SFn ix f) = SFn (fnMap' IntMap.! ix) f
        rename (SSwitch s e) = SSwitch (rename s) (renameE e)
        rename (SShare s) = SShare (rename s)

        renameE :: Event t a -> Event t a
        renameE (ERoot s) = ERoot (rename s)
        renameE (EMap f e) = EMap f (renameE e)
        renameE (ETag e s) = ETag (renameE e) (rename s)

        intCnt = UV.length intMap
        fnCnt  = UV.length fnMap

        intState' = G.generate intCnt $ \ix -> netIntState G.! (intMapAll UV.! ix)

        restrictKeys m s = IntMap.filterWithKey (\k _ -> k `Set.member` s) m
        intDeriv' = fmap rename $ restrictKeys netIntDeriv intReachable

        fnTime' = G.generate fnCnt $ \ix -> netFnTime G.! (fnMap UV.! ix)
        root' = rename netRoot

        -- intDeriv' = execWriter $ getDeriv root'

    in n
        { netIntState = intState'
        , netIntDeriv = IntMap.mapKeys (\k -> intMap' IntMap.! k) intDeriv'
        , netFnTime = fnTime'
        , netRoot = root'
        }

eventOccured :: Time t => Network t o -> Network t o -> Event t a -> Maybe a
eventOccured old new (EMap f e) = f <$> eventOccured old new e
eventOccured old new (ETag e s) = flip (,) (evalSignal new s) <$> eventOccured old new e
eventOccured old new (ERoot s)
    | oldS < 0 && 0 <= newS = Just ()
    | newS <= 0 && 0 < oldS = Just ()
    | otherwise = Nothing
        where
            oldS = evalSignal old s
            newS = evalSignal new s

runSwitches' :: forall v t o. Time t => Network t o -> Network t o -> (Bool, Network t o)
runSwitches' old new =
    let
        go :: Signal t a -> WriterT Any (Sim t) (Signal t a)
        go s@(SPure a) = return s
        go (SMap f s) = go s >>= pure . SMap f
        go (SAdd a b) = do
            a' <- go a
            b' <- go b
            return $ SAdd a' b'
        go (SMul a b) = do
            a' <- go a
            b' <- go b
            return $ SMul a' b'
        go (SDiv a b) = do
            a' <- go a
            b' <- go b
            return $ SDiv a' b'
        go (SAp f a) = do
            f' <- go f
            a' <- go a
            return $ SAp f' a'
        go s@(SInt _ _) = return s
        go s@(SFn ix f) = return s
        go (SSwitch s e) = case eventOccured old new e of
            Just v -> tell (Any True) >> lift v
            Nothing -> do
                s' <- go s
                e' <- goE e
                return $ SSwitch s' e'
        go (SShare s) = go s >>= pure . SShare

        goE :: Event t a -> WriterT Any (Sim t) (Event t a)
        goE (ERoot s) = go s >>= pure . ERoot
        goE (EMap f e) = goE e >>= pure . EMap f
        goE (ETag e s) = do
            e' <- goE e
            s' <- go s
            return (ETag e' s')

        (((root', derivs'), Any changed), net') = addSim new $ runWriterT $ do
            newRoot <- go $ netRoot new
            newDerivs <- mapM go $ netIntDeriv new
            return (newRoot, newDerivs)
    in
        if changed
        then (changed, gc $ net'
            { netRoot = root'
            , netIntDeriv = derivs'
            })
        else (changed, new)

runSwitches old new = snd $ runSwitches' old new
anyEvent old new = fst $ runSwitches' old new

newDoubleSim :: Sim Double (Signal Double a) -> Network Double a
newDoubleSim = newSim

newSim :: Time t => Sim t (Signal t a) -> Network t a
newSim (Sim s) =
    let
        (r,n) = runState s Network
            { netIntState = G.empty
            , netIntDeriv = IntMap.empty
            , netFnTime  = G.empty
            , netRoot = r
            }
    in n

addSim :: Time t => Network t o -> Sim t a -> (a, Network t o)
addSim n (Sim s) = runState s n

-- TODO: Below are three implementations of evalSignal. Each with different
-- sharing characteristics.
--  * evalSignal' does explicit sharing
--  * evalSignal'' does no sharing
--  * evalSignal''' does implicit sharing
-- We need to benchmark and determine which is best.
-- Once we determine which solution to go for, we should integrate it
-- more with the rest of the code such that sharing caries over to
-- events.
--
-- It's probably worth redefining `share` to `id` when not using sharing
-- or using implicit sharing.
evalSignal'
    :: forall t o a. (Time t)
    => Network t o -> Signal t a -> State (IntMap (StableName Prim.Any, Any)) a
evalSignal' c me@(SShare s) = do
    m <- get
    let
        name = unsafePerformIO $ makeStableName me
        nameh = hashStableName name
        mv = do
            (n,v) <- IntMap.lookup nameh m
            if eqStableName name n
            then return $ unsafeCoerce v
            else Nothing
    case mv of
        Just v -> return v
        Nothing -> do
            v' <- evalSignal' c s
            put $ IntMap.insert nameh (unsafeCoerce name, unsafeCoerce v') m
            return v'
evalSignal' c (SPure a) = return a
evalSignal' c (SAdd a b) = do
    a' <- evalSignal' c a
    b' <- evalSignal' c b
    return $ a' + b'
evalSignal' c (SMul a b) = do
    a' <- evalSignal' c a
    b' <- evalSignal' c b
    return $ a' * b'
evalSignal' c (SDiv a b) = do
    a' <- evalSignal' c a
    b' <- evalSignal' c b
    return $ a' / b'
evalSignal' c (SMap f s) = f <$> evalSignal' c s
evalSignal' c (SAp f a) = do
    f' <- evalSignal' c f
    a' <- evalSignal' c a
    return $ f' a'
evalSignal' c (SFn ix f) = return $ f $ netFnTime c G.! ix
evalSignal' c (SInt ix (f :: x -> a)) = do
    let
        len = splen (Proxy :: Proxy t) (Proxy :: Proxy x)
    return $ f $ unsplat $ G.slice ix len (netIntState c)
evalSignal' c (SSwitch s _) = evalSignal' c s

evalSignal'' :: forall t o a. Time t => Network t o -> Signal t a -> a
evalSignal'' c = go
    where
        go :: forall b. Signal t b -> b
        go (SPure a) = a `seq` a
        go (SMap f s) = let x = f $! go s in x `seq` x
        go (SAdd a b) = let x = go a + go b in x `seq` x
        go (SMul a b) = let x = go a * go b in x `seq` x
        go (SDiv a b) = let x = go a / go b in x `seq` x
        go (SAp f a) =
            let
                f' = go f
                a' = go a
                x = f' a'
            in x `seq` x
        go (SFn ix f) = let x = f $! netFnTime c G.! ix in x `seq` x
        go (SInt ix (f :: x -> b)) = let x = f $! unsplat $ G.slice ix (splen (Proxy :: Proxy t) (Proxy :: Proxy x)) (netIntState c) in x `seq` x
        go (SSwitch s _) = go s
        go (SShare s) = go s

evalSignal''' :: forall t o a. Time t => Network t o -> Signal t a -> a
evalSignal''' c = go
    where
        go :: Signal t c -> c
        go = memo gogo

        gogo :: forall b. Signal t b -> b
        gogo (SPure a) = a
        gogo (SMap f s) = f $ go s
        gogo (SAdd a b) = go a + go b
        gogo (SMul a b) = go a * go b
        gogo (SDiv a b) = go a / go b
        gogo (SAp f a) = go f $ go a
        gogo (SFn ix f) = f $ netFnTime c G.! ix
        gogo (SInt ix (f :: x -> b)) = f $ unsplat $ G.slice ix (splen (Proxy :: Proxy t) (Proxy :: Proxy x)) (netIntState c)
        gogo (SSwitch s _) = go s
        gogo (SShare s) = go s

evalSignal :: forall t o a. Time t => Network t o -> Signal t a -> a
--evalSignal n s = evalState (evalSignal' n s) mempty
evalSignal = evalSignal''
--evalSignal = evalSignal'''
--evalSignal n = memo (evalSignal'' n)
--{-# NOINLINE evalSignal #-}

evalRoot :: Time t => Network t o -> o
evalRoot c = evalSignal c (netRoot c)

deltaNet :: (Num t, Time t) => Network t o -> t -> NetStore t -> Network t o
deltaNet n@Network{..} h dv = deltaNetState (deltaNetTime n h) dv

deltaNetState :: (Num t, Time t) => Network t o -> NetStore t -> Network t o
deltaNetState n@Network{..} dv = n
    { netIntState = G.zipWith (+) netIntState dv
    }

deltaNetTime :: (Num t, Time t) => Network t o -> t -> Network t o
deltaNetTime n@Network{..} h = n
    { netFnTime  = G.map (+h) netFnTime
    }

derivs :: (Num t, Time t) => Network t o -> t -> NetStore t -> NetStore t
derivs n h dv =
    let
        n' = deltaNet n h dv
    in derivsNow n'

derivsNow :: (Time t) => Network t o -> NetStore t
derivsNow n = G.concat $ parMapChunk (evalSignal n . snd) $ IntMap.toAscList $ netIntDeriv n

-- TODO: Why is this burning so much CPU and not giving me /that/ much of
-- a speed up. Changing the chunk size doesn't seem to make that much
-- difference.
--
-- 1 -N1  (100% cpu)
-- real	0m14.623s
-- user	0m14.458s
-- sys	0m0.149s
--
-- All the rest 580% cpu
-- 1 -N
-- real	0m10.852s
-- user	0m39.608s
-- sys	0m22.108s
--
-- 20 -N
-- real	0m10.932s
-- user	0m35.629s
-- sys	0m23.757s
--
-- 200 -N
-- real	0m11.231s
-- user	0m37.745s
-- sys	0m24.182s
parMapChunk f as =
    let
        a = map f as
    in a `Par.using` Par.parListChunk 50 (Par.rdeepseq)

-------------------------------------------------------------------
-- Splat
-------------------------------------------------------------------

class Splat t a where
    splen :: Proxy t -> Proxy a -> Int
    splat :: Time t => a -> NetStore t
    unsplat :: Time t => NetStore t -> a

instance Splat t t where
    splen _ _ = 1
    splat t = G.singleton t
    unsplat v = v G.! 0

instance Splat t (t,t) where
    splen t _ = 2
    splat (a,b) = splat a G.++ splat b
    unsplat v = (v G.! 0, v G.! 1)

instance Splat t (t,t,t) where
    splen t _ = 3
    splat (a,b,c) = G.fromList [a,b,c]
    unsplat v = (v G.! 0, v G.! 1, v G.! 2)

instance Splat t (t,t,t,t) where
    splen t _ = 4
    splat (a,b,c,d) = G.fromList [a,b,c,d]
    unsplat v = (v G.! 0, v G.! 1, v G.! 2, v G.! 3)

instance Splat t (Complex t) where
    splen t _ = 2
    splat (a :+ b) = G.fromList [a,b]
    unsplat v = (v G.! 0) :+ (v G.! 1)

instance Splat t (V1 t) where
    splen t _ = splen t (Proxy :: Proxy t)
    splat (V1 a) = splat a
    unsplat = V1 . unsplat

instance Splat t (V2 t) where
    splen t _ = splen t (Proxy :: Proxy (t,t))
    splat (V2 a b) = splat (a,b)
    unsplat v = let (a,b) = unsplat v in V2 a b

instance Splat t (V3 t) where
    splen t _ = splen t (Proxy :: Proxy (t,t,t))
    splat (V3 a b c) = splat (a,b,c)
    unsplat v = let (a,b,c) = unsplat v in V3 a b c

instance Splat t (V4 t) where
    splen t _ = splen t (Proxy :: Proxy (t,t,t,t))
    splat (V4 a b c d) = splat (a,b,c,d)
    unsplat v = let (a,b,c,d) = unsplat v in V4 a b c d

-------------------------------------------------------------------
-- API
-------------------------------------------------------------------

-- | Integrate the input `Signal` with respect to time.
integral :: (Splat t a, Time t) => a -> Signal t a -> Sim t (Signal t a)
integral i s = Sim $ do
    st <- get
    let
        sp = splat i
        ix = G.length (netIntState st)
    put st
        { netIntState = netIntState st G.++ sp
        , netIntDeriv = IntMap.insert ix (fmap splat s) (netIntDeriv st)
        }
    return $ SInt ix id

-- | Create a `Signal` that is a pure function of time. Each `timeFn'` has
-- a local concept of time, and this time starts from the provided time
-- value.
timeFn' :: Time t => t -> (t -> a) -> Sim t (Signal t a)
timeFn' t f = Sim $ do
    st <- get
    let
        ix = G.length (netFnTime st)
    put st
        { netFnTime = netFnTime st `G.snoc` t
        }
    return $ SFn ix f

-- | As `timeFn'` but local time starts at 0.
timeFn :: (Time t, Num t) => (t -> a) -> Sim t (Signal t a)
timeFn = timeFn' 0

-- | Start out as the input `Signal`, but when the `Event` occurs, become
-- the signal that it carries, and remain as that signal forever. Future
-- `Event`s from this event source will be ignored. See `switch` for
-- alternate behavior.
become :: Signal t a -> Event t (Sim t (Signal t a)) -> Signal t a
become = SSwitch

-- | Start out as the input `Signal` and everytime the `Event` occurs, change
-- to the new `Signal` carried in the `Event`. This can easily cause
-- a space leak if used recursively. Look at `become` for a safer solution.
switch :: Signal t a -> Event t (Sim t (Signal t a)) -> Signal t a
switch s e = become s (fmap (\s' -> s' >>= \s'' -> pure (switch s'' e)) e)

-- | Useful combination of `become` and `fmap`
becomeOn :: Signal t a -> Event t b -> (b -> Sim t (Signal t a)) -> Signal t a
becomeOn s e f = become s (fmap f e)

-- | Turns the input `Signal t a` into a `Signal t (Maybe a)`, resulting
-- in `Nothing` as soon as the `Event` fires, and remaining `Nothing` forever
-- after that point.
becomeNothingOn :: Signal t a -> Event t b -> Signal t (Maybe a)
becomeNothingOn s e = becomeOn (fmap Just s) e $ \_ -> return $ pure Nothing

-- | Memoize the result of evaluating this Signal so that repeated
-- uses don't have to re-evaluate everything. Integrations are
-- already shared by default, so avoid sharing them directly.
--
-- TODO: Depending on the signal evaluation strategy (see comments in code)
-- sharing may implicitly happen for all signals, or not at all.
share :: Signal t a -> Signal t a
share = SShare

-- | Emit an event every time the input `Signal` crosses 0.
root :: (Ord t, Num t) => Signal t t -> Event t ()
root = ERoot

-- | Tag the given event with the value of the signal when it occurs.
tag :: Event t a -> Signal t b -> Event t (a,b)
tag  = ETag

-------------------------------------------------------------------
-- Simulation
-------------------------------------------------------------------

data SimResult t o = SimResult
    { stepNumber :: !Integer
    , globalTime :: !t
    , result     :: !o
    } deriving (Eq, Ord, Show, Functor)

asTuple :: SimResult t o -> (t,o)
asTuple (SimResult _ t o) = (t,o)

-- | Run the `Integrator` over the given `Signal` outputting a list of
-- time stamp and `Network`s.
simulate'
    :: forall v t o. (Num t, Time t)
    => Integrator t o -> Sim t (Signal t o) -> [SimResult t (Network t o)]
simulate' integrator s =
    let
        n = newSim s
    in unfoldr (\(cnt,i,n,t) ->
        let
            (n',r,i') =  runIntegrator i n
            t' = t + r
        in
            cnt `seq` t `seq` n  `seq` (Just $!
                ( SimResult cnt t n   -- Previous step
                , (succ cnt, i', n', t')
                ))
        ) (0,integrator,n,0)

-- | Run the `Integrator` over the given `Signal`, outputting a list of
-- time stamp and signal values.
simulate :: (Num t, Time t) => Integrator t o -> Sim t (Signal t o) -> [SimResult t o]
simulate i s = fmap evalRoot <$> simulate' i s

-- | Causes progress messages to be printed every 100 steps. Uses lazy IO,
-- so be careful.
--  TODO: Rejig things to be more streamy.
trace :: (Show t, Show o, PrintfArg t) => [SimResult t o] -> IO [SimResult t o]
trace [] = return []
trace (s@(SimResult c t o):rest)
    | 0 == c `mod` 100 = do
        printf "S: %5d  T: %5.3v  V: " c t
        putStrLn $ show o
        r <- unsafeInterleaveIO $ trace rest
        return $ s : r
    | otherwise = do
        r <- unsafeInterleaveIO $ trace rest
        return $ s : r

-- | `simulate'` specialised to `Double`
simulateDouble'
    :: Integrator Double o
    -> Sim Double (Signal Double o)
    -> [SimResult Double (Network Double o)]
simulateDouble' = simulate'

-- | `simulate` specialised to `Double`
simulateDouble
    :: Integrator Double o
    -> Sim Double (Signal Double o)
    -> [SimResult Double o]
simulateDouble = simulate

-- | `simulate` which stops as soon as the `Signal` value becomes `Nothing`
simulateJust :: (Num t, Time t) => Integrator t (Maybe o) -> Sim t (Signal t (Maybe o)) -> [SimResult t o]
simulateJust i s = fmap (\(SimResult n t (Just x)) -> SimResult n t x) $ takeWhile (isJust . result) $ simulate i s

-- | `simulateJust` specialised to `Double`
simulateJustDouble :: Integrator Double (Maybe o) -> Sim Double (Signal Double (Maybe o)) -> [SimResult Double o]
simulateJustDouble = simulateJust

-- | Simulate until the specified time is reached. The termination time is
-- marked by an `Event`, meaning that adaptive integration methods which
-- narrow down on event locations will be do so to get accurate finishing
-- times.
simulateUntil :: (Ord t, Num t, Time t) => t -> Integrator t (Maybe o) -> Sim t (Signal t o) -> [SimResult t o]
simulateUntil duration i s = simulateJust i $ do
    t <- timer duration
    s' <- s
    return $ s' `becomeNothingOn` t

-- | `simulateUntil` specialised to `Double`
simulateUntilDouble :: Double -> Integrator Double (Maybe o) -> Sim Double (Signal Double o) -> [SimResult Double o]
simulateUntilDouble = simulateUntil

-- | Run a simulation in real time, printing the result values to the terminal.
runRk4RealTime :: forall a. Show a => Sim Double (Signal Double a) -> IO ()
runRk4RealTime s =
    let
        n = newSim s
        go :: UTCTime -> Network Double a -> IO ()
        go prev n = do
            now <- getCurrentTime
            let
                h = realToFrac $ diffUTCTime now prev
                (n',_,_) = runIntegrator (rk4 h) n
            print $ evalRoot n'
            go now n'
    in do
        now <- getCurrentTime
        go now n

-- | Same as `runRk4RealTime` but stops when a `Nothing` value is produced.
runRk4RealTimeJust :: forall a. Show a => Sim Double (Signal Double (Maybe a)) -> IO ()
runRk4RealTimeJust s =
    let
        n = newSim s
        go :: UTCTime -> Network Double (Maybe a) -> IO ()
        go prev n = do
            now <- getCurrentTime
            let
                h = realToFrac $ diffUTCTime now prev
                (n',_,_) = runIntegrator (rk4 h) n
            case evalRoot n' of
                Nothing -> return ()
                Just x -> print x >> go now n'
    in do
        now <- getCurrentTime
        go now n

sample :: (Ord t, Real t) => t -> [(t,a)] -> [(t,a)]
sample freq = go 500
    where
        go _ [] = []
        go n ((t,r):xs)
            | t' < n = (t,r) : go t' xs
            | otherwise = go t' xs
            where
                t' = t `mod'` freq

-------------------------------------------------------------------
-- Integrators
-------------------------------------------------------------------

newtype Integrator t o = Integrator
    { runIntegrator :: Network t o -> (Network t o, t, Integrator t o)
    } deriving Generic

instance NFData (Integrator t o)

eulerSimple :: (Num t, Time t) => t -> Integrator t o
eulerSimple h = integrator
    where
        integrator = Integrator go
        go n =
            let
                ds = derivsNow n
                n' = deltaNet n h (G.map (*h) ds)
            in (runSwitches n n', h, integrator)

-- | Same as `euler` except that it bisects around `Event`s until it is within
-- the provided tolerance.
eulerBisect :: (Fractional t, Ord t, Time t) => t -> t -> Integrator t o
eulerBisect tol h = integrator
    where
        integrator = Integrator $ go h
        go h' n =
            let
                ds = derivsNow n
                n' = deltaNet n h' (G.map (*h') ds)
                (e,n'') = runSwitches' n n'
                res
                    -- TODO: Don't emit nearby points, jump strait to it.
                    | e && h' > tol = go (h'/2) n
                    | otherwise = (n'', h', integrator)
            in res

-- | Configuration for adaptive methods.
--
-- Unchecked assumptions: 0 < `minStep` < `eventTolerance` < `maxStep`
data AdaptiveConfig t = AdaptiveConfig
    { minStep :: Maybe t
    , maxStep :: Maybe t
    , tolerance :: t
    , eventTolerance :: t
    }

instance Fractional t => Default (AdaptiveConfig t) where
    def = AdaptiveConfig
        { minStep = Nothing
        , maxStep = Nothing
        , tolerance = 0.001
        , eventTolerance = 0.001
        }

-- If maxStep < minStep, we choose maxStep always.
adaptiveClamp :: Ord t => AdaptiveConfig t -> t -> t
adaptiveClamp AdaptiveConfig{..} t =
    let
        t' = maybe t (max t) minStep
    in
        maybe t' (min t') maxStep

clamp :: Ord a => a -> a -> a -> a
clamp l h t
    | t < l = l
    | t > h = h
    | otherwise = t

euler :: (Ord t, Fractional t, Time t, Show t) => AdaptiveConfig t -> Integrator t o
euler conf@AdaptiveConfig{..} = Integrator $ go 0.1
    where
        go h n =
            let
                eulStep h n = deltaNet n h (G.map (*h) (derivsNow n))

                n10  = eulStep h n
                n1_2 = eulStep (h/2) n
                n11  = eulStep (h/2) n1_2
                tau  = G.zipWith (-) (netIntState n11) (netIntState n10)
                tau' = G.maximum $ G.map abs tau
                n12  = deltaNetState n11 tau

                h' = adaptiveClamp conf $ 0.9 * h * clamp 0.3 3 (tolerance / tau')

                (e,n') = runSwitches' n n12

                res
                    | e && h > eventTolerance && h > fromMaybe 0 minStep
                        = go (adaptiveClamp conf (h/2)) n
                    | tau' > tolerance && h' > fromMaybe 0 minStep = go h' n
                    | otherwise = (n', h, Integrator $ go h')
            in res

rk4 :: (Fractional t, Time t) => t -> Integrator t o
rk4 h = integrator
    where
        integrator = Integrator go
        go n =
            let
                h2 = h / 2
                h6 = h / 6
                k1 = derivsNow n
                k2 = derivs n h2 (G.map (*h2) k1)
                k3 = derivs n h2 (G.map (*h2) k2)
                k4 = derivs n h  (G.map (*h)  k3)
                n' = deltaNet n h (G.map (* h6)
                    $ G.zipWith4 (\a b c d -> a + 2 * b + 2 * c + d) k1 k2 k3 k4)
            in
                (runSwitches n n', h, integrator)

-------------------------------------------------------------------
-- Visualisation
-------------------------------------------------------------------

class Plot a where
    plot :: [SimResult Double a] -> Chart.EC (Chart.Layout Double Double) ()

newtype a ::: (name :: Symbol) = Named a
    deriving (Read,Enum,Real,RealFrac,RealFloat,Eq,Ord,Show,Num,Fractional,Floating,Chart.PlotValue)

infixr 0 :::

instance Plot Double where
    plot = Chart.plot . Chart.points "" . fmap asTuple

instance (KnownSymbol name, Chart.PlotValue a) => Plot (a ::: name) where
    plot = Chart.plot . Chart.points name . fmap (fmap Chart.toValue) . fmap asTuple
        where
            name = symbolVal (Proxy :: Proxy name)

instance (Plot a, Plot b) => Plot (a,b) where
    plot as = do
        plot (fmap fst <$> as)
        plot (fmap snd <$> as)

instance (Plot a, Plot b, Plot c) => Plot (a,b,c) where
    plot as = do
        plot (fmap (\(a,_,_) -> a) <$> as)
        plot (fmap (\(_,b,_) -> b) <$> as)
        plot (fmap (\(_,_,c) -> c) <$> as)

instance Plot a => Plot (V2 a) where
    plot as = do
        plot (fmap (Chart.^. _x) <$> as)
        plot (fmap (Chart.^. _y) <$> as)

instance Plot a => Plot (V4 a) where
    plot as = do
        plot (fmap (Chart.^. _x) <$> as)
        plot (fmap (Chart.^. _y) <$> as)
        plot (fmap (Chart.^. _z) <$> as)
        plot (fmap (Chart.^. _w) <$> as)

plotSimUntil
    :: Plot o
    => Double
    -> Integrator Double (Maybe o)
    -> Sim Double (Signal Double o)
    -> Chart.EC (Chart.Layout Double Double) ()
plotSimUntil t i s = plot (simulateUntil t i s)

plotSimUntilToFile n title t i s = do
    r <- trace $ simulateUntil t i s
    Chart.toFile Chart.def n $ do
        Chart.layout_title Chart..= title
        plot r

class PlotPara a where
    plotPara :: [SimResult Double a] -> Chart.EC (Chart.Layout Double Double) ()

instance PlotPara (V2 Double) where
    plotPara = Chart.plot . Chart.line "" . (:[]) . fmap (\(SimResult n t (V2 a b)) -> (a,b))

instance (KnownSymbol name) => PlotPara (V2 Double ::: name) where
    plotPara = Chart.plot . Chart.line name . (:[]) . fmap (\(SimResult n t (Named (V2 a b))) -> (a,b))
        where
            name = symbolVal (Proxy :: Proxy name)

instance (PlotPara a, PlotPara b) => PlotPara (a,b) where
    plotPara as = do
        plotPara (fmap fst <$> as)
        plotPara (fmap snd <$> as)

instance (PlotPara a, PlotPara b, PlotPara c) => PlotPara (a,b,c) where
    plotPara as = do
        plotPara (fmap (\(a,_,_) -> a) <$> as)
        plotPara (fmap (\(_,b,_) -> b) <$> as)
        plotPara (fmap (\(_,_,c) -> c) <$> as)

instance PlotPara (Map String (V2 Double)) where
    plotPara [] = return ()
    plotPara as@(SimResult _ _ a : _) = do
        let
            ks = Map.keys a
            kk = if length ks > 10 then \_ -> "" else id
        forM_ a $ \(V2 x y) -> do
            Chart.plot $ Chart.points "" [(x,y)]
        forM_ ks $ \k -> do
            Chart.plot $ Chart.line (kk k) . (:[]) $ fmap (\(SimResult n t m) ->
                let
                    V2 a b = m Map.! k
                in (a,b)
                ) as

plotParaSimUntil
    :: PlotPara o
    => Double
    -> Integrator Double (Maybe o)
    -> Sim Double (Signal Double o)
    -> Chart.EC (Chart.Layout Double Double) ()
plotParaSimUntil t i s = plotPara $ simulateUntil t i s

plotParaSimUntilToFile n title t i s = do
    r <- trace $ simulateUntil t i s
    Chart.toFile Chart.def n $ do
        Chart.layout_title Chart..= title
        plotPara r

-------------------------------------------------------------------
-- Utility
-------------------------------------------------------------------

-- | Applicative tupples
(-:) :: Applicative f => f a -> f b -> f (a,b)
(-:) = liftA2 (,)
infixr 5 -:

-- | Applicative `Named` tuples.
(-::) :: Applicative f => f a -> f b -> f (a ::: an, b ::: bn)
a -:: b = (,) <$> (Named <$> a) <*> (Named <$> b)
infixr 5 -::

threeNames :: Applicative f => f a -> f b -> f c -> f (a ::: an, b ::: nm, c ::: cn)
threeNames a b c = (,,) <$> (Named <$> a) <*> (Named <$> b) <*> (Named <$> c)

-------------------------------------------------------------------
-- Examples/Test
-------------------------------------------------------------------

-- Ex1 demonstrates basic integration with recursion.
ex1 :: (Time t, Fractional t) => Sim t (Signal t (t, t))
ex1 = do
    rec a <- integral 0.5 a
    b <- integral 0 1
    return $ a -: b

-- Ex2 shows how to use `become` to change a singal to a new one upon
-- an event.
ex2 :: forall t. (Fractional t, Ord t, Time t) => Sim t (Signal t t)
ex2 = do
    t <- integral 5 (-1)
    let
        e = root t
        e' = fmap (const $ pure 3) e
    return $ become 1 e'

-- Ex3 show the use of `switch` to control a signal.
ex3 :: (Time t, Floating t, Ord t) => Sim t (Signal t (t,t))
ex3 = do
    t <- timeFn sin
    let
        e = tag (root t) (signum t)
        e' = fmap (\(_,v) -> return $ pure v) e
    return $ (,) <$> t <*> switch 0 e'

-- Ex4 can be used to demonstrate the accuracy of various integrators.
ex4
    :: (Floating t, Time t)
    => Sim t (Signal t
        ( t ::: "sin(t)"
        , t ::: "∫∫∫∫sin(t)"
        ))
ex4 = do
    s <- timeFn sin
    s' <- integral (-1) s >>= integral 0 >>= integral 1 >>= integral 0
    return $ s -:: s'

-- Ex5 demonstrates using `Maybe` to terminate a simulation with
-- `simulateJust`.
--
-- Note: The use of an `Event` rather than `fmap`ing over the signal
-- causes the root finding algorithm of many integrators to kick in allowing
-- for very precise finishing times.
ex5 :: (Floating t, Time t, Ord t) => Sim t (Signal t (Maybe (V4 t)))
ex5 = do
    i <- (integral 0 $ pure (V4 0 1 2 (-0.5)) )
    let
        s = fmap (\i -> 5 - norm i) i
        e = root s
    return $ i `becomeNothingOn` e

-- Ex6 demonstrates the use of `share` to reduce the number
-- of times a signal is evaluated. The trace will show
-- "Calc A" once, but "Calc B" twice when using an explicit
-- sharing eval strategy.
ex6 :: (Time t, Floating t) => Sim t (Signal t (V4 t))
ex6 = do
    let
        a = share $ fmap (Dbg.trace "Calc A") $ 5 + 2
        b = fmap (Dbg.trace "Calc B") $ 5 + 2
    return $ V4 <$> a <*> a <*> b <*> b

ex7 :: (Ord t, Floating t, Time t) => Sim t (Signal t t)
ex7 = do
    s <- timeFn sin
    let
        e = root s
        s1 = becomeOn 0 e (\_ -> pure $ signum s)
    integral 0 s1

timer :: (Num t, Time t, Ord t) => t -> Sim t (Event t ())
timer duration = timeFn (\t -> duration - t) >>= pure . root

stepAndHold :: (Num t, Ord t, Time t) => [(t,Sim t (Signal t a))] -> Sim t (Signal t a) -> Sim t (Signal t a)
stepAndHold [] a = a
stepAndHold ((t,a):ls) last = do
    e <- timer t
    s <- a
    return $ s `becomeOn` e $ \_ -> stepAndHold ls last

ex8 :: (Ord t, Floating t, Time t) => Sim t (Signal t (Maybe (t ::: "Deriv",t ::: "Integral")))
ex8 = do
    t <- timer 6
    d <- stepAndHold [(0.5,1),(0.2, -2),(1,0)] (timeFn sin)
    i <- integral 0 d
    let r = d -:: i
    return $ r `becomeNothingOn` t

ex9
    :: forall t. (Floating t, Time t)
    => Sim t (Signal t
        ( V2 t ::: "Sun"
        , V2 t ::: "Earth"
        , V2 t ::: "Moon"
        ))
ex9 = do
    rec
        let
            degRad :: t -> t
            degRad t = t / 180 * pi

            pos2init = 149.6e9 *^ angle $ degRad 3.201e2
            pos3init = pos2init + 384.4e6 *^ angle (degRad 2.0644e2)

        pos1 <- integral (V2 0 0) vel1
        pos2 <- integral pos2init vel2
        pos3 <- integral pos3init vel3

        let
            vel2init = 29766.42101876582 *^ signorm (perp pos2init)
            vel3init = vel2init + 1023.005 *^ (signorm (perp pos3init))
        vel1 <- integral (V2 0 0) acc1
        vel2 <- integral vel2init acc2
        vel3 <- integral vel3init acc3

        let
            g = 6.674e-11
            m1 = 1.9891e30
            m2 = 5.972e24
            m3 = 7.347e22
            -- f = g * m1 * m2 / r^2 = m1 * a
            dir :: Signal t (V2 t) -> Signal t (V2 t) -> Signal t t -> Signal t (V2 t)
            dir a b c =
                let
                    d = (^-^) <$> b <*> a
                    u = signorm <$> d
                in (*^) <$> c <*> u
            f1 = sum
                [ dir pos1 pos2 $ g * m1 * m2 / (qd <$> pos1 <*> pos2)
                , dir pos1 pos3 $ g * m1 * m3 / (qd <$> pos1 <*> pos3)
                ]
            f2 = sum
                [ dir pos2 pos1 $ g * m2 * m1 / (qd <$> pos2 <*> pos1)
                , dir pos2 pos3 $ g * m2 * m3 / (qd <$> pos2 <*> pos3)
                ]
            f3 = sum
                [ dir pos3 pos1 $ g * m3 * m1 / (qd <$> pos3 <*> pos1)
                , dir pos3 pos2 $ g * m3 * m2 / (qd <$> pos3 <*> pos2)
                ]
            acc1 = (^/) <$> f1 <*> m1
            acc2 = (^/) <$> f2 <*> m2
            acc3 = (^/) <$> f3 <*> m3
    return $ threeNames pos1 pos2 pos3