{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Implements different strategies for rationing (sharing limited
-- quantities of goods between multiple marginal bidders).
module ProductMixAuction.LP.Rationing
  ( -- * Rationing strategies
    RationingStrategy
  , rationNoRationing
  , rationLinearDemand
  , rationLinearDemandPreferPairedBids

    -- * First-order representation
  , RationingMode(..)
  , RationingSteps(..)
  , RationingOptions(..)
  , noRationing
  , linearDemand
  , linearDemandPreferPairedBids
  , runRationingOptions

    -- * Summarising allocations
  , goodAllocations
  , bidderAllocations
  , bidAllocations
  , filterAllocs

   -- * Lenses
  , ro_mode
  , ro_steps
  ) where

import Control.Lens
import Data.Aeson.Types as A
import Data.Default.Class
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.List
import Data.Maybe
import GHC.Generics

import ProductMixAuction.LP.AdditionalConstraint
import ProductMixAuction.LP.Bid
import ProductMixAuction.LP.Core
import ProductMixAuction.LP.Demand
import ProductMixAuction.Supply
import ProductMixAuction.Types


-- | Calculate the total allocation to each good, rounded to the
-- appropriate scale.  All goods are included.
goodAllocations :: ScaleFactor -> Map.Map bid (Map.Map Good Units)
                -> Map.Map Good Units
goodAllocations sf = Map.unionsWith (+) . Map.elems . Map.map (roundMap sf)

-- | Calculate the total allocation of each good to each bidder,
-- rounded to the appropriate scale.  Unsuccessful bidders are
-- excluded from the map.
bidderAllocations :: Ord bidder
                  => ScaleFactor -> Map.Map (bidder, bid) (Map.Map Good Units)
                  -> Map.Map bidder (Map.Map Good Units)
bidderAllocations sf = filterAllocs . Map.mapKeysWith (Map.unionWith (+)) fst . Map.map (roundMap sf)

-- | Round the allocations to the appropriate scale and exclude
-- unsuccessful bids from the map.
bidAllocations :: Ord bid
               => ScaleFactor -> Map.Map bid (Map.Map Good Units)
               -> Map.Map bid (Map.Map Good Units)
bidAllocations sf = filterAllocs . Map.map (roundMap sf)

-- | Make sure that every allocation to a bid is non-zero and
-- unsuccessful bids are excluded from the map.
filterAllocs :: Ord bid
             => Map.Map bid (Map.Map Good Units) -> Map.Map bid (Map.Map Good Units)
filterAllocs = Map.filter (not . Map.null) . Map.map (Map.filter (> 0))

-- | Round the quantities in the map to the given scale factor.
roundMap :: ScaleFactor -> Map.Map x Units -> Map.Map x Units
roundMap sf = Map.map (roundToScale sf)


-- | Convert the values of the primal variables into a map from all
-- bids to their allocations.  Every bid and good will be included.
varsToAllocs :: Ord bid
             => Auction bid -> Map.Map (AVar bid) Units -> Map.Map bid (Map.Map Good Units)
varsToAllocs auction vs = Map.fromList (map bidAllocation bids)
  where
    bids = auctionBidLabels auction
    goods = auctionGoodLabels auction

    bidAllocation bid = (bid, Map.fromList [ (good, fromMaybe 0 (Map.lookup (AllocVar bid good) vs))
                                           | good <- goods ])


-- | A rationing strategy takes an auction and the results of solving
-- its LP (in particular, the auction prices), and computes a
-- possibly-fairer allocation of goods.  It may not change the prices.
type RationingStrategy bid =
  Auction bid -> AuctionResult bid -> IO (Map.Map bid (Map.Map Good Units))

-- | Enumeration of supported choices of rationing strategy.
data RationingMode
  = LinearDemandPreferPairedBids -- ^ See 'rationLinearDemandPreferPairedBids'.
  | LinearDemand                 -- ^ See 'rationLinearDemand'.
  | NoRationing                  -- ^ See 'rationNoRationing'.
  deriving (Eq, Generic, Show)

instance Default  RationingMode where def = LinearDemandPreferPairedBids
instance ToJSON   RationingMode where toJSON = genericToJSON $ jsonOptions ""
instance FromJSON RationingMode where parseJSON = genericParseJSON $ jsonOptions ""

-- | Number of steps to use in rationing strategies that use a
-- step-wise approximation of linear demand.
newtype RationingSteps = RationingSteps { _RationingSteps :: Int }
  deriving (Eq, Ord, Num, Generic, Show, ToJSON, FromJSON)

-- | First-order representation of a choice of rationing strategy and
-- any required parameters.
data RationingOptions = RationingOptions
  { _ro_mode  :: RationingMode
  -- ^ Choice of strategy
  , _ro_steps :: Maybe RationingSteps
  -- ^ Number of steps to use in the approximation of linear demand.
  -- This is used only by 'LinearDemand' and
  -- 'LinearDemandPreferPairedBids'.  If it is absent or is not
  -- positive, a sensible default will be used.
  }
  deriving (Eq, Generic, Show)

-- | See 'rationNoRationing'.
noRationing :: RationingOptions
noRationing                  = RationingOptions NoRationing Nothing

-- | See 'rationLinearDemand'.
linearDemand :: Maybe RationingSteps -> RationingOptions
linearDemand                 = RationingOptions LinearDemand

-- | See 'rationLinearDemandPreferPairedBids'.
linearDemandPreferPairedBids :: Maybe RationingSteps -> RationingOptions
linearDemandPreferPairedBids = RationingOptions LinearDemandPreferPairedBids

instance Default  RationingOptions where def = RationingOptions def (Just def)
instance ToJSON   RationingOptions where toJSON = genericToJSON $ jsonOptions "_ro_"
instance FromJSON RationingOptions where parseJSON = genericParseJSON $ jsonOptions "_ro_"

instance Default  RationingSteps

-- | Translate the first-order representation into the appropriate
-- rationing strategy.
runRationingOptions :: Goodly bid => RationingOptions -> RationingStrategy bid
runRationingOptions (RationingOptions k m) =
  case k of
    NoRationing                  -> rationNoRationing
    LinearDemand                 -> rationLinearDemand  m Nothing
    LinearDemandPreferPairedBids -> rationLinearDemandPreferPairedBids m Nothing

-- | Rationing strategy A: let the linear programme do what it will
-- do, and don't do anything about it.
--
-- The LP solver may make arbitrary choices as to which marginal bids
-- receive goods and which do not.  There is no guarantee of fairness
-- between bids, but this is simple and fast.
rationNoRationing :: Ord bid => RationingStrategy bid
rationNoRationing auction ar =
  return $ varsToAllocs auction (ar ^. ar_quantities)


-- | Rationing strategy D: tweak marginal bids to approximate a linear
-- demand around the bid price, then re-run the auction to find
-- adjusted allocations.
rationLinearDemand :: Goodly bid => Maybe RationingSteps -> Maybe (Map.Map Good Int) -> RationingStrategy bid
rationLinearDemand = rationD_or_Di False

-- | Rationing strategy D(i): tweak multiply-marginal bids to
-- approximate a linear demand around the bid price, re-run the
-- auction to find adjusted allocations, then reallocate to treat
-- equally bids that are marginal on single goods.
--
-- This is the default strategy.
rationLinearDemandPreferPairedBids :: Goodly bid => Maybe RationingSteps -> Maybe (Map.Map Good Int) -> RationingStrategy bid
rationLinearDemandPreferPairedBids = rationD_or_Di True

-- | Rationing strategy D or D(i): tweak marginal bids to approximate
-- a linear demand around the bid price, then re-run the auction to
-- find adjusted allocations.
--
-- If the flag is true, apply the tweaks only to paired bids that are
-- marginal on multiple goods, and subsequently reallocate to treat
-- equally bids that are marginal on single goods.
--
-- The number of steps D affects the fairness of the resulting
-- allocation.  The maximum unfair share (i.e. the amount by which one
-- bidder may be arbitrarily favoured over another) is the maximum
-- quantity of any bid being tweaked, divided by D.  If a step count
-- is not supplied, it will be determined based on the size of the
-- bids and the scale factor; this should give fair results but may
-- lead to large numbers of steps and hence slow calculation.
--
-- Optionally, a map from goods j to weights lambda_j may be provided
-- (see appendix A of "Examples for the Product-Mix Auction").  This
-- affects how goods are allocated to marginal bidders.  If it is not
-- supplied, weights will be assigned sequentially from 1 to the
-- number of goods.
rationD_or_Di :: forall bid . Goodly bid
              => Bool -> Maybe RationingSteps -> Maybe (Map.Map Good Int)
              -> RationingStrategy bid
rationD_or_Di prefer_paired_bids mb_steps mb_weights auction ar
  | needs_rationing = do ar' <- runAuctionCore "ration" auction'
                         return $ share_out (joinBidAllocations (ar' ^. ar_quantities))
  | otherwise       = return $ share_out (ar ^. ar_quantities)
  where
    ps              = ar ^. ar_prices
    steps           = case mb_steps of
                        Just s | s > 0 -> s
                        _              -> defaultSteps prefer_paired_bids ps auction

    weights         = fromMaybe default_weights mb_weights
    default_weights = Map.fromList (zip (auctionGoodLabels auction) [1..])

    auction' = fixSupply (ar ^. ar_quantities) (approximateBids prefer_paired_bids steps ps weights auction)

    bs = Map.fromList [ (bid ^. bid_label, bid) | bid <- _auction_bids auction ]

    share_out :: Map.Map (AVar bid) Units -> Map.Map bid (Map.Map Good Units)
    share_out = shareMatching auction . varsToAllocs auction . redistribute
    redistribute | prefer_paired_bids = evenlyRedistribute bs ps
                 | otherwise          = id

    -- Whether we need to re-run the LP with tweaked bids:
    needs_rationing
         -- If preferring paired bids, this is needed if there are two
         -- multiply-marginal bids, or one multiply-marginal bid and a
         -- singly-marginal bid it needs to beat.
      | prefer_paired_bids = multiply_marginal > 1 || (multiply_marginal == 1 && singly_marginal > 0)
         -- If not, this is needed if there are two or more marginal
         -- bids (regardless of the number of goods involved).
      | otherwise          = singly_marginal + multiply_marginal > 1
    (singly_marginal, multiply_marginal) = countMarginalBids ps (_auction_bids auction)

-- | The initial run of the auction determined the prices and the
-- total quantity to be supplied of each good.  Thus for the rationed
-- auction, we fix the quantity of each good to be supplied, and
-- disable the price and quantity tweaks.
fixSupply :: Map.Map (AVar bid) Units -> Auction bid' -> Auction bid'
fixSupply zs auction = auction & auction_supply       .~ supply
                               & auction_preferences  .~ []    -- Disable price tweaks
                               & auction_tweak_supply .~ False
  where
    supply        = mkFixedSupply [ (good, totalFor good) | good <- auctionGoodLabels auction ]
    totalFor good = sum [ roundToScale sf u | (AllocVar _ good', u) <- Map.toList zs, good == good' ]
    sf = auction ^. auction_scale

-- | Choose a default number of steps D with which to approximate the
-- linear demand of marginal bids.  This should be large enough to
-- ensure that after rounding based on the auction scale, the
-- arbitrary unfairness is limited to a difference in the least
-- significant digit.
defaultSteps :: Bool -> Map.Map Good Price -> Auction bid -> RationingSteps
defaultSteps prefer_paired_bids ps auction = RationingSteps (1 + ceiling (_Units max_k) * 2 * 10 ^ rho)
  where
    -- Maximum quantity of any bid that will be approximated
    -- (TODO: can we do better for generalised bids?)
    max_k           = maximum (0 : map _bid_quantity (filter (isMarginal prefer_paired_bids ps) (_auction_bids auction)))
    ScaleFactor rho = _auction_scale auction

-- | Approximate the marginal (parts of) bids with a linear demand
-- decreasing from the bid quantity to zero in a small interval over
-- the bid price.  The labels of bids gain an integer component for
-- the step number.
approximateBids :: Ord bid
                => Bool -> RationingSteps -> Map.Map Good Price -> Map.Map Good Int
                -> Auction bid -> Auction (bid, Int)
approximateBids prefer_paired_bids (RationingSteps steps) ps weights auction =
    auction { _auction_bids        = concatMap approx              (_auction_bids auction)
            , _auction_constraints = map       transformConstraint (_auction_constraints auction)
            }
  where
    -- If the bid is considered marginal, approximate it by a number
    -- of bid steps; otherwise, leave it alone (modulo changing label).
    approx bid = case marginalGoods prefer_paired_bids ps bid of
                   Nothing -> [bid & bid_label %~ flip (,) 0]
                   Just gs -> map (toStep bid gs) [0..steps-1]

    -- Convert a bid to the l^th step, increasing the prices bid on
    -- marginal goods but leaving non-marginal prices alone.
    toStep bid gs l = Bid { _bid_label    = (bid ^. bid_label, l)
                          , _bid_name     = bid ^. bid_name
                          , _bid_quantity = bid ^. bid_quantity / fromIntegral steps
                          , _bid_values   = Map.mapWithKey (f gs l) (bid ^. bid_values)
                          }
    f gs l good bv
      | good `Set.member` gs = bv & bv_price +~ fromIntegral (l * weight good) / recip_delta_over_d
                                  & bv_quantity //~ fromIntegral steps
      | otherwise            = bv & bv_quantity //~ fromIntegral steps

    -- Reciprocal of the price increase for each step, divided by the
    -- number of steps.
    recip_delta_over_d = fromIntegral (1 + dim * steps)
    dim = Map.size ps

    -- Replace each additional constraint by a constraint on the
    -- allocations to the transformed bids.
    transformConstraint ac = ac & ac_coefficients %~ Map.fromList . concatMap transformCoefficient . Map.toList
    transformCoefficient (b, gs) = case lookupBid b auction of
                                     Just bid | isMarginal prefer_paired_bids ps bid -> [ ((b, l), gs) | l <- [0..steps-1] ]
                                     _                                               -> [((b, 0), gs)]

    weight good = fromMaybe 0 (Map.lookup good weights)


data Marginality = Unsuccessful | NonMarginal | SinglyMarginal Good | MultiplyMarginal (Set.Set Good)

-- | Count the numbers of singly-marginal and multiply-marginal bids
-- in the list of bids.
countMarginalBids :: Map.Map Good Price -> [Bid bid] -> (Int, Int)
countMarginalBids ps = foldl' f (0, 0)
  where
    f c@(!singly_marginal, !multiply_marginal) bid = case bidMarginality ps bid of
        SinglyMarginal{}   -> (singly_marginal+1, multiply_marginal)
        MultiplyMarginal{} -> (singly_marginal, multiply_marginal+1)
        _                  -> c

-- | Determine whether the bid is considered marginal at the given
-- prices (ignoring singly-marginal bids if we are preferring paired
-- bids).
isMarginal :: Bool -> Map.Map Good Price -> Bid bid -> Bool
isMarginal prefer_paired_bids ps bid = isJust (marginalGoods prefer_paired_bids ps bid)

-- | If we are preferring paired bids, require that the bid be
-- marginal on more than one good; otherwise, require that it is
-- marginal on at least one good.
marginalGoods :: Bool -> Map.Map Good Price -> Bid bid -> Maybe (Set.Set Good)
marginalGoods prefer_paired_bids ps bid = case bidMarginality ps bid of
    Unsuccessful                          -> Nothing
    NonMarginal                           -> Nothing
    SinglyMarginal g | prefer_paired_bids -> Nothing
                     | otherwise          -> Just (Set.singleton g)
    MultiplyMarginal gs                   -> Just gs

-- | Determine whether the bid is marginal at the given prices, and if
-- so, calculate the goods on which it is marginal.
--
-- The bid is:
--
--  * unsuccessful if the maximal surplus is negative;
--
--  * singly-marginal if the maximum surplus is zero and
--    is achieved by precisely one good;
--
--  * multiply-marginal if:
--
--    * the bid is not generalised, and more than one good achieves
--      the maximal surplus; or
--
--    * the bid is generalised, and a good achieves a surplus of zero
--      or multiple goods achieve the same non-negative surplus
--      (neither need necessarily be the maximum surplus).
--
--  * non-marginal otherwise.
--
-- Note that a generalised bid may be multiply-marginal on a single
-- good. For example, consider a bid for up to 2 units in total of (up
-- to 1 apple at price 10) or (up to 1 banana at price 0), with
-- auction prices of 0 on both goods.  We consider this
-- multiply-marginal rather than singly-marginal, because if
-- preferring paired bids we cannot simply give the bid a share of the
-- available bananas.
bidMarginality :: Map.Map Good Price -> Bid bid -> Marginality
bidMarginality ps bid
  | Map.null (bid ^. bid_values) = NonMarginal
  | m < 0                        = Unsuccessful
  | m == 0, [g] <- Set.toList gs = SinglyMarginal g
  | is_generalised_marginal      = MultiplyMarginal gs'
  | non_generalised_marginal     = MultiplyMarginal gs
  | otherwise                    = NonMarginal
  where
    (m, gs, m', gs') = bidSurplus ps bid

    -- A generalised bid is marginal if there is a good with surplus
    -- zero or multiple goods with equal surplus.
    is_generalised_marginal = isGeneralised bid && case Set.size gs' of
                                                     0 -> False
                                                     1 -> m' == 0
                                                     _ -> True

    -- A non-generalised bid is marginal if multiple goods achieve the
    -- maximal surplus.
    non_generalised_marginal = not (isGeneralised bid) && Set.size gs >= 2


-- | Collapse the allocation variables for the different bid steps
-- into total allocations for the original bids.
joinBidAllocations :: Ord bid => Map.Map (AVar (bid, Int)) Units -> Map.Map (AVar bid) Units
joinBidAllocations = Map.mapKeysWith (+) f
  where
    f :: AVar (bid, Int) -> AVar bid
    f (AllocVar (bid, _) good) = AllocVar bid good
    f (StepVar good i)         = StepVar good i
    f (ExtraAllocVar good)     = ExtraAllocVar good
    f (TQSSStepVar l)          = TQSSStepVar l

-- | Redistribute the allocations to treat equally bids that are
-- marginal on single goods (paired bids that are marginal on multiple
-- goods will have been preferred already).
evenlyRedistribute :: Ord bid
                   => Map.Map bid (Bid bid) -> Map.Map Good Price
                   -> Map.Map (AVar bid) Units -> Map.Map (AVar bid) Units
evenlyRedistribute bs ps xs = Map.mapWithKey f xs
  where
    f (AllocVar b good)
        | Just bid <- is_marginal_single b good = const (shareFor good * bidQuantityOf bid good)
    f _                                         = id

    -- Test whether this bid label refers to a bid that is marginal on
    -- a single good, and return the bid if so.
    is_marginal_single b good = case Map.lookup b bs of
        Just bid | SinglyMarginal good' <- bidMarginality ps bid
                 , good == good' -> Just bid
        _                        -> Nothing

    -- Calculate the proportion of the bid quantity to allocate to
    -- single marginal bids on a good.  We evenly share the allocation
    -- between all such bids, taking into account the quantity
    -- requested by the bid.
    shareFor good = sum allocated / sum bid_for
      where
        (allocated, bid_for) = unzip [ (x, bidQuantityOf bid good)
                                     | (AllocVar b good', x) <- Map.toList xs
                                     , good == good'
                                     , Just bid <- [is_marginal_single b good]
                                     ]

-- | Specification of a bid without its label.  Two bids are matching
-- iff they have equal specifications.
type BidSpec = (Units, Map.Map Good BidValue)

bidSpec :: Bid bid -> BidSpec
bidSpec bid = (bid ^. bid_quantity, bid ^. bid_values)

-- | Redistribute allocations so that matching bids (those that are
-- identical apart from the label) get equal shares.
--
-- When sharing out the total between the bids, we always round down
-- (unless there is only one bid).  This ensures that if sharing 5
-- units between 3 bids wanting 2 units (accurate to the nearest
-- integer), we get 1,1,1 rather than 2,2,2 (which looks like it
-- allocates a unit we don't have).  It may introduce rounding errors
-- where we report a minor unit less (e.g. 0.49 instead of 0.5).
shareMatching :: forall bid . Ord bid
               => Auction bid
               -> Map.Map bid (Map.Map Good Units) -> Map.Map bid (Map.Map Good Units)
shareMatching auction allocs =
    Map.fromList [ (bid ^. bid_label, x)
                 | bid <- _auction_bids auction
                 , Just x <- [Map.lookup (bidSpec bid) m2]
                 ]
  where
    -- Build a map from bid specifications to the total allocation for
    -- all bids with that specification, and the number of such bids.
    m1 :: Map.Map BidSpec (Map.Map Good Units, Int)
    m1 = Map.fromListWith (\ (x, y) (x', y') -> (Map.unionWith (+) x x', y + y')) $
           map f (_auction_bids auction)

    f bid = (bidSpec bid, (fromMaybe Map.empty (Map.lookup (bid ^. bid_label) allocs), 1))

    -- Divide the totals allocated to each equivalence class by the
    -- number of members to get the per-bid allocation.
    m2 :: Map.Map BidSpec (Map.Map Good Units)
    m2 = Map.map calc_share m1

    calc_share (total_share, num_bids)
      | num_bids > 1 = Map.map (share_between num_bids) total_share
      | otherwise    = total_share

    -- Round the total number of units we have to divide, then divide
    -- them up, then round the resulting quantities downwards.  This
    -- means the shares add up to at most the original (rounded)
    -- total.
    share_between num_bids = floorToScale sf . (/ fromIntegral num_bids) . roundToScale sf

    sf = auction ^. auction_scale

$(makeLenses ''RationingOptions)