{-# LANGUAGE UndecidableInstances #-} -- required due to instance design

module Language.Fortran.Common.Array where

import Control.DeepSeq ( NFData )
import GHC.Generics ( Generic )
import Data.Data ( Data, Typeable )
import Data.Binary ( Binary )
import Text.PrettyPrint.GenericPretty ( Out(..) )

import qualified Text.PrettyPrint as Pretty

import qualified Language.Fortran.PrettyPrint as F

-- | A single array dimension with bounds of type @a@.
--
--   * @'Num' a => 'Dim' a@ is a static, known-size dimension.
--   * @'Dim' ('Language.Fortran.AST.Expression' '()')@ is a dimension with
--     unevaluated bounds expressions. Note that these bounds may be constant
--     expressions, or refer to dummy variables, or be invalid.
--   * @'Num' a => 'Dim' ('Maybe' a)@ is a dimension where some bounds are
--     known, and others are not. This may be useful to record some information
--     about dynamic explicit-shape arrays.
data Dim a = Dim
  { dimLower :: a -- ^ Dimension lower bound.
  , dimUpper :: a -- ^ Dimension upper bound.
  } deriving stock (Show, Generic, Data, Eq)
    deriving stock (Functor, Foldable, Traversable)
    deriving anyclass (NFData, Binary)

    -- | This instance is purely for convenience. No definition of ordering is
    --   provided, and the implementation may change at any time.
    deriving stock Ord

-- | Fortran syntax uses @lower:upper@, so only provide an 'Out' instance for
--   that style.
instance Out a => Out (Dim a) where
    doc (Dim lb ub) = doc lb <> Pretty.char ':' <> doc ub

instance Out (Dim a) => F.Pretty (Dim a) where
    pprint' _ = doc

-- | Fortran array dimensions, defined by a list of 'Dim's storing lower and
--   upper bounds.
--
-- You select the list type @t@ (which should be 'Functor', 'Foldable' and
-- 'Traversable') and the bound type @a@ (e.g. 'Int').
--
-- Using a non-empty list type such as 'Data.List.NonEmpty.NonEmpty' will
-- disallow representing zero-dimension arrays, providing extra soundness.
--
-- Note the following excerpt from the F2018 standard (8.5.8.2 Explicit-shape
-- array):
--
-- > If the upper bound is less than the lower bound, the range is empty, the
-- > extent in that dimension is zero, and the array is of zero size.
--
-- Note that the 'Foldable' instance does not provide "dimension-like" access to
-- this type. That is, @'length' (a :: 'Dims' t a)@ will _not_ tell you how many
-- dimensions 'a' represents. Use 'dimsLength' for that.
data Dims t a
  -- | Explicit-shape array. All dimensions are known.
  = DimsExplicitShape
      (t (Dim a)) -- ^ list of all dimensions

  -- | Assumed-size array. The final dimension has no upper bound (it is
  --   obtained from its effective argument). Earlier dimensions may be defined
  --   like explicit-shape arrays.
  | DimsAssumedSize
      (Maybe (t (Dim a))) -- ^ list of all dimensions except last
      a                   -- ^ lower bound of last dimension

  -- | Assumed-shape array. Shape is taken from effective argument. We store the
  --   lower bound for each dimension, and thus also the rank (via list length).
  | DimsAssumedShape
      (t a) -- ^ list of lower bounds

    deriving stock (Generic)
    deriving stock (Functor, Foldable, Traversable)

-- We have to standalone derive most instances due to the @t@ list-like.
deriving stock instance (Show a, Show (t a), Show (t (Dim a)))
  => Show (Dims t a)
deriving anyclass instance (NFData a, NFData (t a), NFData (t (Dim a)))
  => NFData (Dims t a)
deriving stock instance (Data a, Data (t a), Data (t (Dim a)), Typeable t)
  => Data (Dims t a)
deriving stock instance (Eq a, Eq (t a), Eq (t (Dim a)))
  => Eq (Dims t a)
deriving anyclass instance (Binary a, Binary (t a), Binary (t (Dim a)))
  => Binary (Dims t a)

-- | This instance is purely for convenience. No definition of ordering is
--   provided, and the implementation may change at any time.
deriving stock instance (Ord a, Ord (t a), Ord (t (Dim a)))
  => Ord (Dims t a)

instance (Foldable t, Functor t, Out (Dim a), Out a)
  => Out (Dims t a) where
    docPrec _ = doc
    doc = Pretty.parens . \case
      DimsExplicitShape ds ->
        prettyIntersperse dimSep $ fmap doc ds
      DimsAssumedShape ss ->
        prettyIntersperse dimSep $ fmap go ss
        where
          go s = doc s <> Pretty.char ':'
      DimsAssumedSize mds d ->
        -- A bit fragile, but hopefully won't explode on empty 'Just's.
        case mds of
          Nothing -> prettyLast
          Just ds -> prettyAfter dimSep (fmap doc ds) <> prettyLast
        where
          prettyLast = doc d <> Pretty.text ":*"
      where
        dimSep = Pretty.text ", "

instance Out (Dims t a) => F.Pretty (Dims t a) where
    pprint' _ = doc

-- Faster is possible for non @[]@ list-likes, but this is OK for the general
-- case.
prettyIntersperse :: Foldable t => Pretty.Doc -> t Pretty.Doc -> Pretty.Doc
prettyIntersperse dBetween ds =
    case foldMap (\d -> [dBetween, d]) ds of
      []    -> mempty
      _:ds' -> mconcat ds'

prettyAfter :: Foldable t => Pretty.Doc -> t Pretty.Doc -> Pretty.Doc
prettyAfter dAfter = foldMap (\d -> d <> dAfter)

-- | Traverse over the functor in a 'Dims' value with a functor bound type.
--
-- For example, to turn a @'Dims' t ('Maybe' a)@ into a @'Maybe' ('Dims' t a)@.
dimsTraverse :: (Traversable t, Applicative f) => Dims t (f a) -> f (Dims t a)
dimsTraverse = traverse id
-- TODO provide a SPECIALIZE clause for the above Maybe case. performance! :)

-- | How many dimensions does the given 'Dims' represent?
dimsLength :: Foldable t => Dims t a -> Int
dimsLength = \case
  DimsExplicitShape ds -> length ds
  DimsAssumedShape ss -> length ss
  DimsAssumedSize mds _d ->
    case mds of
      Nothing -> 1
      Just ds -> length ds + 1