-- | This module implements a data type for constructing a new 'Solution'.

{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE TemplateHaskell       #-}
module AM3.Solution (
  -- * Types
  Solution
  , Candidate
    -- * Construction
  , empty
    -- * Queries
  , cost
  , connections
  , correct
    -- * Other Queries
    , spaceLeft
    , dataLeft
    , connection
    -- * Modifiers
  , updateCon
  , insertCon
    -- * GRASP parameters
  , gparams
    -- * GRASP modifiers
  , appendCand
    -- * GRASP queries
  , neighborhood
  , candidates
  ) where

import           AM3.Instance
import           Control.Arrow
import           Control.Lens
import           Control.Monad.Random
import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as H
import           Data.HashSet (HashSet)
import qualified Data.HashSet as S
import           Data.Hashable
import qualified Data.List as L
import           Data.Maybe
import           Data.Vector.Unboxed (findIndex, unsafeFreeze, (!))
import           Debug.Trace
import           GRASP
import           System.Random.Shuffle

-- | Represents a solution to an 'Instance' of the problem.
data Solution = Solution {
  _inst     :: Instance -- ^ Instance.
  , _distr  :: HashMap (CenterId, OfficeId) Int -- ^ amount of data per connection.
  , _conns :: HashMap OfficeId (HashSet CenterId) -- ^ list of connected centers per office.
  , _oUsage :: HashMap OfficeId Int -- ^ data stored per office.
  , _cUsage :: HashMap CenterId Int -- ^ storage usage per center.
  } deriving (Show, Read)
makeLenses ''Solution

-- | A candidate that can be added to a 'Solution'.
type Candidate = Connection

-- | Connection between center and office of a certain amount.
type Connection = (CenterId, OfficeId, Int)


centerTotalCost :: Instance
         -> CenterId -- ^ Id of the center.
         -> Int -- ^ Amount of data.
         -> Cost -- ^ total cost.
centerTotalCost _ _ 0 = 0
centerTotalCost i c k = centerCost i c + case findIndex (<=k) (segmentThresholds i) of
  Nothing -> error "no suitable segment exists"
  Just six -> k * segmentCost i six

-- | Tells if an office has some data stored in a center.
existsCon :: Solution -> CenterId -> OfficeId -> Bool
existsCon s c o = fromMaybe False (S.member o <$> H.lookup c (s ^. conns))

-- | List of 'Connection's in a 'Solution'.
connections :: Solution -> [Connection]
connections s = [ (c, o, k) | ((c, o), k) <- H.toList (s ^. distr)]

-- | Returns how much data from some office is assigned to a center.
connection :: Solution -> CenterId -> OfficeId -> Int
connection s c o = fromMaybe 0 (H.lookup (c, o) (s ^. distr))

-- | Strict sum
sum' :: [Cost] -> Cost
sum' = L.foldl' (+) 0

-- | Computes the total cost of a solution.
cost :: Solution -> Cost
cost sol = (sum' . map (uncurry $ centerTotalCost i)) $ H.toList (sol ^. cUsage)
  where
    i = sol ^. inst

-- | Creates an empty 'Solution'.
empty :: Instance -> Solution
empty insta = Solution
    { _inst = insta
    , _distr = H.empty
    , _cUsage = H.empty
    , _oUsage = H.empty
    , _conns = H.fromList [ (o, S.empty) | o <- offices insta ]
    }

-- | Returns (old, new, updatedMap)
update :: (Eq k, Hashable k) =>
          k -- ^ key.
          -> (Maybe a -> Maybe a) -- ^ update function.
          -> HashMap k a -- ^ The hashmap.
          -> (Maybe a, Maybe a, HashMap k a) -- ^ (old, new, updatedMap).
update k f hm =
      case f old of
          Nothing -> (old, Nothing, H.delete k hm)
          Just new -> (old, Just new, H.insert k new hm)
  where
    old = H.lookup k hm

addUsage :: CenterId -> OfficeId -> Int -> Solution -> Solution
addUsage c o k = over oUsage (H.insertWith (+) o k) . over cUsage (H.insertWith (+) c k)

assertConnExists :: CenterId -> OfficeId -> Bool -> Solution -> Solution
assertConnExists c o exists = over conns (H.adjust (fun c) o)
  where fun | exists = S.insert
            | otherwise = S.delete

-- | Returns True if the solution satisfies all the restrictions.
--
-- At the moment only checks if every office has stored its data.
correct :: Solution -> Bool
correct s = all ((== 0) . dataLeft s) (offices (s ^. inst))

-- | Updates a connection if present.
updateCon :: CenterId -> OfficeId -> (Maybe Int -> Int) -> Solution -> Solution
updateCon c o upd s =
  let (mold, mnew, hm') = update (c, o) f (s ^. distr)
      diff = new - old
      old = fromMaybe 0 mold
      new = fromMaybe 0 mnew
      s1
        | diff == 0 = s
        | otherwise = addUsage c o diff s
      s2
        | old == 0 && new > 0 = assertConnExists c o True s1
        | old > 0 && new == 0 = assertConnExists c o False s1
        | otherwise = s1
  in set distr hm' s2
  where
    f x
      | updx == 0 = Nothing
      | otherwise = Just updx
      where updx = upd x

-- | Inserts a new connection.
insertCon :: Solution -> CenterId -> OfficeId -> Int -> Solution
insertCon s c o v = assertConnExists c o True $ addUsage c o v $ over distr (H.insert (c, o) v) s

-- | Appends a 'Candidate' to the 'Solution'. Assumes that the center and the
-- office are not already connected.
appendCand :: Solution -> Candidate -> Solution
appendCand s (c, o, v) = insertCon s c o v

-- | Space left in a center.
spaceLeft :: Solution -> CenterId -> Int
spaceLeft s c = centerCapacity (s ^. inst) c - centerUsage s c

-- | Data left in an office. Replications are taken into account.
dataLeft :: Solution -> OfficeId -> Int
dataLeft s o = replications (s ^. inst) * officeData (s ^. inst) o - officeUsage s o

-- | Generates in a random order the list of neighbors.
--
-- Neighbors are generated as follows:
--
-- * Move as many data as possible from an office in one center to another center.
--
-- * Interchange two offices from two centers.
-- @(c1, o1, d1) (c2, o2, d2) --> (c1, o2, d2) (c2, o1, d1)@
neighborhood :: MonadRandom m => Solution -> m [Solution]
neighborhood s = do
  rcon <- shuffleM (connections s)
  return [ moveData c1 c2 o s | (c1, o, _) <- rcon, c2 <- centers (s ^. inst)]
    where
      -- | Moves as many data as possible from o in c1 to c2.
      moveData :: CenterId -> CenterId -> OfficeId -> Solution -> Solution
      moveData c1 c2 o s = (updateCon c1 o (\x -> fromMaybe 0 x - much)
                            >>> updateCon c2 o (\x -> fromMaybe 0 x + much)) s
        where
          much = min (spaceLeft s c2) (connection s c1 o)
      swapData c1 o1 c2 o2 = undefined

-- | How much data is stored in a center.
centerUsage :: Solution -> CenterId -> Int
centerUsage s c = fromMaybe 0 (H.lookup c (s ^. cUsage))

-- | How much data from an office is already stored.
officeUsage :: Solution -> OfficeId -> Int
officeUsage s o = fromMaybe 0 (H.lookup o (s ^. oUsage))

-- | Generates the list of candidates and their estimated cost.
--
-- Candidates are generated as follows:
--
-- * All data from an office is stored in a center.
-- * All remaining space from a center is occupied with some data from an office.
candidates :: Solution -> [(Candidate, Cost)]
candidates sol = mapMaybe (uncurry mkCandidate) (allowedConnections i)
  where
    i = view inst sol
    mkCandidate c o
      | add == 0 = Nothing
      | add > 0 && not (existsCon sol c o) =
        let addCost = centerTotalCost i c (cx + add) - centerTotalCost i c cx
        in Just ((c, o, add), addCost)
      | otherwise = Nothing
      where
        cx = centerUsage sol c
        add = minimum [officeData i o, spaceLeft sol c, dataLeft sol o]

-- | Creates the parameters 'GParams' for 'grasp' specifics for this problem.
gparams :: MonadRandom m =>
  Int -- ^ Maximum number of iterations.
  -> Maybe Int -- ^ Optional maximum number of candidates.
  -> Double -- ^ Alpha.
  -> Instance -- ^ Problem instance.
  -> GParams Solution Candidate m
gparams maxi maxCand alph inst = GParams {
  alpha = alph
  , maxitr = maxi
  , costf = return . cost
  , correctf = return . correct
  , start = return (empty inst)
  , append = \s -> return . appendCand s
  , genCandidates = \s -> let c = candidates s in return $ maybe id take maxCand c
  , neighbors = neighborhood
}