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

-- | This is the main module for solving standard Product-Mix Auctions
-- using linear programming, possibly involving multiple runs of the
-- LP solver for some features.  See 'ProductMixAuction.LP.Core' for
-- translating and running basic auctions using a single LP solver
-- iteration.
module ProductMixAuction.LP
  ( -- * Inputs
    AuctionInput(..)
  , Maximise(..)
  , Shuffle(..)
  , mkAuctionInput
  , flattenBidders

    -- * Outputs
  , AuctionOutput(..)

    -- * Running auctions
  , runAuction

    -- * Lenses
  , ac_supply
  , ac_bidders
  , ac_scale
  , ac_tqss
  , ac_demand_curves
  , ac_ration
  , ac_preferences
  , ac_verbosity
  , ac_maximise

  , ao_prices
  , ao_bid_allocations
  , ao_bidder_allocations
  , ao_good_allocations
  , ao_tqss_total_supply
  , ao_normalised_price
  , ao_tqss_points
  , ao_demand_curves
  , ao_profit
  ) where

import Control.Exception
import Control.Lens hiding ((.=))
import Control.Monad
import Data.Aeson.Types as A
import Data.Default.Class
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Text as T
import qualified Data.Vector as V
import GHC.Generics
import qualified Test.QuickCheck as QC

import qualified ProductMixAuction.BudgetConstraints as BC
import qualified ProductMixAuction.BudgetConstraints.Types as BC
import ProductMixAuction.LP.AdditionalConstraint
import ProductMixAuction.LP.Bid
import ProductMixAuction.LP.Core
import ProductMixAuction.LP.Demand
import ProductMixAuction.LP.Rationing
import ProductMixAuction.LP.TQSS
import ProductMixAuction.Supply
import ProductMixAuction.Types


data Maximise = MaxEfficiency  -- ^ A standard LP auction, in which efficiency is maximised
              | MaxProfit      -- ^ An auction in which the auctioneer's profit is maximised, not solved using an LP
  deriving (Bounded, Enum, Eq, Generic, Show)

instance Default Maximise where
  def = MaxEfficiency

instance FromJSON Maximise where
  parseJSON = genericParseJSON $ jsonOptions ""

instance ToJSON Maximise where
  toJSON = genericToJSON $ jsonOptions ""

instance QC.Arbitrary Maximise where
  arbitrary = QC.arbitraryBoundedEnum


data Shuffle = NoShuffle    -- ^ Keep order of bid labels as in input
             | ShuffleBids  -- ^ Randomize the order of bid labels (across all bidders)
  deriving (Bounded, Enum, Eq, Generic, Show)

instance Default Shuffle where
  def = NoShuffle

instance FromJSON Shuffle where
  parseJSON = genericParseJSON $ jsonOptions ""

instance ToJSON Shuffle where
  toJSON = genericToJSON $ jsonOptions ""


-- | Inputs for a Product-Mix Auction to be solved using linear
-- programming.  This supports advanced features that may require
-- multiple iterations of the LP solver, unlike 'Auction', which
-- represents the data provided for a single iteration.
data AuctionInput bid =
    AuctionInput
        { _ac_supply  :: Supply Units
          -- ^ Supply of goods available in the auction
        , _ac_bidders :: [Bidder bid]
          -- ^ Bid data from participants in the auction
        , _ac_scale   :: ScaleFactor
          -- ^ Precision with which results should be reported
        , _ac_tqss    :: Maybe (TQSS TQSSTable)
          -- ^ Optionally use a TQSS to determine total quantity sold in the auction
        , _ac_demand_curves :: [DemandCurveType]
          -- ^ Calculate the given demand curve types
        , _ac_ration  :: RationingOptions
          -- ^ Approach for rationing limited quantities of goods between tied bidders
        , _ac_preferences :: [Good]
          -- ^ Preference order for allocating goods (permutation of
          -- the good labels in descending order of preference)
        , _ac_constraints :: [AdditionalConstraint (BidderName, bid)]
          -- ^ Additional linear constraints on the quantities allocated to bids
        , _ac_verbosity :: Verbosity
          -- ^ Controls how much debug output is generated when
          -- solving the auction
        , _ac_maximise :: Maximise
          -- ^ Whether to maximise efficiency or profit
        , _ac_shuffle :: Shuffle
          -- ^ Whether to shuffle bids on input
        }
  deriving (Generic, Show)

makeLenses ''AuctionInput

instance Default (AuctionInput bid)

instance ( Ord bid, FromJSONKey bid, FromJSON bid
         ) => FromJSON (AuctionInput bid) where
  parseJSON = genericParseJSON $ jsonOptions "_ac_"

instance ( ToJSONKey bid, ToJSON bid
         ) => ToJSON (AuctionInput bid) where
  toJSON = genericToJSON $ jsonOptions "_ac_"

-- | Given a supply of goods and some bidders, create an auction input
-- with sensible defaults.
mkAuctionInput :: Ord bid => Supply Units -> [Bidder bid] -> Maybe (TQSS TQSSTable)
               -> AuctionInput bid
mkAuctionInput supply bidders mb_tqss =
    AuctionInput
        { _ac_supply           = supply
        , _ac_bidders          = bidders
        , _ac_scale            = def
        , _ac_ration           = def
        , _ac_tqss             = mb_tqss
        , _ac_demand_curves    = defaultDemandCurves supply mb_tqss
        , _ac_preferences      = defaultPreferences supply
        , _ac_constraints      = []
        , _ac_verbosity        = def
        , _ac_maximise         = def
        , _ac_shuffle          = def
        }

-- | Given an 'AuctionInput', extract the relevant data for a
-- single-iteration 'Auction'.
inputToAuction :: AuctionInput bid -> Auction (BidderName, bid)
inputToAuction ac = Auction { _auction_supply       = ac ^. ac_supply
                            , _auction_bids         = flattenBidders (ac ^. ac_bidders)
                            , _auction_scale        = ac ^. ac_scale
                            , _auction_preferences  = ac ^. ac_preferences
                            , _auction_constraints  = ac ^. ac_constraints
                            , _auction_size_limit   = Nothing
                            , _auction_tweak_supply = True
                            , _auction_verbosity    = ac ^. ac_verbosity
                            }

-- | Flatten a list of bidders to produce a list of bids annotated with the
-- bidder name.
flattenBidders :: [Bidder bid] -> [Bid (BidderName, bid)]
flattenBidders = concatMap flattenBidder
  where
    flattenBidder :: Bidder bid -> [Bid (BidderName, bid)]
    flattenBidder bidder = (bidder ^. bidder_bids) & each . bid_label %~ (,) (bidder ^. bidder_name)

-- | Convert an 'AuctionInput' to an 'Auction' as in 'inputToAuction', but
-- optionally shuffle the order of bids/constraints and the bid labels.
-- Produces an auction with integer bid labels, and an injective function to map
-- them back to the original labels.
--
-- It is important to shuffle the order of the bids/constraints lists *and* to
-- produce a random permutation of the bid labels, because the solvers may
-- depend on both.  For example, 'maxProfitAuction' allocates goods to bids in
-- the order of the list, while the LP generated by 'maxEfficiencyAuction'
-- depends on the ordering on bid labels.
--
inputToAuctionShuffle :: Ord bid => AuctionInput bid -> IO (Auction Int, Int -> (BidderName, bid))
inputToAuctionShuffle ac = do
    bids <- shuffleList (auction0 ^. auction_bids)
    constraints <- shuffleList (auction0 ^. auction_constraints)
    xys <- zip [1..] <$> shuffleList (map _bid_label bids)
    let intMap     = Map.fromList xys
        labelMap   = Map.fromList (map (\(x,y)->(y,x)) xys)
        intToLabel = (Map.!) intMap
        labelToInt = (Map.!) labelMap
        auction    = mapAuction labelToInt
                                (auction0 { _auction_bids = bids, _auction_constraints = constraints })
    return (auction, intToLabel)
  where
    auction0 = inputToAuction ac
    shuffleList = case ac^.ac_shuffle of
                    ShuffleBids -> QC.generate . QC.shuffle
                    NoShuffle   -> return


-- | Outputs from running a Product-Mix Auction.
data AuctionOutput bid =
    AuctionOutput
        { _ao_prices       :: Map.Map Good Price
          -- ^ Prices for each good as determined by the auction.
          -- Bids above these prices are successful and will be
          -- allocated goods (modulo rationing).
        , _ao_good_allocations :: Map.Map Good Units
          -- ^ For each good, the total quantity allocated.  Due to
          -- rounding errors, the total of the allocation to
          -- individual bids may not match this figure.
        , _ao_bidder_allocations :: Map.Map BidderName (Map.Map Good Units)
          -- ^ For each bidder, the total quantity allocated of each
          -- good.  Due to rounding errors, the total of the
          -- allocation to individual bids may not match this figure.
        , _ao_bid_allocations  :: Map.Map (BidderName, bid) (Map.Map Good Units)
          -- ^ For each bid, the quantities of each good it is
          -- allocated.  Usually this will be a single type of good
          -- with the quantity requested by the bid, but rationing may
          -- result in fulfilling a single bid by allocating multiple
          -- goods.
        , _ao_tqss_total_supply :: Maybe Units
          -- ^ Total supply size resulting from solving the TQSS (or
          -- 'Nothing', if a TQSS is not in use).
        , _ao_normalised_price :: Maybe Price
          -- ^ Normalised price, if a normalised TQSS is in use.
        , _ao_tqss_points  :: Maybe [TQSSPoint]
          -- ^ Points of the TQSS evaluated (or 'Nothing', if a
          -- non-normalised TQSS is not in use).
        , _ao_demand_curves :: Map.Map Good [DemandCurve]
          -- ^ For each good, the demand curves for that good (if
          -- demand curves have been calculated).
        , _ao_profit :: Maybe Price
          -- ^ Auctioneer's profit from the auction, where this is
          -- being maximised.
        }
  deriving (Eq, Show, Generic)

-- | The completely uninformative auction output, used when we are
-- asked to run an auction with no goods or no bids.
emptyOutput :: AuctionOutput bid
emptyOutput = AuctionOutput Map.empty Map.empty Map.empty Map.empty Nothing Nothing Nothing Map.empty Nothing

makeLenses ''AuctionOutput

instance Default (AuctionOutput bid) where
  def = emptyOutput

instance (Ord bid, FromJSONKey bid, FromJSON bid) => FromJSON (AuctionOutput bid) where
  parseJSON = genericParseJSON $ jsonOptions "_ao_"

instance (Ord bid, ToJSONKey bid, ToJSON bid) => ToJSON (AuctionOutput bid) where
  toJSON = genericToJSON $ jsonOptions "_ao_"

instance (QC.Arbitrary bid, Ord bid) => QC.Arbitrary (AuctionOutput bid) where
  arbitrary = AuctionOutput <$> QC.arbitrary <*> QC.arbitrary <*> QC.arbitrary
                            <*> QC.arbitrary <*> QC.arbitrary <*> QC.arbitrary
                            <*> QC.arbitrary <*> QC.arbitrary <*> QC.arbitrary


-- | Run a Product-Mix Auction as specified by the inputs, and produce
-- the corresponding outputs.
runAuction :: Goodly bid => AuctionInput bid -> IO (AuctionOutput bid)
runAuction ac
  | empty     = return emptyOutput
  | otherwise = do
    checkValidAuction ac
    case ac ^. ac_maximise of
      MaxEfficiency -> maxEfficiencyAuction ac
      MaxProfit     -> maxProfitAuction ac
  where
    bids = flattenBidders (ac ^. ac_bidders)
    empty = null bids || null (supplyGoods (_ac_supply ac))

-- | Check that supply curves are all non-decreasing, throwing an
-- exception if this is not the case.
checkValidAuction :: AuctionInput bid -> IO ()
checkValidAuction ac = do
    case decreasingSupply (ac ^. ac_supply) of
        Just good -> throwIO (SupplyCurveDecreasing good)
        Nothing   -> return ()

-- | Solve an auction that maximises efficiency, by running one or
-- more iterations of a linear programme.
maxEfficiencyAuction :: Goodly bid => AuctionInput bid -> IO (AuctionOutput bid)
maxEfficiencyAuction ac = do
    (auction, unshuffle) <- inputToAuctionShuffle ac
    (auction', ar, mb_r, mb_points) <- case ac ^. ac_tqss of
        Nothing   -> do ar <- runAuctionCore "runAuction" auction
                        return (auction, ar, Nothing, Nothing)
        Just tqss -> do (auction', ar, r, mb_points) <- runTQSS auction tqss
                        return (auction', ar, Just r, mb_points)
    alloc <- Map.mapKeys unshuffle <$> runRationingOptions (ac ^. ac_ration) auction' ar
    let good_allocs   = goodAllocations   sf alloc
        bidder_allocs = bidderAllocations sf alloc
        bid_allocs    = bidAllocations    sf alloc
    demand <- calcDemand auction' (ac^.ac_tqss) good_allocs (ac^.ac_demand_curves)
    return AuctionOutput { _ao_prices             = ar ^. ar_prices
                         , _ao_good_allocations   = good_allocs
                         , _ao_bidder_allocations = bidder_allocs
                         , _ao_bid_allocations    = bid_allocs
                         , _ao_tqss_total_supply  = mb_r
                         , _ao_normalised_price   = ar ^. ar_normalised_price
                         , _ao_tqss_points        = mb_points
                         , _ao_demand_curves      = demand
                         , _ao_profit             = Nothing
                         }
  where
    sf      = ac ^. ac_scale


-- | Solve an auction that maximises the auctioneer's profit, by
-- converting it into the format expected by the budget-constrained
-- auction solver.
--
-- There is no support for generalised or asymmetric bids, supply
-- orderings other than horizontal, additional constraints or a TQSS.
-- Rationing options and the preference order for allocating goods
-- will be silently ignored.
maxProfitAuction :: Goodly bid => AuctionInput bid -> IO (AuctionOutput bid)
maxProfitAuction ac = do
    unless (isHorizontalSupply (ac ^. ac_supply)) $ throwIO MaxProfitNotHorizontal
    unless (isNothing (ac ^. ac_tqss))            $ throwIO MaxProfitTQSS
    (auction, unshuffle) <- inputToAuctionShuffle ac
    debugPutStrLn (ac ^. ac_verbosity) "maxProfitAuction"  $ showAuction auction
    when   (any badBid (auction ^. auction_bids)) $ throwIO MaxProfitComplexBid
    case BC.runBCAuction (convertAuction auction) of
        Just bcar -> return $ convertAuctionResult (BC.mapBCAuctionResult unshuffle bcar)
        Nothing   -> throwIO MaxProfitNoSolution
  where
    badBid bid = isGeneralised bid || isAsymmetric bid

convertAuction :: Show bid => Auction bid -> BC.BCAuction bid
convertAuction auction =
    BC.BCAuction { BC._bca_kind          = BC.Standard
                 , BC._bca_filter_prices = BC.AllPrices
                 , BC._bca_supply        = convertSupply supply
                 , BC._bca_bids          = map (convertBid num_goods) (auction ^. auction_bids)
                 }
  where
    supply    = auction ^. auction_supply
    num_goods = length (supplyGoods supply)

convertSupply :: Supply Units -> BC.SupplyCurveVector
convertSupply = V.fromList . map snd . listSupplyCurves

convertBid :: Show bid => BC.Dimension -> Bid bid -> BC.Bid bid
convertBid num_goods bid =
    BC.Bid { BC._bid_label  = bid ^. bid_label
           , BC._bid_name   = T.pack (fromMaybe (show (bid ^. bid_label)) (bid ^. bid_name))
           , BC._bid_budget = BC.Budget (_Units (bid ^. bid_quantity))
           , BC._bid_prices = ps
           }
  where
    ps = V.fromList [ BC.Price (_TweakedPrice (bidValue bid j))
                    | j <- [G 1..G num_goods]
                    ]

convertAuctionResult :: Ord bid => BC.BCAuctionResult (BidderName, bid) -> AuctionOutput bid
convertAuctionResult bcar =
    AuctionOutput { _ao_prices             = toGoodMap (fmap round (bcar ^. BC.bcar_prices))
                  , _ao_good_allocations   = toGoodMap (bcar ^. BC.bcar_quantities)
                  , _ao_bidder_allocations = filterAllocs (Map.mapKeysWith (Map.unionWith (+)) fst bid_allocs)
                  , _ao_bid_allocations    = filterAllocs bid_allocs
                  , _ao_tqss_total_supply  = Nothing
                  , _ao_normalised_price   = Nothing
                  , _ao_tqss_points        = Nothing
                  , _ao_demand_curves      = Map.empty
                  , _ao_profit             = Just (round (bcar ^. BC.bcar_profit))
                 }
  where
    toGoodMap :: V.Vector a -> Map.Map Good a
    toGoodMap = Map.fromList . zip [G 1..] . V.toList

    bid_allocs = fmap toGoodMap (bcar ^. BC.bcar_allocations)