{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric      #-}

module Camfort.Specification.Units.MonadTypes where

import           Control.Monad.Reader
import           Control.Monad.State.Strict
import           Data.Binary                             (Binary)
import           Data.Char                               (toLower)
import           Data.Data                               (Data)
import qualified Data.IntMap.Strict                      as IM
import           Data.List                               (find, isPrefixOf)
import qualified Data.Map.Strict                         as M
import qualified Data.Set                                as S
import           Data.Typeable                           (Typeable)
import           GHC.Generics                            (Generic)
import qualified Language.Fortran.AST                    as F
import           Language.Fortran.Util.ModFile           (ModFiles)

import           Camfort.Analysis
import           Camfort.Analysis.Annotations            (Annotation)
import           Camfort.Input
import           Camfort.Specification.Units.Annotation  (UA)
import           Camfort.Specification.Units.Environment (Constraints, PP,
                                                          UnitInfo, VV)

--------------------------------------------------------------------------------
--  Environment
--------------------------------------------------------------------------------

-- | Some options about how to handle literals.
data LiteralsOpt
  = LitPoly     -- ^ All literals are polymorphic.
  | LitUnitless -- ^ All literals are unitless.
  | LitMixed    -- ^ The literal "0" or "0.0" is fully parametric
                -- polymorphic. All other literals are monomorphic,
                -- possibly unitless.
  deriving (Show, Eq, Ord, Data)

instance Read LiteralsOpt where
  readsPrec _ s = case find ((`isPrefixOf` map toLower s) . fst) ms of
                    Just (str, con) -> [(con, drop (length str) s)]
                    Nothing         -> []
    where
      ms = [ ("poly", LitPoly), ("unitless", LitUnitless), ("mixed", LitMixed)
           , ("litpoly", LitPoly), ("litunitless", LitUnitless), ("litmixed", LitMixed) ]

-- | Options for the unit solver
data UnitOpts = UnitOpts
  { uoLiterals :: LiteralsOpt               -- ^ how to handle literals
  }
  deriving (Show, Data, Eq, Ord)

data UnitEnv = UnitEnv
  { unitOpts     :: UnitOpts
  , unitProgramFile :: F.ProgramFile Annotation
  }

--------------------------------------------------------------------------------
--  State
--------------------------------------------------------------------------------

-- | Function/subroutine name -> associated, parametric polymorphic constraints
type TemplateMap = M.Map F.Name Constraints

-- | Things that can be exported from modules
data NameParamKey
  = NPKParam PP Int     -- ^ Function/subroutine name, position of parameter
  | NPKVariable VV      -- ^ variable
  deriving (Ord, Eq, Show, Data, Typeable, Generic)

instance Binary NameParamKey

-- | mapped to a list of units (to be multiplied together)
type NameParamMap = M.Map F.ProgramUnitName (M.Map NameParamKey [UnitInfo])

-- | Variable => unit
type VarUnitMap   = M.Map VV UnitInfo
-- | Set of variables given explicit unit annotations
type GivenVarSet  = S.Set F.Name
-- | Alias name => definition
type UnitAliasMap = M.Map String UnitInfo
-- | Map of CallId to CallId
type CallIdMap    = IM.IntMap Int

-- | Working state for the monad
data UnitState = UnitState
  { usProgramFile  :: F.ProgramFile UA
  , usVarUnitMap   :: VarUnitMap
  , usGivenVarSet  :: GivenVarSet
  , usUnitAliasMap :: UnitAliasMap
  , usTemplateMap  :: TemplateMap
  , usNameParamMap :: NameParamMap
  , usCallIdRemap  :: CallIdMap
    -- | Next number to returned by 'freshId'.
  , usNextUnique   :: Int
  , usConstraints  :: Constraints }
  deriving (Show, Data)

--------------------------------------------------------------------------------
--  Monads
--------------------------------------------------------------------------------

-- | Analysis with access to 'UnitEnv' information.
type UnitAnalysis = ReaderT UnitEnv (AnalysisT () () IO)

-- | UnitSolvers are analyses annotated with unit information.
type UnitSolver = StateT UnitState UnitAnalysis