{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TemplateHaskell            #-}

-- | Representation of supply specifications (including supply curves
-- and supply orderings) for all kinds of Product-Mix Auction.
module ProductMixAuction.Supply
  ( -- * Individual supply curves
    Step(..)
  , SupplyCurve(..)

  , mkEmptySupplyCurve
  , mkFixedSupplyCurve
  , mkFiniteSupplyCurve
  , mkInfiniteSupplyCurve
  , mkSupplyCurve

  , unpackSupplyCurve
  , totalSupplyCurve
  , showSupplyCurve
  , isNondecreasingSupplyCurve
  , firstReservePrice

    -- * Sets of supply curves
  , Supply(..)
  , SupplyCurveAndCover(..)
  , SupplyOrdering(..)
  , SupplyDir
  , mkHorizontalSupply
  , mkVerticalSupply
  , mkFixedSupply
  , mkTabularSupply
  , mkTabularSupplyWithBase

  , lookupSupplyCurve
  , listSupplyCurves
  , numSupplyCurves
  , supplyGoods
  , supplyGoodsAbove
  , supplyGoodsAboveAll
  , supplyGoodBelow
  , firstGoods
  , isFirstGood
  , isLastGood
  , isHorizontalSupply
  , isVerticalSupply
  , totalSupply
  , showSupply
  , allocAbove
  , relativePrice
  , setSupplyCurve
  , decreasingSupply

    -- * Scaling/rounding supply curves
  , SupplyScaleFunction
  , SupplyScaleLambda(..)
  , scaleSupplyCurve
  , scaleSupply
  , scaleSupplyRounded
  , ceilingSupply
  , ceilingSupplyCurve

    -- * Arbitrary supply curve generation
  , arbitrarySupply
  , arbitrarySupplyWithDir
  , arbitraryStep
  , arbitrarySteps
  , arbitrarySupplyCurve
  , arbitraryFiniteSupplyCurve

    -- * Lenses
  , step_units
  , step_price
  , sc_steps
  , sc_final_price
  , scac_curve
  , scac_cover
  , supply_ordering
  , supply_curves
  ) where

import Control.Lens hiding ((.=))
import Control.Monad
import Data.Aeson.Types
import Data.Csv (FromField(..), FromRecord(..), (.!))
import Data.Default.Class
import Data.List
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set
import qualified Test.QuickCheck as QC
import GHC.Generics

import ProductMixAuction.Types


-- | A single step of a supply curve, parameterised in the
-- representation of quantities of goods.  Note that the step length
-- is relative, i.e. the width of this single step rather than the
-- total number of units sold up to this point.
data Step u = Step { _step_units :: u      -- ^ @s^q_j@, the step length
                   , _step_price :: Price  -- ^ @mu^q_j@, the margin at this step
                   }
  deriving (Eq, Functor, Read, Show, Generic)

$(makeLenses ''Step)

instance Default u => Default (Step u)

instance FromField u => FromRecord (Step u) where
  parseRecord v
    | length v == 2 = Step <$> v .! 0 <*> v .! 1
    | otherwise     = mzero

instance ToJSON u => ToJSON (Step u) where
  toJSON = genericToJSON $ jsonOptions "_step_"

instance FromJSON u => FromJSON (Step u) where
  parseJSON = genericParseJSON $ jsonOptions "_step_"

-- | A supply "curve" is in fact a step function, with the price
-- increasing at each step.  A final price may optionally be
-- specified, indicating that there is an unlimited quantity available
-- at that price.
--
-- Quantities are those for the good to which this supply curve
-- belongs, and all later goods (if any).  Prices are expressed
-- relative to the preceding good (or to "no sale", if there is no
-- preceding good).
data SupplyCurve u = SupplyCurve { _sc_steps       :: [Step u]    -- ^ Finite-width steps in the supply curve
                                 , _sc_final_price :: Maybe Price -- ^ Optional price for infinite-width step (not fully supported)
                                 }
  deriving (Eq, Functor, Read, Show, Generic)

$(makeLenses ''SupplyCurve)

instance Default (SupplyCurve u) where
  def = mkEmptySupplyCurve

instance ToJSON u => ToJSON (SupplyCurve u) where toJSON = genericToJSON $ jsonOptions "_sc_"

instance FromJSON u => FromJSON (SupplyCurve u) where parseJSON = genericParseJSON $ jsonOptions "_sc_"

-- | Make a supply curve with no quantity available at all.
mkEmptySupplyCurve :: SupplyCurve u
mkEmptySupplyCurve = SupplyCurve [] Nothing

-- | Make a single-step supply curve with a fixed quantity available
-- and no price margin over the previous good.
mkFixedSupplyCurve :: (Ord u, Num u) => u -> SupplyCurve u
mkFixedSupplyCurve u = mkFiniteSupplyCurve [Step u 0]

-- | Make a stepped supply curve with a finite total quantity
-- (i.e. going off vertically to infinity).  Steps of width 0 will be
-- ignored.
mkFiniteSupplyCurve :: (Ord u, Num u) => [Step u] -> SupplyCurve u
mkFiniteSupplyCurve steps = mkSupplyCurve steps Nothing

-- | Make a stepped supply curve with an infinite quantity available
-- (i.e. going off horizontally to infinity).  Steps of width 0 will
-- be ignored.
--
-- This is not fully supported and may be removed.
mkInfiniteSupplyCurve :: (Ord u, Num u) => [Step u] -> Price -> SupplyCurve u
mkInfiniteSupplyCurve steps p = mkSupplyCurve steps (Just p)

-- | Make a supply curve with an optional final price, filtering out
-- any zero-width steps.
mkSupplyCurve :: (Ord u, Num u) => [Step u] -> Maybe Price -> SupplyCurve u
mkSupplyCurve steps mb_price = SupplyCurve (filter ((> 0) . view step_units) steps) mb_price

-- | View a supply curve as a list of indexed steps, satisfying the invariant
-- that only the final step may have 'Nothing' for the quantity.
unpackSupplyCurve :: SupplyCurve u -> [(Int, Step (Maybe u))]
unpackSupplyCurve sc = zip [1..] (help (sc ^. sc_steps))
  where
    help (step:steps) = fmap Just step : help steps
    help []           = case sc ^. sc_final_price of
                          Nothing -> []
                          Just p  -> [Step Nothing p]

-- | Total quantity available in (the finite part of) the supply
-- curve, i.e. the total width of the steps.
totalSupplyCurve :: Num u => SupplyCurve u -> u
totalSupplyCurve = sumOf (sc_steps . each . step_units)

-- | Multiply all the step lengths in the supply curve by the given
-- scaling factor.
scaleSupplyCurve :: Ratio -> SupplyCurve Units -> SupplyCurve Units
scaleSupplyCurve x = sc_steps . each . step_units *~ Units (_Ratio x)

-- | Render a supply curve as a list of (quantity, price) pairs.
showSupplyCurve :: Show u => SupplyCurve u -> String
showSupplyCurve sc = intercalate ", " (map showStep (unpackSupplyCurve sc))
  where
    showStep (_, Step mb_s mu) = "(" ++ maybe "Infinity" show mb_s
                                     ++ ", " ++ show (_Price mu) ++ ")"

-- | Check that a supply curve is non-decreasing.
isNondecreasingSupplyCurve :: SupplyCurve u -> Bool
isNondecreasingSupplyCurve sc = go 0 (sc ^. sc_steps)
  where
    go p []     = maybe True (>= p) (sc ^. sc_final_price)
    go p (s:ss) = s^.step_price >= p && go (s^.step_price) ss

-- | Extract the (first) reserve price from a supply curve.  Returns 'Nothing'
-- if the supply curve has no steps.
firstReservePrice :: SupplyCurve u -> Maybe Price
firstReservePrice sc = case sc ^. sc_steps of
  Step _ p : _ -> Just p
  []           -> Nothing


-- | Data stored about each good in a 'Supply'.
data SupplyCurveAndCover u = SupplyCurveAndCover
  { _scac_curve :: SupplyCurve u
    -- ^ Supply curve for the good.
  , _scac_cover :: [Good]
    -- ^ Covers (i.e. immediate successors) of the good in the partial
    -- order.
  }
  deriving (Eq, Read, Show, Functor, Generic)

$(makeLenses ''SupplyCurveAndCover)

instance Default (SupplyCurveAndCover u) where
  def = SupplyCurveAndCover def []

instance ToJSON u => ToJSON (SupplyCurveAndCover u) where
  toJSON = genericToJSON $ jsonOptions "_scac_"
instance FromJSON u => FromJSON (SupplyCurveAndCover u) where
  parseJSON = genericParseJSON $ jsonOptions "_scac_"


-- | Ordering on goods in a 'Supply'.
data SupplyOrdering
    = Horizontal -- ^ Horizontal ordering: each good is priced
                 -- relative to selling nothing at all, using the
                 -- identity relation as the partial order.
    | Vertical   -- ^ Vertical ordering: each good is priced relative
                 -- to the previous good and below all succeeding
                 -- goods.
    | Other      -- ^ An arbitrary partial order.
  deriving (Eq, Show, Generic)

instance Default SupplyOrdering where
  def = Vertical

instance ToJSON   SupplyOrdering where toJSON    = genericToJSON $ jsonOptions ""
instance FromJSON SupplyOrdering where parseJSON = genericParseJSON $ jsonOptions ""


-- | The supply for an auction consists of:
--
--  * a set of goods (represented by the domain of the map);
--
--  * a partial order on this set (represented by listing the covers
--  for each good, i.e. its immediate successors);
--
--  * a 'SupplyCurve' for each good, expressing how the reserve price
--  relative to the previous good increases as the quantity sold of
--  this good and all larger goods increases.
--
-- The partial order must be a subset of the total order induced by
-- the 'Ord' instance on goods (i.e. if @j1@ has a cover @j2@ then we
-- must have @j1 < j2 == True@).
--
-- It is not yet clear that arbitrary partial orders make sense
-- economically, but it is a convenient generalisation for
-- implementation purposes.  The smart constructors below implement
-- partial orders that lead to sensible results.
--
-- Note that the representation of the partial order by its covers
-- does not enforce that it is antisymmetric.  Violating this property
-- may create a cyclic graph, leading to nontermination or other
-- nonsense.
data Supply u =
    Supply { _supply_ordering :: SupplyOrdering
             -- ^ Whether the supply ordering is horizontal, vertical
             -- or an arbitrary partial order.  This must be
             -- consistent with the partial order defined by the
             -- covers for each good.
           , _supply_curves   :: Map.Map Good (SupplyCurveAndCover u)
             -- ^ For each good in the supply, its supply curve and
             -- the set of covers that defines the partial order.
           }
  deriving (Eq, Show, Functor, Generic)

instance Default (Supply u) where
  def = Supply def mempty

instance ToJSON u => ToJSON   (Supply u) where toJSON = genericToJSON $ jsonOptions "_supply_"
instance FromJSON u => FromJSON (Supply u) where parseJSON = genericParseJSON $ jsonOptions "_supply_"

$(makeLenses ''Supply)

type SupplyDir u = [(G, SupplyCurve u)] -> Supply u

-- | Make a horizontal supply, where all goods are priced relative to
-- selling nothing at all, using the identity relation as the partial
-- order.  The order of the list is unimportant.
mkHorizontalSupply :: [(Good, SupplyCurve u)] -> Supply u
mkHorizontalSupply = Supply Horizontal . Map.fromList . map (fmap (\ sc -> SupplyCurveAndCover sc []))

-- | Make horizontal supply curves corresponding to a list of fixed
-- total quantities of each good.
mkFixedSupply :: (Ord u, Num u) => [(Good, u)] -> Supply u
mkFixedSupply = mkHorizontalSupply . map (fmap mkFixedSupplyCurve)

-- | Make a vertical supply, where each good is priced relative to the
-- previous element of the list and below all the succeeding elements.
mkVerticalSupply :: [(Good, SupplyCurve u)] -> Supply u
mkVerticalSupply = Supply Vertical . Map.fromList . withTails
  where
    withTails xs = zipWith f xs (map (pure . fst) (tail xs) ++ [[]])
    f (g, sc) ys = (g, SupplyCurveAndCover sc ys)

-- | Make a tabular supply, where each element of the outer list of
-- lists gives a vertical column, and the columns themselves are
-- ordered horizontally.
mkTabularSupply :: [[(Good, SupplyCurve u)]] -> Supply u
mkTabularSupply = Supply Other . Map.unions . map (_supply_curves . mkVerticalSupply)

-- | Make a tabular supply, where each element of the outer list of
-- lists gives a vertical column, and the columns themselves are
-- ordered horizontally.  A single good is below all the columns.
mkTabularSupplyWithBase :: (Good, SupplyCurve u) -> [[(Good, SupplyCurve u)]] -> Supply u
mkTabularSupplyWithBase (g, sc) cols =
    Supply Other (Map.insert g (SupplyCurveAndCover sc covers) (_supply_curves (mkTabularSupply cols)))
  where
    covers = map fst (mapMaybe listToMaybe cols)

-- | Look up the supply curve for a particular good, returning the
-- empty supply curve if the good is not available at all.
lookupSupplyCurve :: Good -> Supply u -> SupplyCurve u
lookupSupplyCurve g supply = maybe mkEmptySupplyCurve _scac_curve $ Map.lookup g (_supply_curves supply)

-- | List all the goods with their corresponding supply curves.
listSupplyCurves :: Supply u -> [(Good, SupplyCurve u)]
listSupplyCurves = map (fmap _scac_curve) . Map.toList . _supply_curves

-- | Number of supply curves in the supply.
numSupplyCurves :: Supply u -> Int
numSupplyCurves = Map.size . _supply_curves

-- | List all the goods from the supply.
supplyGoods :: Supply u -> [Good]
supplyGoods = Map.keys . _supply_curves

-- | List all goods that are immediate successors of the given good.
supplyGoodsAbove :: Good -> Supply u -> [Good]
supplyGoodsAbove g supply = maybe [] _scac_cover (Map.lookup g (_supply_curves supply))

-- | List all goods that are (transitively) non-strictly larger than
-- the given good.
supplyGoodsAboveAll :: Good -> Supply u -> [Good]
supplyGoodsAboveAll g0 supply = help g0
  where
    help g = g : concatMap help (supplyGoodsAbove g supply)

-- | Find the good that is the immediate predecessor of the given
-- good, if there is one.
supplyGoodBelow :: Good -> Supply u -> Maybe Good
supplyGoodBelow g supply = fst <$> find (\ (_, SupplyCurveAndCover _ gs) -> g `elem` gs) (Map.toList (_supply_curves supply))

-- | Calculate the set of all "bottom" goods in the supply, i.e. those
-- priced relative to selling nothing at all.  This will be all the
-- goods for a horizontal supply, or just the first good for a
-- vertical supply.
firstGoods :: Supply u -> Set.Set Good
firstGoods supply = Map.keysSet (_supply_curves supply) Set.\\ covers
  where
    covers = Set.fromList (concat (map _scac_cover (Map.elems (_supply_curves supply))))

-- | Is this good a "bottom" good?
isFirstGood :: Good -> Supply u -> Bool
isFirstGood g supply = g `Set.member` firstGoods supply

-- | Is this good a "top" good?
isLastGood :: Good -> Supply u -> Bool
isLastGood g supply = null (supplyGoodsAbove g supply)

-- | Is this supply ordered horizontally?
isHorizontalSupply :: Supply u -> Bool
isHorizontalSupply supply = _supply_ordering supply == Horizontal

-- | Is this supply ordered vertically?
isVerticalSupply :: Supply u -> Bool
isVerticalSupply supply = _supply_ordering supply == Vertical

-- | Calculate the total quantity available from (the finite parts of)
-- the supply curves for the first goods.
totalSupply :: Num u => Supply u -> u
totalSupply supply = sum (firstGoodSupplies supply)

-- | List the total supply curve widths for all the first goods.
firstGoodSupplies :: Num u => Supply u -> [u]
firstGoodSupplies supply = [ totalSupplyCurve sc
                           | (good, sc) <- listSupplyCurves supply
                           , good `Set.member` goods
                           ]
  where
    goods = firstGoods supply

-- | Scale a supply to the given total size.  First goods will be
-- scaled in proportion to the new total size.  Subsequent goods will
-- be scaled depending on the parameter @lambda@, which must be in the
-- interval @[0,1]@:
--
--  * @lambda = 0@ scales all goods in proportion to the new total size;
--
--  * @lambda = 1@ does not scale goods other than the first goods.
scaleSupply :: SupplyScaleLambda -> Units -> Supply Units -> Supply Units
scaleSupply (SupplyScaleLambda (Ratio lambda)) (Units new_r) supply = over supply_curves (Map.mapWithKey f) supply
  where
    firsts = firstGoods supply
    this_r = _Units (totalSupply supply)

    first_ratio = Ratio (new_r / this_r)
    other_ratio = Ratio ((lambda*this_r +(1-lambda)*new_r) / this_r)

    f good (SupplyCurveAndCover sc covers) = SupplyCurveAndCover (scaleSupplyCurve ratio sc) covers
      where
        ratio | good `Set.member` firsts = first_ratio
              | otherwise                = other_ratio


-- | Function to rescale a set of supply curves, given a new total
-- size of the supply.
type SupplyScaleFunction = Units -> Supply Units -> Supply Units

-- | Parameter to control how supply is scaled when a TQSS is in use.
-- See 'scaleSupply' for how this is used.
newtype SupplyScaleLambda = SupplyScaleLambda { _ssl_lambda :: Ratio }
  deriving (Eq, Generic, Show, QC.Arbitrary)

instance Default  SupplyScaleLambda where def = SupplyScaleLambda 0
instance ToJSON   SupplyScaleLambda where toJSON = genericToJSON $ jsonOptions "_ssl_"
instance FromJSON SupplyScaleLambda where parseJSON = genericParseJSON $ jsonOptions "_ssl_"

-- | Scale the supply curves based on the 'SupplyScaleLambda', then
-- round the step widths to the 'ScaleFactor'.
scaleSupplyRounded :: ScaleFactor -> SupplyScaleLambda -> SupplyScaleFunction
scaleSupplyRounded sf lambda u = ceilingSupply sf . scaleSupply lambda u

-- | Round all the supply curves using 'ceilingSupplyCurve'.
ceilingSupply :: RealFrac u => ScaleFactor -> Supply u -> Supply u
ceilingSupply sf = over (supply_curves.each.scac_curve) (ceilingSupplyCurve sf)

-- | Round the widths of the supply curve steps to a precision given
-- by the 'ScaleFactor'.  This rounds upwards the cumulative total
-- width of all the steps.
ceilingSupplyCurve :: RealFrac u => ScaleFactor -> SupplyCurve u -> SupplyCurve u
ceilingSupplyCurve sf = over sc_steps (go 0)
  where
    -- Accumulator x represents error carried forward from
    -- approximating previous step.
    go _ []           = []
    go x (step:steps)
      | y > 0     = (step & step_units .~ y)
                    : go (y - z) steps
      | otherwise = go (y - z) steps
      where
        z = step ^. step_units - x
        y = ceilingToScale sf z


-- | Render a set of supply curves in a vaguely human-readable format.
showSupply :: Show u => Supply u -> String
showSupply supply = unlines (map showGood (Map.toList (_supply_curves supply)))
  where
    showGood (g, (SupplyCurveAndCover sc covers)) = show g ++ ": " ++ showSupplyCurve sc ++ showCovers covers

    showCovers []     = ""
    showCovers covers = " -- covered by: " ++ show covers


-- | Given a supply and the amounts allocated to each good, calculate
-- the total allocated to a good and all its successors.
allocAbove :: Num u => Supply u -> Map.Map Good u -> Good -> u
allocAbove supply good_allocs good = my_alloc + successor_allocs
  where
    my_alloc         = Map.findWithDefault 0 good good_allocs
    successor_allocs = sum (map (allocAbove supply good_allocs) (supplyGoodsAbove good supply))

-- | Look up the price of a good relative to the preceding good (or to
-- no sale, if there is no preceding good).
relativePrice :: Supply u -> Map.Map Good Price -> Good -> Price
relativePrice supply ps good = case supplyGoodBelow good supply of
                                 Just good' -> lookupPrice good - lookupPrice good'
                                 Nothing    -> lookupPrice good
  where
    lookupPrice g = fromMaybe 0 (Map.lookup g ps)

-- | Set the supply of a single good within the set of supply curves.
setSupplyCurve :: Good -> SupplyCurve u -> Supply u -> Supply u
setSupplyCurve good sc supply = supply { _supply_curves = _supply_curves supply & at good . _Just . scac_curve .~ sc }


-- | Find the first good with a decreasing supply curve.  For a valid
-- supply this should return 'Nothing'.
decreasingSupply :: Supply u -> Maybe Good
decreasingSupply supply = fst <$> ifind (const (not . isNondecreasingSupplyCurve . _scac_curve)) (_supply_curves supply)


-- | Generate an arbitrary supply curve step, given generators for the
-- step width (quantity available) and height (price margin).
arbitraryStep :: QC.Gen u -> QC.Gen Price -> QC.Gen (Step u)
arbitraryStep arbU arbP = Step <$> arbU <*> arbP

-- | Generate an arbitrary supply curve, given generators for the
-- number of supply curve steps, step width and step height.
arbitrarySupplyCurve :: Num u => QC.Gen Int -> QC.Gen u -> QC.Gen Price -> QC.Gen (SupplyCurve u)
arbitrarySupplyCurve arbS arbU arbP =
    makeIncreasing <$> (flip replicateM (arbitraryStep arbU arbP) =<< arbS)
                   <*> QC.oneof [pure Nothing, Just <$> arbP]

-- | Generate an arbitrary finite supply curve, given generators for
-- the number of supply curve steps, step width and step height.
arbitraryFiniteSupplyCurve :: Num u => QC.Gen Int -> QC.Gen u -> QC.Gen Price -> QC.Gen (SupplyCurve u)
arbitraryFiniteSupplyCurve arbS arbU arbP
  = flip makeIncreasing Nothing <$> (flip replicateM (arbitraryStep arbU arbP) =<< arbS)

arbitrarySteps :: Num u => QC.Gen Int -> QC.Gen u -> QC.Gen Price -> QC.Gen [Step u]
arbitrarySteps arbS arbU arbP = makeIncreasingSteps <$> (flip replicateM (arbitraryStep arbU arbP) =<< arbS)

-- | Make the list of relative steps increasing.
makeIncreasingSteps :: [Step u] -> [Step u]
makeIncreasingSteps (step:steps) = scanl (\ s1 s2 -> s2 & step_price +~ s1 ^. step_price) step steps
makeIncreasingSteps []           = []

-- | Turn the list of relative steps into an increasing supply curve.
makeIncreasing :: [Step u] -> Maybe Price -> SupplyCurve u
makeIncreasing (step:steps) mb_fp = SupplyCurve (scanl (\ s1 s2 -> s2 & step_price +~ s1 ^. step_price) step steps)
                                                    ((sumOf (each.step_price) (step:steps) +) <$> mb_fp)
makeIncreasing [] mb_fp           = SupplyCurve [] mb_fp

-- | Generate an arbitrary set of supply curves for the given goods,
-- with either a horizontal or a vertical ordering.
arbitrarySupply :: Num u => [Good] -> QC.Gen Int -> QC.Gen u -> QC.Gen Price -> QC.Gen (Supply u)
arbitrarySupply = arbitrarySupplyWithDir (QC.elements [mkVerticalSupply, mkHorizontalSupply])

-- | Generate an arbitrary set of supply curves for the given goods, with a
-- supplied generator for supply orderings, as well as generators for the number
-- of steps, step width and step height.
arbitrarySupplyWithDir :: Num u => QC.Gen (SupplyDir u) -> [Good] -> QC.Gen Int -> QC.Gen u -> QC.Gen Price -> QC.Gen (Supply u)
arbitrarySupplyWithDir arbDir goods arbS arbU arbP =
    arbDir <*> mapM supplyCurveFor goods
  where
    supplyCurveFor good = (,) good <$> arbitrarySupplyCurve arbS arbU arbP