content: {- |
  Module      :  Camfort.Specification.Units.ModFile
  Description :  Helpers for working with units-relevant ModFiles.
  Copyright   :  (c) 2017, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish
  License     :  Apache-2.0
  
  Maintainer  :  dom.orchard@gmail.com
  Stability   :  experimental
  -}
  
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE DeriveGeneric      #-}
  {-# LANGUAGE TupleSections #-}
  
  module Camfort.Specification.Units.ModFile
    (
      genUnitsModFile
    , initializeModFiles
    , runCompileUnits
    , dumpModFileCompiledUnits
    ) where
  
  import           Camfort.Analysis (analysisModFiles)
  import           Camfort.Specification.Units.Annotation (UA)
  import           Camfort.Specification.Units.Environment (Constraint(..), foldUnits, UnitInfo(..), colSort, Constraints)
  import           Camfort.Specification.Units.InferenceBackend (flattenConstraints, flattenUnits, genUnitAssignments')
  import           Camfort.Specification.Units.Monad
  import           Control.Monad.State (get, gets, lift)
  import           Data.Binary (Binary, decodeOrFail, encode)
  import           Data.Data (Data)
  import           Data.Generics.Uniplate.Operations (universeBi)
  import           Data.List (partition)
  import qualified Data.Map as M
  import           Data.Maybe (mapMaybe, catMaybes)
  import qualified Data.Set as S
  import           Data.Typeable (Typeable)
  import           GHC.Generics (Generic)
  import qualified Language.Fortran.AST as F
  import qualified Language.Fortran.Analysis as FA
  import           Language.Fortran.Util.ModFile
  import           Prelude hiding (mod)
  
  -- | The data-structure stored in 'fortran-src mod files'
  data CompiledUnits = CompiledUnits
    { cuTemplateMap  :: TemplateMap
    , cuNameParamMap :: NameParamMap
    } deriving (Ord, Eq, Show, Data, Typeable, Generic)
  
  instance Binary CompiledUnits
  
  emptyCompiledUnits :: CompiledUnits
  emptyCompiledUnits = CompiledUnits M.empty M.empty
  
  combinedCompiledUnits :: ModFiles -> CompiledUnits
  combinedCompiledUnits mfs = CompiledUnits { cuTemplateMap = M.unions tmaps
                                            , cuNameParamMap = M.unions nmaps }
    where
      cus = map mfCompiledUnits mfs
      tmaps = map cuTemplateMap cus
      nmaps = map cuNameParamMap cus
  
  -- | Name of the labeled data within a ModFile containing unit-specific info.
  unitsCompiledDataLabel :: String
  unitsCompiledDataLabel = "units-compiled-data"
  
  mfCompiledUnits :: ModFile -> CompiledUnits
  mfCompiledUnits mf = case lookupModFileData unitsCompiledDataLabel mf of
    Nothing -> emptyCompiledUnits
    Just bs -> case decodeOrFail bs of
      Left _ -> emptyCompiledUnits
      Right (_, _, cu) -> cu
  
  -- | Initialize units-relevant ModFile information.
  initializeModFiles :: UnitSolver ()
  initializeModFiles = do
    mfs <- lift . lift $ analysisModFiles
    let compiledUnits = combinedCompiledUnits mfs
    modifyTemplateMap  . const . cuTemplateMap  $ compiledUnits
    modifyNameParamMap . const . cuNameParamMap $ compiledUnits
  
  -- | Produce information for a "units-mod" file.
  runCompileUnits :: UnitSolver CompiledUnits
  runCompileUnits = do
    cons <- usConstraints `fmap` get
    pf <- usProgramFile  `fmap` get
  
    let unitAssigns = flattenConstraints cons
    let epsilon = 0.001 -- arbitrary
    let approxEq a b = abs (b - a) < epsilon
  
    let variables = M.fromList [ (NPKVariable var, units) | ([UnitPow (UnitVar var) k], units) <- unitAssigns
                                                          , k `approxEq` 1 ]
  
    -- Create sets of relevant program-unit and variable names.
    let getName pu | F.Named n <- FA.puName pu = Just n | otherwise = Nothing
    let puNameSet = S.fromList $ catMaybes [ getName pu | pu <- universeBi pf :: [F.ProgramUnit UA] ] ++
                                 [ FA.varName e | F.BlInterface _ _ (Just e) _ _ _ <- universeBi pf :: [F.Block UA] ]
  
    let puVarNameSet pu = S.fromList $
          [ (FA.varName v, FA.srcName v) | F.Declarator _ _ v _ _ _ <- universeBi pu :: [F.Declarator UA] ]
  
    -- Map of modules -> associated declared variables
    let puVarNameMap :: M.Map F.ProgramUnitName (S.Set VV)
        puVarNameMap = M.fromListWith S.union [ (F.getName pu, puVarNameSet pu)
                                              | pu@F.PUModule {} <- universeBi pf :: [F.ProgramUnit UA] ]
  
    -- Filter functions that cut out any information not having to do
    -- with the current modules being compiled.
    let filterPUs = M.filterWithKey (const . (`S.member` puNameSet))
    -- FIXME: investigate whether we should include vars that are linked
    -- transitively to current pf vars should also be included.
    tmap <- (filterPUs . M.map optimiseTemplate) `fmap` gets usTemplateMap
  
    let findNPK vv = ( (NPKVariable vv), ) <$> M.lookup (NPKVariable vv) variables
  
    -- 'Name Param Map': module names -> (variables -> unit info)
    let npm = M.filter (not . M.null) $ flip M.map puVarNameMap (M.fromList . mapMaybe findNPK . S.toList)
  
    pure CompiledUnits { cuTemplateMap = tmap, cuNameParamMap = npm }
  
  -- | Cut out unnecessary constraints in the template using the solver.
  optimiseTemplate :: Constraints -> Constraints
  optimiseTemplate cons = map (\ (l, r) -> ConEq (foldUnits l) r) optimised
    where
      unitAssigns  = genUnitAssignments' (compileColSort) cons
      unitPows     = map (fmap flattenUnits) unitAssigns
      optimised    = filter cull $ map (fmap foldUnits . shiftTerms) unitPows
  
      cull (lhs, _) = or [ True | (UnitPow (UnitParamPosAbs _) _) <- universeBi lhs ]
  
      isUnitRHS (UnitPow (UnitName _) _)        = True
      isUnitRHS (UnitPow (UnitParamEAPAbs _) _) = True
      isUnitRHS (UnitPow (UnitParamImpAbs _) _) = True
      isUnitRHS (UnitPow (UnitParamPosAbs _) _) = False
      isUnitRHS _                               = False
  
      negateCons = map (\ (UnitPow u k) -> UnitPow u (-k))
  
      shiftTerms :: ([UnitInfo], [UnitInfo]) -> ([UnitInfo], [UnitInfo])
      shiftTerms (lhs, rhs) = (lhsOk ++ negateCons rhsShift, rhsOk ++ negateCons lhsShift)
        where
          (lhsOk, lhsShift) = partition (not . isUnitRHS) lhs
          (rhsOk, rhsShift) = partition isUnitRHS rhs
  
      compileColSort = flip colSort
  
  -- | Generate a new ModFile containing Units information.
  genUnitsModFile :: F.ProgramFile UA -> CompiledUnits -> ModFile
  genUnitsModFile pf cu = alterModFileData f unitsCompiledDataLabel $ genModFile pf
    where
      f _ = Just $ encode cu
  
  -- | Return information about compiled units for debugging purposes
  dumpModFileCompiledUnits :: ModFile -> Maybe String
  dumpModFileCompiledUnits mf = do
    bs <- lookupModFileData unitsCompiledDataLabel mf
    cu <- case decodeOrFail bs of
      Left _ -> Nothing
      Right (_, _, cu) -> Just cu
    pure . unlines $ [ "Template Map (size=" ++ show (M.size (cuTemplateMap cu)) ++ "):"
                     , concat [ unlines (i 2 fname':map (i 4 . show) temp)
                              | (fname, temp) <- M.toList (cuTemplateMap cu)
                              , let fname' = "Template for " ++ show fname ]
                     , "NameParam Map  (size=" ++ show (M.size (cuNameParamMap cu)) ++ "):"
                     , concat [ unlines (i 2 mod':[ i 4 (show ui ++ " :: " ++ show v)
                                                  | (NPKVariable (v,_), ui) <- M.toList npkmap ])
                              | (mod, npkmap) <- M.toList (cuNameParamMap cu)
                              , let mod' = "Module " ++ show mod ]
                     ]
   where
     i n s = replicate n ' ' ++ s
  