{-
   Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
-}

{-# LANGUAGE PatternGuards, DoAndIfThenElse, ConstraintKinds, ScopedTypeVariables #-}

module Camfort.Specification.Units.Synthesis
  (runSynthesis)
where

import Data.Maybe
import qualified Data.Set as S
import Data.Generics.Uniplate.Operations
import Control.Monad.State.Strict hiding (gets)

import qualified Language.Fortran.AST as F
import qualified Language.Fortran.Analysis as FA
import qualified Language.Fortran.Util.Position as FU

import Camfort.Analysis.Annotations hiding (Unitless)
import Camfort.Specification.Units.Environment
import Camfort.Specification.Units.Monad
import Camfort.Specification.Units.InferenceFrontend (puName, puSrcName)

-- | Insert unit declarations into the ProgramFile as comments.
runSynthesis :: Char -> [(VV, UnitInfo)] -> UnitSolver [(VV, UnitInfo)]
runSynthesis marker vars = do
  -- descendBiM finds the head of lists
  modifyProgramFileM $ descendBiM (synthProgramUnits marker vars) <=< descendBiM (synthBlocks marker vars)
  return vars

-- Should be invoked on the beginning of a list of blocks
synthBlocks :: Char -> [(VV, UnitInfo)] -> [F.Block UA] -> UnitSolver [F.Block UA]
synthBlocks marker vars = fmap reverse . foldM (synthBlock marker vars) []

-- Process an individual block while building up a list of blocks (in
-- reverse order) to ultimately replace the original list of
-- blocks. We're looking for blocks containing declarations, in
-- particular, in order to possibly insert a unit annotation before
-- them.
synthBlock :: Char -> [(VV, UnitInfo)] -> [F.Block UA] -> F.Block UA -> UnitSolver [F.Block UA]
synthBlock marker vars bs b@(F.BlStatement a (FU.SrcSpan lp _) _ (F.StDeclaration _ _ _ _ decls)) = do
  pf    <- usProgramFile `fmap` get
  gvSet <- usGivenVarSet `fmap` get
  newBs <- fmap catMaybes . forM (universeBi decls) $ \ e -> case e of
    e@(F.ExpValue _ _ (F.ValVariable _))
      | vname `S.notMember` gvSet                     -- not a member of the already-given variables
      , Just u <- lookup (vname, sname) vars -> do    -- and a unit has been inferred
        -- Create new annotation which labels this as a refactored node.
        let newA  = a { FA.prevAnnotation = (FA.prevAnnotation a) {
                           prevAnnotation = (prevAnnotation (FA.prevAnnotation a)) {
                               refactored = Just lp } } }
            -- Create a zero-length span for the new comment node.
            newSS = FU.SrcSpan (lp {FU.posColumn = 0}) (lp {FU.posColumn = 0})
            -- Build the text of the comment with the unit annotation.
            txt   = marker:" " ++ showUnitDecl (FA.srcName e, u)
            space = FU.posColumn lp - 1
            (F.ProgramFile mi _) = pf
            newB  = F.BlComment newA newSS . F.Comment $ buildCommentText mi space txt
        return $ Just newB
      where
        vname = FA.varName e
        sname = FA.srcName e
    (_ :: F.Expression UA) -> return Nothing
  return (b:reverse newBs ++ bs)
synthBlock _ _ bs b = return (b:bs)

-- Should be invoked on the beginning of a list of program units
synthProgramUnits :: Char -> [(VV, UnitInfo)] -> [F.ProgramUnit UA] -> UnitSolver [F.ProgramUnit UA]
synthProgramUnits marker vars pus = do
  fmap reverse . foldM (synthProgramUnit marker vars) [] $ pus

-- Process an individual program unit while building up a list of
-- program units (in reverse order) to ultimately replace the original
-- list of program units. We're looking for functions, in particular,
-- in order to possibly insert a unit annotation before them.
synthProgramUnit :: Char -> [(VV, UnitInfo)] -> [F.ProgramUnit UA] -> F.ProgramUnit UA -> UnitSolver [F.ProgramUnit UA]
synthProgramUnit marker vars pus pu@(F.PUFunction a (FU.SrcSpan lp _) _ _ _ _ mret _ _) = do
  pf    <- usProgramFile `fmap` get
  gvSet <- usGivenVarSet `fmap` get
  let (vname, sname) = case mret of Just e  -> (FA.varName e, FA.srcName e)
                                    Nothing -> (puName pu, puSrcName pu)
  case lookup (vname, sname) vars of
    -- if return var has a unit & not a member of the already-given variables
    Just u | vname `S.notMember` gvSet -> do
      let newA  = a { FA.prevAnnotation = (FA.prevAnnotation a) {
                         prevAnnotation = (prevAnnotation (FA.prevAnnotation a)) {
                             refactored = Just lp } } }
      -- Create a zero-length span for the new comment node.
      let newSS = FU.SrcSpan (lp {FU.posColumn = 0}) (lp {FU.posColumn = 0})
      -- Build the text of the comment with the unit annotation.
          txt   = marker:" " ++ showUnitDecl (sname, u)
          space = FU.posColumn lp - 1
          (F.ProgramFile mi _) = pf
          newPU = F.PUComment newA newSS . F.Comment $ buildCommentText mi space txt

      -- recursively descend to find program units inside of current one
      fmap (:newPU:pus) $ descendBiM (synthProgramUnits marker vars) pu

    -- otherwise, nevermind, but still recursively descend to find
    -- program units inside of current one
    _ -> fmap (:pus) $ descendBiM (synthProgramUnits marker vars) pu
synthProgramUnit marker vars pus pu = fmap (:pus) $ descendBiM (synthProgramUnits marker vars) pu

-- Pretty print a unit declaration.
showUnitDecl (sname, u) = "unit(" ++ show u ++ ") :: " ++ sname