{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE TypeApplications #-}

module Ouroboros.Network.PeerSelection.LedgerPeers.Utils
  ( bigLedgerPeerQuota
  , accumulateBigLedgerStake
  , recomputeRelativeStake
  , AccPoolStake (..)
  , PoolStake (..)
  , RelayAccessPoint (..)
  ) where

import Control.Exception (assert)
import Data.Bifunctor (first)
import Data.List as List (foldl', sortOn)
import Data.List.NonEmpty (NonEmpty)
import Data.Ord (Down (..))
import Data.Ratio ((%))

import Ouroboros.Network.PeerSelection.LedgerPeers.Type
import Ouroboros.Network.PeerSelection.RelayAccessPoint

-- | The total accumulated stake of big ledger peers.
--
bigLedgerPeerQuota :: AccPoolStake
bigLedgerPeerQuota :: AccPoolStake
bigLedgerPeerQuota = AccPoolStake
0.9

-- | Sort ascendingly a given list of pools with stake,
-- and tag each one with cumulative stake, with a cutoff
-- at 'bigLedgerPeerQuota'
--
accumulateBigLedgerStake :: [(PoolStake, NonEmpty RelayAccessPoint)]
                         -> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
accumulateBigLedgerStake :: [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
accumulateBigLedgerStake =
    ((AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint)) -> Bool)
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
forall a. (a -> Bool) -> [a] -> [a]
takeWhilePrev (\(AccPoolStake
acc, (PoolStake, NonEmpty RelayAccessPoint)
_) -> AccPoolStake
acc AccPoolStake -> AccPoolStake -> Bool
forall a. Ord a => a -> a -> Bool
<= AccPoolStake
bigLedgerPeerQuota)
    ([(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
 -> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))])
-> ([(PoolStake, NonEmpty RelayAccessPoint)]
    -> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))])
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AccPoolStake
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
go AccPoolStake
0
    ([(PoolStake, NonEmpty RelayAccessPoint)]
 -> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))])
-> ([(PoolStake, NonEmpty RelayAccessPoint)]
    -> [(PoolStake, NonEmpty RelayAccessPoint)])
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((PoolStake, NonEmpty RelayAccessPoint) -> Down PoolStake)
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (PoolStake -> Down PoolStake
forall a. a -> Down a
Down (PoolStake -> Down PoolStake)
-> ((PoolStake, NonEmpty RelayAccessPoint) -> PoolStake)
-> (PoolStake, NonEmpty RelayAccessPoint)
-> Down PoolStake
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PoolStake, NonEmpty RelayAccessPoint) -> PoolStake
forall a b. (a, b) -> a
fst)
    ([(PoolStake, NonEmpty RelayAccessPoint)]
 -> [(PoolStake, NonEmpty RelayAccessPoint)])
-> ([(PoolStake, NonEmpty RelayAccessPoint)]
    -> [(PoolStake, NonEmpty RelayAccessPoint)])
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LedgerPeersKind
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
recomputeRelativeStake LedgerPeersKind
BigLedgerPeers
  where
    takeWhilePrev :: (a -> Bool) -> [a] -> [a]
    takeWhilePrev :: forall a. (a -> Bool) -> [a] -> [a]
takeWhilePrev a -> Bool
f [a]
as =
        ((Maybe a, a) -> a) -> [(Maybe a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe a, a) -> a
forall a b. (a, b) -> b
snd
      ([(Maybe a, a)] -> [a])
-> ([(Maybe a, a)] -> [(Maybe a, a)]) -> [(Maybe a, a)] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Maybe a, a) -> Bool) -> [(Maybe a, a)] -> [(Maybe a, a)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\(Maybe a
a, a
_) -> Bool -> (a -> Bool) -> Maybe a -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True a -> Bool
f Maybe a
a)
      ([(Maybe a, a)] -> [a]) -> [(Maybe a, a)] -> [a]
forall a b. (a -> b) -> a -> b
$ [Maybe a] -> [a] -> [(Maybe a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Maybe a
forall a. Maybe a
Nothing Maybe a -> [Maybe a] -> [Maybe a]
forall a. a -> [a] -> [a]
: (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> [a] -> [Maybe a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a]
as)) [a]
as

    -- natural fold
    go :: AccPoolStake
       -> [(PoolStake, NonEmpty RelayAccessPoint)]
       -> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
    go :: AccPoolStake
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
go AccPoolStake
_acc [] = []
    go !AccPoolStake
acc (a :: (PoolStake, NonEmpty RelayAccessPoint)
a@(PoolStake
s, NonEmpty RelayAccessPoint
_) : [(PoolStake, NonEmpty RelayAccessPoint)]
as) =
      let acc' :: AccPoolStake
acc' = AccPoolStake
acc AccPoolStake -> AccPoolStake -> AccPoolStake
forall a. Num a => a -> a -> a
+ Rational -> AccPoolStake
AccPoolStake (PoolStake -> Rational
unPoolStake PoolStake
s)
      in (AccPoolStake
acc', (PoolStake, NonEmpty RelayAccessPoint)
a) (AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
forall a. a -> [a] -> [a]
: AccPoolStake
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(AccPoolStake, (PoolStake, NonEmpty RelayAccessPoint))]
go AccPoolStake
acc' [(PoolStake, NonEmpty RelayAccessPoint)]
as

-- | Not all stake pools have valid \/ usable relay information. This means that
-- we need to recalculate the relative stake for each pool.
--
recomputeRelativeStake :: LedgerPeersKind
                -> [(PoolStake, NonEmpty RelayAccessPoint)]
                -> [(PoolStake, NonEmpty RelayAccessPoint)]
recomputeRelativeStake :: LedgerPeersKind
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
recomputeRelativeStake LedgerPeersKind
ledgerPeersKind [(PoolStake, NonEmpty RelayAccessPoint)]
pl =
    let pl' :: [(PoolStake, NonEmpty RelayAccessPoint)]
pl'   = (PoolStake -> PoolStake)
-> (PoolStake, NonEmpty RelayAccessPoint)
-> (PoolStake, NonEmpty RelayAccessPoint)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first PoolStake -> PoolStake
adjustment ((PoolStake, NonEmpty RelayAccessPoint)
 -> (PoolStake, NonEmpty RelayAccessPoint))
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(PoolStake, NonEmpty RelayAccessPoint)]
pl
        total :: PoolStake
total = (PoolStake -> PoolStake -> PoolStake)
-> PoolStake -> [PoolStake] -> PoolStake
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' PoolStake -> PoolStake -> PoolStake
forall a. Num a => a -> a -> a
(+) PoolStake
0 ((PoolStake, NonEmpty RelayAccessPoint) -> PoolStake
forall a b. (a, b) -> a
fst ((PoolStake, NonEmpty RelayAccessPoint) -> PoolStake)
-> [(PoolStake, NonEmpty RelayAccessPoint)] -> [PoolStake]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(PoolStake, NonEmpty RelayAccessPoint)]
pl')
        pl'' :: [(PoolStake, NonEmpty RelayAccessPoint)]
pl''  = (PoolStake -> PoolStake)
-> (PoolStake, NonEmpty RelayAccessPoint)
-> (PoolStake, NonEmpty RelayAccessPoint)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (PoolStake -> PoolStake -> PoolStake
forall a. Fractional a => a -> a -> a
/ PoolStake
total) ((PoolStake, NonEmpty RelayAccessPoint)
 -> (PoolStake, NonEmpty RelayAccessPoint))
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(PoolStake, NonEmpty RelayAccessPoint)]
pl'
    in
    Bool
-> [(PoolStake, NonEmpty RelayAccessPoint)]
-> [(PoolStake, NonEmpty RelayAccessPoint)]
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (let total' :: PoolStake
total' = [PoolStake] -> PoolStake
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([PoolStake] -> PoolStake) -> [PoolStake] -> PoolStake
forall a b. (a -> b) -> a -> b
$ ((PoolStake, NonEmpty RelayAccessPoint) -> PoolStake)
-> [(PoolStake, NonEmpty RelayAccessPoint)] -> [PoolStake]
forall a b. (a -> b) -> [a] -> [b]
map (PoolStake, NonEmpty RelayAccessPoint) -> PoolStake
forall a b. (a, b) -> a
fst [(PoolStake, NonEmpty RelayAccessPoint)]
pl''
            in PoolStake
total PoolStake -> PoolStake -> Bool
forall a. Eq a => a -> a -> Bool
== PoolStake
0 Bool -> Bool -> Bool
|| (PoolStake
total' PoolStake -> PoolStake -> Bool
forall a. Ord a => a -> a -> Bool
> (Rational -> PoolStake
PoolStake (Rational -> PoolStake) -> Rational -> PoolStake
forall a b. (a -> b) -> a -> b
$ Integer
999999 Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
% Integer
1000000) Bool -> Bool -> Bool
&&
                  PoolStake
total' PoolStake -> PoolStake -> Bool
forall a. Ord a => a -> a -> Bool
< (Rational -> PoolStake
PoolStake (Rational -> PoolStake) -> Rational -> PoolStake
forall a b. (a -> b) -> a -> b
$ Integer
1000001 Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
% Integer
1000000))
           )
    [(PoolStake, NonEmpty RelayAccessPoint)]
pl''
  where
    adjustment :: PoolStake -> PoolStake
    adjustment :: PoolStake -> PoolStake
adjustment =
      case LedgerPeersKind
ledgerPeersKind of
        LedgerPeersKind
AllLedgerPeers ->
          -- We do loose some precision in the conversion. However we care about
          -- precision in the order of 1 block per year and for that a Double is
          -- good enough.
          Rational -> PoolStake
PoolStake (Rational -> PoolStake)
-> (PoolStake -> Rational) -> PoolStake -> PoolStake
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Rational
forall a. Real a => a -> Rational
toRational (Double -> Rational)
-> (PoolStake -> Double) -> PoolStake -> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt @Double (Double -> Double) -> (PoolStake -> Double) -> PoolStake -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> Double
forall a. Fractional a => Rational -> a
fromRational (Rational -> Double)
-> (PoolStake -> Rational) -> PoolStake -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PoolStake -> Rational
unPoolStake
        LedgerPeersKind
BigLedgerPeers ->
          PoolStake -> PoolStake
forall a. a -> a
id