{-
   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 ImplicitParams #-}
{-# LANGUAGE DeriveDataTypeable #-}

module Analysis.Loops where

import Data.Data
import Data.List
import Data.Ord

import Language.Fortran
import Language.Fortran.Pretty


import Data.Generics.Uniplate.Operations
import Control.Monad.State.Lazy
import Debug.Trace

import Analysis.LVA
import Analysis.Annotations
import Analysis.Syntax
import Analysis.Types

import Helpers
import Traverse

import Transformation.Syntax -- <- for doing reassociation

import qualified Data.Map.Lazy as Map hiding (map, (\\))

--  when travesing whole program collect all declarations with bounds 
--  collect all constants (#1) 
--  identify all loop 'variables' (#2) 
--   - identify all variables indexed by the loop variables

-- loopBody :: Fortran t -> State (TypeEnvStack t) (Fortran ([String], [String], [String]))
-- loopBody (For _ v@(VarName _ s) e1 e2 e3 body) = 
--     let
--         anno = (
--     in For anno v e1 e2 e3 body
--  
-- newFrame gammas = []:gammas
-- pushVar v t (g:gs) = ((v, t):g):gs
-- popVar (((v,t):g):gs) = (g:gs)
-- popFrame (g:gs) = (g, gs)

-- ap (fmap ((,[""]),[""]))

loopAnalyse :: Program a -> Program Annotation
loopAnalyse p = map ((descendBi arrayIndices) . ix . lvaOnUnit . (transformBi reassociate) . (fmap (const unitAnnotation))) p

analyse' :: Program Annotation -> Program Annotation
analyse' p = map ((descendBi arrayIndices) . ix . lvaOnUnit . (transformBi reassociate))  p


collect :: (Eq a, Ord k) => [(k, a)] -> Map.Map k [a]
collect = collect' Map.empty 
          where collect' as []                         = as
                collect' as ((v, n):es) | Map.member v as = collect' (Map.insert v (nub $ n : ((Map.!) as v)) as) es
                                        | otherwise   = collect' (Map.insert v [n] as) es

arrayIndices :: Block Annotation -> Block Annotation
arrayIndices x = 
    let tenv = typeEnv x
        
        arrIxsF :: Fortran Annotation -> Annotation
        arrIxsF y = let readIxs = [(v, mfmap (const ()) e) | 
                                     (Var _ _ [(VarName _ v, e)]) <- rhsExpr y,
                                     length e > 0,
                                     isArrayTypeP' tenv v]

                        writeIxs = [(v, mfmap (const ()) e) |
                                     (Var _ _ [(VarName _ v, e)]) <- lhsExpr y,
                                     length e > 0,
                                     isArrayTypeP' tenv v]

                    in (tag y) { arrsRead = (collect readIxs), arrsWrite = (collect writeIxs) } 
    in extendBi arrIxsF x               

ix :: ProgUnit Annotation -> ProgUnit Annotation
ix = let ixF :: Fortran Annotation -> Annotation
         ixF f = (tag f) { indices = (nub [v | (For _ _ (VarName _ v) _ _ _ _) <- ((universeBi f)::[Fortran Annotation])])}
     in extendBi ixF