{-# LANGUAGE OverloadedStrings, TemplateHaskell #-}
module ProductMixAuction.Graphics.Scatter
  ( scatterBidsChart
  , stackedBidsChart

  , StackedPoint
  , BidPoint
  , ScatterPoint(..)
  , sp_label
  , sp_size
  , sp_x
  , sp_y
  , sp_z
  , sp_colour
  , sp_symbol
  ) where

import Control.Lens
import Control.Monad (guard)
import Data.Aeson (toJSON)
import Data.Csv (ToField)
import Data.List (sortBy, elemIndex)
import qualified Data.Map as Map
import Data.Maybe
import Data.Monoid hiding (All)
import Data.Ord (comparing, Down(..))
import qualified Data.Set as Set
import qualified Data.Text as T
import Graphics.Plotly.Base hiding (x,y,z,i,j,k)
import qualified Graphics.Plotly.Base as P

import ProductMixAuction.Graphics.Types
import ProductMixAuction.LP.Bid
import ProductMixAuction.LP.Demand
import ProductMixAuction.Types

type StackedPoint = ScatterPoint () Color
type BidPoint     = ScatterPoint Double ()
data ScatterPoint z c = ScatterPoint { _sp_label  :: Maybe T.Text
                                     , _sp_size   :: Double
                                     , _sp_x      :: Double
                                     , _sp_y      :: Double
                                     , _sp_z      :: z
                                     , _sp_colour :: c
                                     , _sp_symbol :: Symbol
                                     }

$(makeLenses ''ScatterPoint)

traceBidCore :: (p -> Maybe T.Text) -- ^ Optional label for each point
             -> (p -> Double)       -- ^ Size for each point
             -> (p -> Double)       -- ^ x coordinate
             -> (p -> Double)       -- ^ y coordinate
             -> Maybe (p -> Double) -- ^ z coordinate, if any
             -> Maybe (p -> Color)  -- ^ specified colour, if any
             -> Maybe (p -> Symbol) -- ^ custom symbol, if any
             -> Env -> Bidder bid -> [p] -> Trace
traceBidCore get_lbl get_size get_x get_y mb_get_z mb_get_col mb_get_symbol env bidder ps =
    scatterNd & P.x .~ Just (map (toJSON . get_x) ps)
              & P.y .~ Just (map (toJSON . get_y) ps)
              & P.z .~ (map toJSON <$> mb_zs)
              & text ?~ map (fromMaybe empty_label . get_lbl) ps
              & mode ?~ prefixMode [Markers]
              & name ?~ _BidderName (bidder ^. bidder_name)
              & marker ?~ (defMarker & size ?~ P.List (map (toJSON . (* dotsize) . get_size) ps)
                                     & markercolor .~ (P.List . map toJSON <$> mb_cs)
                                     & markerline ?~ (defMarkerLine & markerlinewidth ?~ All 2)
                                     & symbol .~ (P.List <$> mb_ss) )
              & hoverinfo ?~ HoverNone
              & textposition ?~ textPos
  where
    mb_zs = map <$> mb_get_z <*> pure ps
    mb_cs = map <$> mb_get_col <*> pure ps
    mb_ss = map <$> mb_get_symbol <*> pure ps

    prefixMode  = if env ^. env_options . go_show_bids_q then (ModeText :) else id
    scatterNd   = if isNothing mb_get_z then scatter else scatter3d

    dotsize = env ^. env_options ^. go_dotsize

    -- TopRight looks slightly better on 3d scatter plots, because it is further away from the dot
    textPos = if isNothing mb_get_z then MiddleRight else TopRight

    -- Plotly's Chart Studio tends to display "undefined" in place of
    -- the empty label (at least for 3D graphs), so we use a single space
    empty_label = " "

unitsLabel :: ScaleFactor -> Units -> T.Text
unitsLabel sf = dropSuffix0 . uiLabel . _Units . roundToScale sf
  where
    dropSuffix0 = \x -> T.stripSuffix ".0" x ?| x

unitsFractionLabel :: ScaleFactor -> Units -> Units -> T.Text
unitsFractionLabel sf x y = unitsLabel sf x <> "/" <> unitsLabel sf y

-- | Plot a scatter chart of the bids, provided the auction has two or
-- three goods.
scatterBidsChart :: (ToField bid, Ord bid) => Env -> [Bidder bid] -> BidderAllocMap bid -> Maybe Plotly
scatterBidsChart env bidders allocs = do
    guard (d == 2 || d == 3)
    pure $ plotly ("pma-scatter-bids-" <> T.pack (show d) <> "d") (bid_trs ++ price_trs)
      & layout . title ?~ desc
      & layout . xaxis ?~ axis' gex
      & layout . yaxis ?~ axis' gey
      & layout . zaxis .~ (axis' <$> mb_gez)
  where
    gex    = env ^? env_x_good ?! "Missing good x"
    gey    = env ^? env_y_good ?! "Missing good y"
    mb_gez = env ^? env_z_good

    d = length (env ^. env_goods)
    axis' ge = defAxis & axistitle   ?~ ("Price on good " <> uiLabel good)
                      & range       ?~ (min_bound, max_bound)
                      & axisvisible ?~ True
      where
        -- Calculate the axis bounds by finding the minimum and
        -- maximum prices bid for the relevant good, and adding an
        -- offset of 5% of the range at each end.
        good = ge ^. ge_good
        min_price = fromIntegral (ge ^. ge_min_price)
        max_price = fromIntegral (ge ^. ge_max_price)
        offset = 0.05 * (max_price - min_price)
        min_bound = min_price - offset
        max_bound = max_price + offset
    desc = "Bids in price space"
    bid_trs   = map (traceBidScatter env allocs gex gey mb_gez) bidders
    price_trs = case mb_gez of
                  Just gez -> tracePricesScatter3D gex gey gez
                  Nothing  -> tracePricesScatter2D gex gey

traceBidScatter :: (ToField bid, Ord bid)
                => Env -> BidderAllocMap bid -> GoodEnv -> GoodEnv -> Maybe GoodEnv -> Bidder bid -> Trace
traceBidScatter env allocs gex gey mb_gez bidder =
    traceBidCore _sp_label _sp_size _sp_x _sp_y (_sp_z <$ mb_gez) Nothing (Just _sp_symbol) env bidder ps
  where
    ps = scatterBidData (env ^. env_scale) allocs gex gey mb_gez bidder

scatterBidData :: (ToField bid, Ord bid)
               => ScaleFactor -> BidderAllocMap bid -> GoodEnv -> GoodEnv -> Maybe GoodEnv -> Bidder bid -> [BidPoint]
scatterBidData sf allocs gex gey mb_gez bidder =
    map fst $ sortBy (comparing (Down . snd))
                     [ (sp, a)
                     | bid <- bidder ^. bidder_bids
                     , let (q, a) = overallAllocation sf allocs bidder bid
                     , sp <- adjustPoint a (mkScatterPoint bid q a)
                     ]
  where
    mkScatterPoint bid q a =
        ScatterPoint { _sp_label  = if a == None then Nothing else Just (unitsLabel sf q)
                     , _sp_x      = look bid gex
                     , _sp_y      = look bid gey
                     , _sp_z      = maybe 0 (look bid) mb_gez
                     , _sp_size   = 1
                     , _sp_colour = ()
                     , _sp_symbol = Circle
                     }

    look bid ge = _TweakedPrice (bidValue bid (_ge_good ge))

tracePricesScatter2D :: GoodEnv -> GoodEnv -> [Trace]
tracePricesScatter2D gex gey =
  [ line2d [(xp,  0), (xp, yp), (xp+dp, yp+dp)] & hoverinfo ?~ HoverNone
                                                & name      ?~ "LIP 1"
  , line2d [(0,  yp), (xp, yp)] & hoverinfo ?~ HoverNone
                                & name      ?~ "LIP 2"
  , line2d [(xp, yp), (xp, yp)] & hoverinfo ?~ HoverPlus [HoverX, HoverY]
                                & name      ?~ "Auction prices"
  ]
  where
    xp  = fromIntegral (gex ^. ge_price)
    mxp = 1.05 * fromIntegral (gex ^. ge_max_price)
    yp  = fromIntegral (gey ^. ge_price)
    myp = 1.05 * fromIntegral (gey ^. ge_max_price)
    dp  = min (mxp-xp) (myp-yp)

tracePricesScatter3D :: GoodEnv -> GoodEnv -> GoodEnv -> [Trace]
tracePricesScatter3D gex gey gez =
  [ line3d [(xp,  0, zp), (xp, yp, zp)] & hoverinfo ?~ HoverNone
  , line3d [(0,  yp, zp), (xp, yp, zp)] & hoverinfo ?~ HoverNone
  , line3d [(xp, yp,  0), (xp, yp, zp)] & hoverinfo ?~ HoverNone
  , line3d [(xp, yp, zp), (xp+dp, yp+dp, zp+dp)] & hoverinfo ?~ HoverNone
  , line3d [(xp, yp, zp), (xp, yp, zp)] & hoverinfo ?~ HoverPlus [HoverX, HoverY, HoverZ]
  , quadrilateral3d ((xp,  0, zp), (xp, yp, zp), (xp+dp, yp+dp, zp+dp), (xp+dp, 0,  zp+dp)) (ColRGB 255 0 0)
  , quadrilateral3d ((0,  yp, zp), (xp, yp, zp), (xp+dp, yp+dp, zp+dp), (0, yp+dp,  zp+dp)) (ColRGB 0 255 0)
  , quadrilateral3d ((xp, yp,  0), (xp, yp, zp), (xp+dp, yp+dp, zp+dp), (xp+dp, yp+dp,  0)) (ColRGB 0 0 255)
  , quadrilateral3d ((xp,  0, zp), (xp, yp, zp), (xp, yp,  0), (xp,  0,  0)) (ColRGB 255 0 255)
  , quadrilateral3d ((xp,  0, zp), (xp, yp, zp), ( 0, yp, zp), ( 0,  0, zp)) (ColRGB 255 255 0)
  , quadrilateral3d ((xp, yp,  0), (xp, yp, zp), ( 0, yp, zp), ( 0, yp,  0)) (ColRGB 0 255 255)
  ]
  where
    xp  = fromIntegral (gex ^. ge_price)
    mxp = 1.05 * fromIntegral (gex ^. ge_max_price)
    yp  = fromIntegral (gey ^. ge_price)
    myp = 1.05 * fromIntegral (gey ^. ge_max_price)
    zp  = fromIntegral (gez ^. ge_price)
    mzp = 1.05 * fromIntegral (gez ^. ge_max_price)
    dp = maximum [mxp-xp, myp-yp, mzp-zp]


-- | Plot a stacked column chart of the bids.
stackedBidsChart :: (Ord bid, ToField bid)
                 => Env -> [Bidder bid] -> Map.Map (BidderName, bid) (Map.Map Good Units) -> Plotly
stackedBidsChart env bidders allocs =
  plotly "pma-stacked-bids" trs
    & layout . title ?~ "Allocation of bids to goods"
    & layout . xaxis ?~ (axis "Good"  & tickvals ?~ (toJSON <$> [1 .. length goods])
                                      & ticktext ?~ ((("Good " <>) . uiLabel) <$> goods))
    & layout . yaxis ?~ axis "Price"
    & layout . showlegend ?~ (not $ env ^. env_options. go_simple_colors)
  where
    goods = env ^.. env_goods . traverse . ge_good
    trs = map (traceBidStacked env allocs pm) bidders ++ tracePricesStacked env
    pm = goodPricePreferredMap (env ^. env_scale) (env ^. env_prices) allocs bidders

-- | How much of an allocation (perhaps of a partial good) did a bid
-- receive, relative to the quantity it requested?  This takes account
-- of asymmetric and generalised bids.
data Allocation = None    -- ^ No units allocated
                | Partial -- ^ Bid received less than its demanded quantity
                | Full    -- ^ Bid received the full amount it demanded
  deriving (Eq, Ord)

instance Monoid Allocation where
  mempty  = Full
  mappend = min

adjustPoint :: Allocation -> ScatterPoint z c -> [ScatterPoint z c]
-- Not allocated anything: draw an open circle
adjustPoint None    sp = [ sp & sp_symbol .~ CustomSymbol "circle-open" ]
-- Partially allocated: draw an open circle containing a half-size filled circle
adjustPoint Partial sp = [ sp & sp_symbol .~ CustomSymbol "circle-open"
                         , sp & sp_label  .~ Nothing
                              & sp_size   .~ 13/20
                         ]
-- Fully allocated: draw an open circle on top of the default filled circle
adjustPoint Full    sp = [ sp & sp_symbol .~ CustomSymbol "circle-open"
                         , sp
                         ]

-- | Determine the goods preferred by this bid (i.e. the goods for which
-- the bid price minus the auction price is maximised), given the
-- auction prices and allocations.  If the bid receives an allocation,
-- return all the allocated goods.  If the bid receives nothing,
-- choose the lowest-numbered good that achieves the maximal surplus.
--
-- For each good, also return:
--
--  * the quantity allocated of that good (if there are any
--    allocations), otherwise the quantity requested;
--
--  * whether this represents a full 'Allocation'.
preferredGoods :: Ord bid => ScaleFactor -> Map.Map Good Price -> BidderAllocMap bid -> Bidder bid -> Bid bid -> [(Good, Units, Allocation)]
preferredGoods sf ps alloc bidder bid
  | Just m <- lookupBidAlloc bidder bid alloc
  , not (Map.null m) = map (\ (g, a) -> (g, a, partial_or_full g a)) (Map.toList m)
  | otherwise = case Set.toList gs of
                  g:_ -> [(g, bidDemandQuantity bid g, None)]
                  []  -> []
  where
    (_, gs, _, _) = bidSurplus ps bid

    partial_or_full g a = if roundToScale sf a < roundToScale sf (bidDemandQuantity bid g) then Partial else Full

-- | Return the total quantity allocated to this bid, and whether this
-- represents a full 'Allocation'.
--
-- A bid is considered "fully allocated" if the overall bid quantity
-- constraint is binding (i.e. the sum of the per-good quantities
-- multiplied by the 'bidFraction's meets the 'bid_quantity') or if at
-- least one good's maximum constraint is binding (i.e. the quantity
-- of the good meets the corresponding 'bidDemandQuantity').
--
-- Note that this is tested after rounding, so in some circumstances
-- we may show a bid as fully allocated that actually receives less
-- than its demand in an exact analytic solution.
overallAllocation :: Ord bid => ScaleFactor -> BidderAllocMap bid -> Bidder bid -> Bid bid -> (Units, Allocation)
overallAllocation sf alloc_map bidder bid = (s, alloc)
  where
    -- Map from goods to allocated quantities for this bid, and the total allocated.
    m = fromMaybe Map.empty (lookupBidAlloc bidder bid alloc_map)
    s = sum m

    alloc | overall_full                     = Full
          | any full_for_good (Map.toList m) = Full
          | s > 0                            = Partial
          | otherwise                        = None

    -- Multiply quantity of each good by the corresponding trade-off,
    -- and compare to the overall bid quantity.
    overall_full = roundToScale sf adjusted_total >= roundToScale sf (bid ^. bid_quantity)
    adjusted_total = sum [ Units (_TradeOff (bidFraction bid g)) * u | (g, u) <- Map.toList m ]

    -- Compare good quantity to the demand (adjusted for the
    -- trade-off). Since we are comparing rounded quantities, we
    -- shouldn't multiply through by the trade-off, otherwise we end
    -- up considering an allocation of 3.3 to a bid demanding 10/3 to
    -- be partial rather than full.
    full_for_good (g, u) = roundToScale sf u >= roundToScale sf (bidDemandQuantity bid g)


-- | Construct a map from (good, price) pairs to the list of bid
-- identities for bids whose preferred goods include the given good at
-- the given price.
goodPricePreferredMap :: Ord bid
                      => ScaleFactor
                      -> Map.Map Good Price
                      -> BidderAllocMap bid
                      -> [Bidder bid]
                      -> Map.Map (Good, Price) [(BidderName, bid)]
goodPricePreferredMap sf ps allocs bidders =
    Map.fromListWith (++) [ ((g, round (bidValue bid g)), [bidIdentity bidder bid])
                          | bidder <- bidders
                          , bid <- bidder ^. bidder_bids
                          , (g, _, _) <- preferredGoods sf ps allocs bidder bid
                          ]

bidIdentity :: Bidder bid -> Bid bid -> (BidderName, bid)
bidIdentity bidder bid = (bidder ^. bidder_name, bid ^. bid_label)

type BidderAllocMap bid = Map.Map (BidderName, bid) (Map.Map Good Units)

lookupBidAlloc :: Ord bid => Bidder bid -> Bid bid -> BidderAllocMap bid -> Maybe (Map.Map Good Units)
lookupBidAlloc bidder bid = Map.lookup (bidIdentity bidder bid)


traceBidStacked :: (ToField bid, Ord bid)
                => Env -> BidderAllocMap bid
                -> Map.Map (Good, Price) [(BidderName, bid)]
                -> Bidder bid -> Trace
traceBidStacked env allocs pm bidder =
    traceBidCore _sp_label _sp_size _sp_x _sp_y Nothing mb_sp_colour (Just _sp_symbol) env bidder ps
  where
    -- Plot highest bids first, so that legend reflects fill state of highest bid
    ps = sortBy (comparing (Down . _sp_y)) $
             stackedBidderData (env ^. env_scale) (env ^. env_prices) allocs pm bidder

    -- Specify the colour of each bid only if using simple colours
    mb_sp_colour
      | env ^. env_options . go_simple_colors = Just _sp_colour
      | otherwise                             = Nothing

stackedBidderData ::  (ToField bid, Ord bid)
                => ScaleFactor
                -> Map.Map Good Price -> BidderAllocMap bid
                -> Map.Map (Good, Price) [(BidderName, bid)]
                -> Bidder bid -> [StackedPoint]
stackedBidderData sf ps allocs pm bidder =
    [ sp
    | bid       <- bidder ^. bidder_bids
    , (g, q, a) <- preferredGoods sf ps allocs bidder bid
    , sp        <- adjustPoint a (mkStackedPoint bid g q a)
    ]
  where
    mkStackedPoint bid g q a =
        ScatterPoint { _sp_label  = Just (label bid g q a)
                     , _sp_x      = fromIntegral (_good_index g) + nudge bid g
                     , _sp_y      = _TweakedPrice (bidValue bid g)
                     , _sp_z      = ()
                     , _sp_size   = 1
                     , _sp_colour = color_for bid g
                     , _sp_symbol = Circle
                     }

    -- If the bid is allocated nothing, show the quantity demanded.
    -- If it is partially allocated, show the quantity allocated as a
    -- fraction of the quantity demanded.
    -- If it is fully allocated, show the quantity allocated.
    label bid g q a = case a of
                        None    -> unitsLabel sf demand_q
                        Partial -> unitsFractionLabel sf q demand_q
                        Full    -> unitsLabel sf q
      where
        demand_q = bidDemandQuantity bid g

    -- If there are multiple bids preferring this good/price, offset
    -- each bid evenly across the width of the bar
    nudge bid g
      | Just bns <- Map.lookup (g, round (bidValue bid g)) pm
      , let l = length bns
      , l > 1
      , Just i <- elemIndex (bidder ^. bidder_name, bid ^. bid_label) bns
      = nudge_offset l i
      | otherwise = 0

    -- Given the number of bids "l" for the same good/price, and the
    -- index "i" of this bid within the list, calculate how much to
    -- offset this point from the center.
    nudge_offset l i = fromIntegral i * (w / fromIntegral (l-1)) - w/2
      where
        -- Reduce the spread width when there is only one good
        w | num_goods > 1 = thresholdwidth
          | otherwise     = thresholdwidth/2

        num_goods = Set.size (Map.keysSet ps)

    color_for bid g = case compare (round (bidValue bid g)) (fromMaybe 0 (Map.lookup g ps)) of
                               GT -> lightgreen
                               EQ -> lightbrown
                               LT -> lightred


tracePricesStacked :: Env -> [Trace]
tracePricesStacked env = map tracePriceGood (Map.toList (env ^. env_prices))
  where
    tracePriceGood (g, p0) =
      line2d [(i - thresholdwidth, p), (i, p), (i + thresholdwidth, p)]
        & hoverinfo ?~ HoverPlus [HoverY]
        & hoveron   ?~ [HoverPoints]
        & fill      ?~ ToSelf
        & name      ?~ ("Price of good " <> T.pack (show (_good_index g)))
      where
        p = fromIntegral p0
        i = fromIntegral (_good_index g)

thresholdwidth :: Double
thresholdwidth = 0.2