{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Representation of additional constraints imposed on bids.
module ProductMixAuction.LP.AdditionalConstraint
  ( AdditionalConstraint(..)
  , mkAbsoluteQuantityConstraint
  , mkRelativeQuantityConstraint
  , showConstraint
  , mapConstraint

    -- * Lenses
  , ac_c
  , ac_c'
  , ac_coefficients
  ) where

import Control.Lens
import Data.Aeson
import Data.Default.Class
import Data.List
import qualified Data.Map as Map
import GHC.Generics

import ProductMixAuction.Types

-- | Represents an additional linear constraint on the bids.  These
-- are always of the form
--
-- > sum_{i,j} c^i_j x^i_j <= C R + C'
--
-- for some coefficients @c^i_j >= 0@ and constants @C, C' >= 0@.
-- Here @x^i_j@ is the amount of good @j@ allocated to bid @i@, and
-- @R@ is the total size of the auction supply.
data AdditionalConstraint bid =
    AdditionalConstraint
        { _ac_coefficients :: Map.Map bid (Map.Map Good Double)
           -- ^ Coefficient @c^i_j@ for each bid and good on LHS
           -- (absence from the map indicates a coefficient of 0).
        , _ac_c  :: Ratio
           -- ^ Coefficient of size of auction on RHS.
        , _ac_c' :: Double
           -- ^ Constant on RHS.
        }
  deriving (Generic, Show)

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

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

instance Default (AdditionalConstraint bid) where
  def = AdditionalConstraint Map.empty 1 0


-- | Make a constraint on the quantity allocated to the given bids,
-- across the given goods.  If all goods are listed, this is
-- essentially equivalent to reducing the '_bid_quantity' for the bids.
mkAbsoluteQuantityConstraint :: Ord bid => [Good] -> [bid] -> Units -> AdditionalConstraint bid
mkAbsoluteQuantityConstraint goods bids u = AdditionalConstraint coeffs 0 (_Units u)
  where
    m      = Map.fromList (map (\ g -> (g, 1)) goods)
    coeffs = Map.fromList (map (\ b -> (b, m)) bids)

-- | Make a constraint on the quantity allocated to the given bids,
-- across the given goods, as a fraction of the total auction size.
mkRelativeQuantityConstraint :: Ord bid => [Good] -> [bid] -> Ratio -> AdditionalConstraint bid
mkRelativeQuantityConstraint goods bids r = AdditionalConstraint coeffs r 0
  where
    m      = Map.fromList (map (\ g -> (g, 1)) goods)
    coeffs = Map.fromList (map (\ b -> (b, m)) bids)

-- | Render an additional constraint as a human-readable string.
showConstraint :: Show bid => AdditionalConstraint bid -> String
showConstraint ac = summands as ++ " <= " ++ summands [showy (_ac_c ac) "R", showy (_ac_c' ac) "1"]
  where
    as = [ showy x ("Alloc(" ++ show bid ++ "," ++ show good ++ ")")
         | (bid, gxs) <- Map.toList (_ac_coefficients ac)
         , (good, x) <- Map.toList gxs
         ]
    showy x s
      | x == 0    = []
      | s == "1"  = [show x]
      | x == 1    = [s]
      | otherwise = [show x ++ " * " ++ s]
    summands xs = intercalate " + " (concat xs)


$(makeLenses ''AdditionalConstraint)


-- | Change the representation of bid labels in an additional constraint.
-- Requires the codomain to have an 'Ord' instance, hence not a 'Functor'.
mapConstraint :: Ord bid' => (bid -> bid') -> AdditionalConstraint bid -> AdditionalConstraint bid'
mapConstraint f = over ac_coefficients (Map.mapKeys f)