-- | Data type representing an instance of the problem.

{-# LANGUAGE RecordWildCards #-}

module AM3.Instance (
  -- * Types
  Instance
  , CenterId
  , OfficeId
  , SegmentId
    -- * Construction
  , newInstance
    -- * Importing
  , fromFile
    -- * Exporting
  , toDat
  , toFile
    -- * Queries
  , isAllowed
  , numCenters
  , numOffices
  , numSegments
  , officeData
  , centerCost
  , segmentCost
  , segmentThreshold
  , centerCosts
  , segmentCosts
  , segmentThresholds
  , centers
  , offices
  , allowedConnections
  , centerCapacity
  , replications
  , centerCapacities
  , officeDatas
  , index2CO
   ) where

import AM3.DatParser
import Control.Arrow
import Data.Foldable
import Data.Ix
import Data.List.Split (chunksOf)
import Data.Vector.Unboxed (Vector, fromList, (!))
import qualified Data.Vector.Unboxed as V
import GRASP
import Text.Megaparsec hiding (fromFile)

-- | Identifier of a Center.
type CenterId = Int

-- | Identifier of an office.
type OfficeId = Int

-- | Identifier of a segment.
type SegmentId = Int

-- | Container for all parameters of the problem.
data Instance = Instance {
  _nC :: Int -- ^ Number of centers.
  , _nO :: Int -- ^ Number of offices.
  , _nP :: Int -- ^ Number of segments
  , _r :: Int -- ^ Number of replications.
  , _d :: Vector Int -- ^ Needed storage per office.
  , _k :: Vector Int -- ^ Center capacities.
  , _f :: Vector Cost -- ^ Fixed cost of each center.
  , _s :: Vector Cost -- ^ Segment costs.
  , _m :: Vector Int -- ^ Segments thresholds.
  , _u :: Vector Bool -- ^ Allowed connections. Stored by rows (row per center).
  , _us :: [(CenterId, OfficeId)] -- ^ List of allowed connections.
  } deriving (Show, Read)

newInstance ::
  Int -- ^ Number of replications.
  -> [Int] -- ^ Needed storage per office.
  -> [Int] -- ^ Center capacities.
  -> [Cost] -- ^ Fixed costs.
  -> [Cost] -- ^ Segment costs.
  -> [Int] -- ^ Segment thresholds.
  -> [Bool] -- ^ Allowed connections.
  -> Instance
newInstance rr dd kk ff ss mm uu =
  let x =
        Instance {
        _r = rr
        , _nC = length kk
        , _nO = length dd
        , _nP = length mm
        , _d = fromList dd
        , _k = fromList kk
        , _f = fromList ff
        , _s = fromList ss
        , _m = fromList mm
        , _u = fromList uu
        , _us = calcAllowedConnections x
        }
  in x

allowedConnections :: Instance -> [(CenterId, OfficeId)]
allowedConnections = _us

emptyInstance :: Instance
emptyInstance = newInstance 0 [] [] [] [] [] []

replications :: Instance -> Int
replications Instance{..} = _r

-- | List of all 'CenterId'.
centers :: Instance -> [CenterId]
centers Instance{..} = range (0, _nC - 1)

-- | List of all 'OfficeId'.
offices :: Instance -> [OfficeId]
offices Instance{..} = range (0, _nO - 1)


-- | True if the connection is allowed.
isAllowed :: Instance -> CenterId -> OfficeId -> Bool
isAllowed Instance{..} c o = _u ! ix
  where ix = _nO*c + o

-- | Unidimensional to bidimensional index.
index2CO :: Instance -> Int -> (CenterId, OfficeId)
index2CO Instance{..} = (`divMod`_nO)

-- | Cost for using a center.
centerCost :: Instance -> CenterId -> Cost
centerCost Instance{..} = (_f !)

-- | Vector of center costs.
centerCosts :: Instance -> Vector Cost
centerCosts = _f

centerCapacity :: Instance -> CenterId -> Int
centerCapacity Instance{..} = (_k !)

centerCapacities :: Instance -> Vector Int
centerCapacities = _k

officeDatas :: Instance -> Vector Int
officeDatas = _d

-- | Needed storage of an office.
officeData :: Instance -> OfficeId -> Int
officeData Instance{..} = (_d !)

-- | Number of offices.
numOffices :: Instance -> Int
numOffices Instance{..} = _nO

-- | Number of centers.
numCenters :: Instance -> Int
numCenters Instance{..} = _nC

-- | Number of segments
numSegments :: Instance -> Int
numSegments Instance{..} = _nP

-- | Cost of a segment.
segmentCost :: Instance -> SegmentId -> Int
segmentCost Instance{..} = (_s !)

segmentCosts :: Instance -> Vector Int
segmentCosts = _s

-- | Threshold of a segment.
segmentThreshold :: Instance -> SegmentId -> Int
segmentThreshold Instance{..} = (_m !)

segmentThresholds :: Instance -> Vector Int
segmentThresholds = _m

-- | List of all allowed connections.
calcAllowedConnections :: Instance -> [(CenterId, OfficeId)]
calcAllowedConnections i = [ (c, o) | c <- centers i, o <- offices i
                                    , isAllowed i c o ]

mkInstance :: [Assig] -> Instance
mkInstance as = let parsed = foldl' f emptyInstance as
                in parsed {_us = calcAllowedConnections parsed}
  where
    f i (Assig "r" (Num x)) = i {_r = x}
    f i (Assig "NO" (Num x)) = i {_nO = x}
    f i (Assig "NC" (Num x)) = i {_nC = x}
    f i (Assig "NP" (Num x)) = i {_nP = x}
    f i (Assig "d" xs) = i {_d = fromList $ flatten xs}
    f i (Assig "k" xs) = i {_k = fromList $ flatten xs}
    f i (Assig "f" xs) = i {_f = fromList $ flatten xs}
    f i (Assig "s" xs) = i {_s = fromList $ flatten xs}
    f i (Assig "u" xs) = i {_u = fromList $ map toEnum $ flatten xs}
    f i (Assig "m" xs) = i {_m = fromList $ flatten xs}
    f i _ = i


-- | Creates and 'Instance' from a @.dat@ file.
fromFile :: FilePath -> IO (Either ParseError Instance)
fromFile path = right mkInstance <$> parseFromFile pdat path

-- | Shows an 'Instance' in @.dat@ format.
toDat :: Instance -> String
toDat Instance{..} = unlines [
  assigInt "r" _r
  , assigInt "NO" _nO
  , assigInt "NP" _nP
  , assigInt "NC" _nC
  , assigVector "d" _d
  , assigVector "k" _k
  , assigVector "f" _f
  , assigVector "m" _m
  , assigVector "s" _s
  , assigVector2 "u" (V.map fromEnum _u) _nO
  ]
  where
    showAssig name val = name ++ " = " ++ val ++ ";"
    assigInt name x = showAssig name (show x)
    assigVector name v = showAssig name (noCommas $ show v)
    assigVector2 name v2 numCols =
      showAssig name (noCommas $ show (chunksOf numCols (V.toList v2)))
    noCommas = map (\case ',' -> ' '; x -> x)

-- | Exports an 'Instance' to a @.dat@ file.
toFile :: FilePath -> Instance -> IO ()
toFile p = writeFile p . toDat