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

-- | This module is responsible for calculating data for supply and demand
-- curves of different types.
module ProductMixAuction.LP.Demand
  ( -- * Representation of demand curves
    DemandCurveType(..)
  , DemandCurve(..)
  , defaultDemandCurves

    -- * Calculating demand curves
  , calcDemand
  , normalisedTotalDemandPoints

    -- * Utilities
  , stepsToPoints
  , bidSurplus

    -- * Lenses
  , dc_type
  , dc_supply
  , dc_demand
  , dc_max_price
  ) where

import Control.Lens
import Data.Aeson
import Data.Aeson.TH
import Data.Function (on)
import Data.List
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set
import Data.Ord
import GHC.Generics
import qualified Test.QuickCheck as QC

import ProductMixAuction.LP.Bid
import ProductMixAuction.LP.Core
import ProductMixAuction.LP.TQSS
import ProductMixAuction.Supply
import ProductMixAuction.Types

-- | The demand curve type determines how to calculate supply and
-- demand.  See 'calcDemand' for more details.  Where applicable, the
-- 'Int' field gives the number of points to use when calculating
-- demand by repeated runs of the LP.
data DemandCurveType
      = VariableTotalAllocation Int
          -- ^ Base goods keep their original supply curve, and for
          -- other goods, we fix the supply based on the allocations
          -- resulting from the original run of the auction.
      | FixedTotalAllocation Int
          -- ^ Supply curves of all goods are fixed based on the
          -- original allocations.
      | AdjustedDemand Int
          -- ^ Only the supply curve for the current good is fixed,
          -- and it uses a large negative price to avoid quantities
          -- other than the target quantity.
      | UnadjustedDemand
          -- ^ Aggregate demand of all the bids on this good, ignoring
          -- bidding on other goods (so paired bids are counted for
          -- multiple goods).  Only really makes sense for horizontal
          -- auctions.
      | UnadjustedDemandWithTQSS
          -- ^ Same demand curve values as 'UnadjustedDemand', but the
          -- supply curve includes the TQSS as if it were additional
          -- supply.
      | NoDemand
          -- ^ Supply curve only, without any demand curve points.
  deriving (Eq, Read, Show, Generic)

-- | Data for plotting a supply and demand curve graph.
data DemandCurve =
    DemandCurve
        { _dc_type      :: DemandCurveType
          -- ^ Approach used to generate these data.
        , _dc_supply    :: [(Units, Price)]
          -- ^ Points representing the supply curve step function.
        , _dc_demand    :: [(Units, Price)]
          -- ^ Points on the demand curve, which gives possible total
          -- allocations to the good and all its successors, along
          -- with the resulting price differences compared to the
          -- previous good.
        , _dc_max_price :: Price
          -- ^ Maximum value for the y-axis.
        }
  deriving (Eq, Generic, Show)

$(makeLenses ''DemandCurve)

instance FromJSON DemandCurve
  where parseJSON = genericParseJSON $ jsonOptions "_dc_"

instance ToJSON DemandCurve
  where toJSON = genericToJSON $ jsonOptions "_dc_"

-- TODO: We have to use TH rather than Generics here, because the
-- Generic-derived instances don't round-trip.  We should investigate
-- why this is and file an aeson bug.
$(deriveJSON (jsonOptions "") ''DemandCurveType)

{-
instance FromJSON DemandCurveType where
  parseJSON = genericParseJSON $ jsonOptions ""

instance ToJSON DemandCurveType where
  toJSON = genericToJSON $ jsonOptions ""
-}

instance QC.Arbitrary DemandCurve where
  arbitrary = DemandCurve <$> QC.arbitrary <*> QC.arbitrary <*> QC.arbitrary <*> QC.arbitrary

instance QC.Arbitrary DemandCurveType where
  arbitrary = QC.oneof [ VariableTotalAllocation <$> QC.arbitrary
                       , FixedTotalAllocation <$> QC.arbitrary
                       , AdjustedDemand <$> QC.arbitrary
                       , pure UnadjustedDemand
                       , pure UnadjustedDemandWithTQSS
                       , pure NoDemand
                       ]


-- | Determine which graphs to generate for this supply by default.
-- For a horizontal supply we generate 'UnadjustedDemand' (and
-- 'UnadjustedDemandWithTQSS' if there is a constraint-based TQSS),
-- which do not require the number of points to be specified.  For a
-- vertical supply we generate 'AdjustedDemand' with a small number of
-- points.  Do not generate demand curves for more complex supply
-- orderings.
defaultDemandCurves :: Supply u -> Maybe (TQSS fun) -> [DemandCurveType]
defaultDemandCurves supply mb_tqss
  | isHorizontalSupply supply
  , Just tqss <- mb_tqss
  , isConstrainedTQSS tqss     = [UnadjustedDemand, UnadjustedDemandWithTQSS]
  | isHorizontalSupply supply  = [UnadjustedDemand]
  | isVerticalSupply supply    = [AdjustedDemand defaultDemandPoints]
  | otherwise                  = [NoDemand]
  where
    -- Default number of demand points to plot on supply/demand graphs.
    defaultDemandPoints :: Int
    defaultDemandPoints = 20


-- | Calculate a list of supply and demand curves for each good.
--
-- The number of units (horizontal axis) gives the amount allocated to
-- the good and all successive goods.  The price (vertical axis) gives
-- the demand price of the good relative to the demand price of the
-- previous good.
--
-- For 'UnadjustedDemand' and 'UnadjustedDemandWithTQSS', the demand
-- of all the bids on this good is aggregated, ignoring bidding on
-- other goods (so paired bids are counted for multiple goods).  This
-- only really makes sense for horizontal auctions.  The latter adds
-- the TQSS vertically to the supply curve.
--
-- For other demand curve types, this uses 'calcDemandCurveLP' to run
-- the auction repeatedly over a range of quantities.  The supply of
-- the current good is fixed to be the desired quantity.  The supply
-- of other goods is adjusted depending upon the 'DemandCurveType':
--
--  * For 'VariableTotalAllocation', the base goods keep their
--    original supply curve, and for other goods, we fix the supply
--    based on the allocations resulting from the original run of the
--    auction.  Fixed supply curves have a single step of height zero.
--
--  * For 'FixedTotalAllocation', the supply curves of all goods are
--    fixed based on the allocations.  Fixed supply curves have a
--    single step of height zero.
--
--  * For 'AdjustedDemand', all goods other than the current good keep
--    their original supply curves.  The fixed supply curve for the
--    current good has a single step with a large negative height.
calcDemand :: Goodly bid
           => Auction bid -> Maybe (TQSS TQSSTable) -> Map.Map Good Units -> [DemandCurveType]
           -> IO (Map.Map Good [DemandCurve])
calcDemand auction mb_tqss good_allocs curve_types
    = Map.fromList <$> mapM demandFor (auctionGoodLabels auction)
  where
    supply = _auction_supply auction

    -- For this good, calculate each relevant graph
    demandFor good = (,) good <$> mapM (calcGraph good) curve_types

    -- For this good and graph, calculate the price for each quantity point
    calcGraph good t = mkDemandCurve good t <$> case t of
      UnadjustedDemand               -> return $ unadjustedDemandCurvePoints (auction ^. auction_bids) good
      UnadjustedDemandWithTQSS       -> return $ unadjustedDemandCurvePoints (auction ^. auction_bids) good
      VariableTotalAllocation points -> calc points good 0              (flip isFirstGood supply)
      FixedTotalAllocation    points -> calc points good 0              (const False)
      AdjustedDemand          points -> calc points good (-2^(20::Int)) (const True)
      NoDemand                       -> return []

    calc = calcDemandCurveLP auction good_allocs

    -- Given the demand curve points, calculate the supply curve
    -- points (adding in the TQSS for 'UnadjustedDemandWithTQSS')
    mkDemandCurve good t demand_points = DemandCurve { _dc_type      = t
                                                     , _dc_supply    = supply_points
                                                     , _dc_demand    = demand_points
                                                     , _dc_max_price = max_price
                                                     }
      where
        sc0 = lookupSupplyCurve good supply
        sc = case (t, mb_tqss) of
               (UnadjustedDemandWithTQSS, Just tqss) -> addSupplyCurveSteps (tqss^.tqss_fun) sc0
               _                                     -> sc0
        supply_points = supplyCurvePoints max_price sc

        -- Add 10% to the maximum price of the supply and demand curves,
        -- for lines that extend vertically "to infinity".
        max_price    = round (extra * fromIntegral (maximum ( 0 : maybeToList (sc ^. sc_final_price)
                                              ++ (sc ^.. sc_steps . each . step_price)
                                              ++ map snd demand_points)))
        extra = 1.1 :: Double


-- | Add a finite list of steps vertically to an existing supply
-- curve.
addSupplyCurveSteps :: [Step Units] -> SupplyCurve Units -> SupplyCurve Units
addSupplyCurveSteps ss (SupplyCurve ss' mb_p) = SupplyCurve (addSteps ss ss') mb_p

-- | Add two lists of steps vertically, splitting them up as needed.
-- The result will be the same width as the narrower of the two inputs.
addSteps :: [Step Units] -> [Step Units] -> [Step Units]
addSteps [] _  = []
addSteps _  [] = []
addSteps (x:xs) (y:ys) = case compare (x^.step_units) (y^.step_units) of
    LT -> Step (x^.step_units) (x^.step_price + y^.step_price)
            : addSteps xs ((y & step_units -~ (x^.step_units)) : ys)
    EQ -> Step (x^.step_units) (x^.step_price + y^.step_price)
            : addSteps xs ys
    GT -> Step (y^.step_units) (x^.step_price + y^.step_price)
            : addSteps ((x & step_units -~ (y^.step_units)) : xs) ys

-- | Subtract the second list of steps from the first, dropping any
-- steps that would have a negative price.
subtractSteps :: [Step Units] -> [Step Units] -> [Step Units]
subtractSteps xs ys =
    filter ((>= 0) . _step_price) (addSteps xs (ys & each.step_price %~ negate))

-- | Calculate the demand curves for each good, at a fixed number of
-- points given as an argument.
--
-- The number of units (horizontal axis) gives the amount allocated to
-- the good and all successive goods.  The price (vertical axis) gives
-- the demand price of the good relative to the demand price of the
-- previous good.
--
-- This works by running the auction over a range of quantities,
-- fixing the supply curves appropriately and finding the prices that
-- result.  We fix the supply of the good itself to be the desired
-- quantity (using a fixed supply curve of the appropriate width, and
-- where the height is given as an argument).  The supply of other
-- goods is adjusted only for goods on which the predicate is false.
calcDemandCurveLP :: Goodly bid
                  => Auction bid
                  -> Map.Map Good Units
                  -> Int
                  -> Good
                  -> Price
                  -> (Good -> Bool)
                  -> IO [(Units, Price)]
calcDemandCurveLP auction good_allocs points good supply_price skip =
    mapM priceFor quantityPoints
  where
    supply = _auction_supply auction

    -- Points at which to calculate demand
    quantityPoints = basicPoints ++ extraPoints

    -- Evenly distribute a fixed number of points over the width of
    -- the supply curve
    basicPoints
      | step_size > 0 = [step_size,2*step_size..max_quantity]
      | otherwise     = []
    max_quantity = totalSupplyCurve (lookupSupplyCurve good supply)
    step_size    = max_quantity / fromIntegral points

    -- If this good (and its successors) have a non-zero allocation,
    -- make sure we plot the demand at the allocation and a tiny step
    -- below and above it, for a clearer graph when the demand price
    -- drops close to the allocation
    extraPoints
      | Just m <- Map.lookup good good_allocs'
      , m > 0     = [m, m - epsilon, m + epsilon]
      | otherwise = []
    epsilon = doubleToUnits (scaleDistance (_auction_scale auction))

    -- Calculate the demand price for this good and graph at this
    -- point
    priceFor u = do ar <- runAuctionCore "calcDemand" auction { _auction_supply = supply' }
                    return (u, max 0 (relativePrice supply (ar ^. ar_prices) good))
      where
        supply' = Map.foldrWithKey f supply (Map.insert good u good_allocs')

        f g u'
          | skip g, g /= good = id
          | otherwise         = setSupplyCurve g (mkFiniteSupplyCurve [Step u' supply_price])

    -- Calculate the total allocation to each good and all its successors
    good_allocs' = Map.mapWithKey (\ g _ -> allocAbove supply good_allocs g) good_allocs

-- | Calculate the intervals for the unadjusted demand curve for this
-- good.  This is a decreasing function, where prices are the bid
-- prices and quantities arise from aggregating the quantities
-- demanded by each bid, ignoring any bidding on other goods.
unadjustedDemandCurvePoints :: [Bid bid] -> Good -> [(Units, Price)]
unadjustedDemandCurvePoints bids good = go 0 xs0
  where
    xs0 = sortBy (comparing (Down . fst))
        . ((0, 0) :)
        . filter ((> 0) . fst)
        . map (\ bid -> (bidValue bid good, bidDemandQuantity bid good))
        $ bids

    -- TODO: factor out a function to do this aggregation
    go _ []         = []
    go u ((p,w):xs) = (u,p') : (u', p') : go u' xs
      where
        u' = u + w
        p' = round p


-- | Convert a supply curve into a list of points.
supplyCurvePoints :: Price -> SupplyCurve Units -> [(Units, Price)]
supplyCurvePoints max_price sc = stepsToPoints (fromMaybe max_price (sc ^. sc_final_price))  (sc ^. sc_steps)

-- | Convert a list of steps into a list of points.  This includes two
-- points for each step (giving the left and right ends for each given
-- price), plus an extra point indicating the line going off to
-- infinity.
stepsToPoints :: Price -> [Step Units] -> [(Units, Price)]
stepsToPoints max_price = go 0
  where
    go u []           = [(u, max_price)]
    go u (step:steps) = (u, p) : (u', p) : go u' steps
      where
        u' = u + step ^. step_units
        p  = step ^. step_price


-- | Calculate the normalised total demand as a decreasing step
-- function from quantities to prices, represented as a list of
-- points.
--
-- This works essentially as follows:
--
--  * Assign each bid's quantity demanded to one or more goods, giving
--    an absolute demand curve for each good.
--
--  * Normalise each absolute demand curve by subtracting the supply
--    curve for the corresponding good.
--
--  * Add the normalised demand curves horizontally.
--
-- For the first step, we must convert a bid into the demand on one or
-- more goods.  If the bid is definitely fully allocated (its surplus
-- is strictly positive), its demand corresponds to what it is
-- actually allocated.  If the bid has a zero surplus, but is
-- allocated something, its demand includes what it is actually
-- allocated and uses an arbitrary allocated good for the remainder.
-- Otherwise, if the bid is not allocated anything, use an arbitrary
-- good that maximises its (possibly negative) surplus.
--
-- Note that for bids whose surplus is not strictly positive this is
-- an approximation: making different arbitrary choices would change
-- the demand graph (but only for quantities beyond the total
-- allocated).
normalisedTotalDemandPoints :: forall bid . Ord bid
                            => Supply Units
                            -> [Bid bid]
                            -> Map.Map Good Price
                            -> Map.Map bid (Map.Map Good Units)
                            -> [(Units, Price)]
normalisedTotalDemandPoints supply bids ps allocs =
    stepsToPoints 0 $ addHoriz $ map normalise $ fstOut $ concatMap toDemandSteps bids
  where
    -- Calculate the good(s) demanded by the bid, and the quantity and
    -- absolute price for each one.
    toDemandSteps :: Bid bid -> [(Good, Step Units)]
    toDemandSteps bid =
        case Map.lookup (bid ^. bid_label) allocs of
          Just xs
            -- Bid has a strictly positive surplus: treat it as
            -- demanding what it actually received.
            | m > 0                    -> allocationToDemand bid xs
            -- Bid received an allocation (so will have zero surplus):
            -- treat it as demanding what it actually received, plus a
            -- remainder step on an arbitrary good that it received.
            | (g,_):_ <- Map.toList xs -> allocationToDemand bid xs
                                       ++ [ mkDemandStep bid g r | let r = bidDemandQuantity bid g - sum xs, r > 0 ]
            -- Bid did not receive an allocation (zero or negative
            -- surplus): treat it as demanding an arbitrary good that
            -- is closest to being allocated.
          _ | g:_ <- Set.toList gs     -> [mkDemandStep bid g (bidDemandQuantity bid g)]
          _                            -> []
      where
        (m, gs, _, _) = bidSurplus ps bid

    allocationToDemand :: Bid bid -> Map.Map Good Units -> [(Good, Step Units)]
    allocationToDemand bid xs = map (uncurry (mkDemandStep bid)) (Map.toList xs)

    mkDemandStep :: Bid bid -> Good -> Units -> (Good, Step Units)
    mkDemandStep bid g u = (g, Step u (round (bidValue bid g)))

    -- Sort the demand steps to produce a decreasing absolute demand
    -- curve, then subtract the supply curve to give normalised
    -- prices.
    normalise :: (Good, [Step Units]) -> [Step Units]
    normalise (g, ss) = sortBy (comparing (Down . _step_price)) ss `subtractSteps` (sc ^. sc_steps)
      where
        sc = lookupSupplyCurve g supply

    -- Merge the normalised demand curves for each good horizontally,
    -- by concatenating and sorting again (we could probably do this
    -- more efficiently, but this will do).
    addHoriz :: [[Step Units]] -> [Step Units]
    addHoriz = sortBy (comparing (Down . _step_price)) . concat

fstOut :: Ord a => [(a,b)] -> [(a,[b])]
fstOut = map f . groupBy ((==) `on` fst) . sortBy (comparing fst)
  where
    f xs = (fst (head xs), map snd xs)

-- | Given the auction prices and a bid, calculate:
--
-- * the maximal surplus @m@ achieved by the bid (i.e. the value to
--   the bid of its preferred goods at the given prices);
--
-- * the set of goods @gs@ that achieve surplus @m@;
--
-- * the maximal surplus @m'@ achieved on goods that may be marginal
--   for generalised bids (i.e. goods with a surplus of zero or a
--   non-unique positive value);
--
-- * the set of goods @gs'@ that achieve surplus @m'@.
bidSurplus :: Map.Map Good Price -> Bid bid -> (Double, Set.Set Good, Double, Set.Set Good)
bidSurplus ps bid
  | Map.null xs = (0, Set.empty, 0, Set.empty)
  | Map.null ys = (m, gs, 0, Set.empty)
  | otherwise   = (m, gs, m', gs')
  where
    -- Construct a map from surplus values to sets of goods.
    xs = Map.fromListWith Set.union $ map (\ (good, bv) -> (w good bv, Set.singleton good))
                                          (Map.toList (bid ^. bid_values))

    -- Find the maximal surplus, m, and the goods that achieve it.
    (m, gs) = Map.findMax xs

    -- Find all goods that may be marginal for generalised bids, and
    -- the maximal surplus achieved by such goods.
    ys  = Map.filterWithKey (\ k s -> k == 0 || (k > 0 && Set.size s > 1)) xs
    m'  = fst (Map.findMax ys)
    gs' = Set.unions (Map.elems ys)

    -- Calculate the value of the good to the bid at the given prices.
    w good bv = _TweakedPrice (bv ^. bv_price - maybe 0 fromIntegral (Map.lookup good ps))
              / _TradeOff (bv ^. bv_fraction)