content: {-# LANGUAGE OverloadedStrings #-}
  
  module Futhark.IR.Mem.Interval
    ( Interval (..),
      distributeOffset,
      expandOffset,
      intervalOverlap,
      selfOverlap,
      primBool,
      intervalPairs,
      justLeafExp,
    )
  where
  
  import Data.Function (on)
  import Data.List (maximumBy, minimumBy, (\\))
  import Futhark.Analysis.AlgSimplify qualified as AlgSimplify
  import Futhark.Analysis.PrimExp.Convert
  import Futhark.IR.Prop
  import Futhark.IR.Syntax hiding (Result)
  import Futhark.Util
  
  data Interval = Interval
    { lowerBound :: TPrimExp Int64 VName,
      numElements :: TPrimExp Int64 VName,
      stride :: TPrimExp Int64 VName
    }
    deriving (Show, Eq)
  
  instance FreeIn Interval where
    freeIn' (Interval lb ne st) = freeIn' lb <> freeIn' ne <> freeIn' st
  
  distributeOffset :: (MonadFail m) => AlgSimplify.SofP -> [Interval] -> m [Interval]
  distributeOffset [] interval = pure interval
  distributeOffset offset [] = fail $ "Cannot distribute offset " <> show offset <> " across empty interval"
  distributeOffset offset [Interval lb ne 1] = pure [Interval (lb + TPrimExp (AlgSimplify.sumToExp offset)) ne 1]
  distributeOffset offset (Interval lb ne st0 : is)
    | st <- AlgSimplify.Prod False [untyped st0],
      Just (before, quotient, after) <- focusMaybe (`AlgSimplify.maybeDivide` st) offset =
        distributeOffset (before <> after) $
          Interval (lb + TPrimExp (AlgSimplify.sumToExp [quotient])) ne st0 : is
    | [st] <- AlgSimplify.simplify0 $ untyped st0,
      Just (before, quotient, after) <- focusMaybe (`AlgSimplify.maybeDivide` st) offset =
        distributeOffset (before <> after) $
          Interval (lb + TPrimExp (AlgSimplify.sumToExp [quotient])) ne st0 : is
    | otherwise = do
        rest <- distributeOffset offset is
        pure $ Interval lb ne st0 : rest
  
  findMostComplexTerm :: AlgSimplify.SofP -> (AlgSimplify.Prod, AlgSimplify.SofP)
  findMostComplexTerm prods =
    let max_prod = maximumBy (compare `on` (length . AlgSimplify.atoms)) prods
     in (max_prod, prods \\ [max_prod])
  
  findClosestStride :: [PrimExp VName] -> [Interval] -> (PrimExp VName, [PrimExp VName])
  findClosestStride offset_term is =
    let strides = map (untyped . stride) is
        p =
          minimumBy
            ( compare
                `on` ( termDifferenceLength
                         . minimumBy (compare `on` \s -> length (offset_term \\ AlgSimplify.atoms s))
                         . AlgSimplify.simplify0
                     )
            )
            strides
     in ( p,
          (offset_term \\) $
            AlgSimplify.atoms $
              minimumBy (compare `on` \s -> length (offset_term \\ AlgSimplify.atoms s)) $
                AlgSimplify.simplify0 p
        )
    where
      termDifferenceLength (AlgSimplify.Prod _ xs) = length (offset_term \\ xs)
  
  expandOffset :: AlgSimplify.SofP -> [Interval] -> Maybe AlgSimplify.SofP
  expandOffset [] _ = Nothing
  expandOffset offset i1
    | (AlgSimplify.Prod b term_to_add, offset_rest) <- findMostComplexTerm offset, -- Find gnb
      (closest_stride, first_term_divisor) <- findClosestStride term_to_add i1, -- find (nb-b, g)
      target <- [AlgSimplify.Prod b $ closest_stride : first_term_divisor], -- g(nb-b)
      diff <- AlgSimplify.sumOfProducts $ AlgSimplify.sumToExp $ AlgSimplify.Prod b term_to_add : map AlgSimplify.negate target, -- gnb - gnb + gb = gnb - g(nb-b)
      replacement <- target <> diff -- gnb = g(nb-b) + gnb - gnb + gb
      =
        Just (replacement <> offset_rest)
  
  intervalOverlap :: [(VName, PrimExp VName)] -> Names -> Interval -> Interval -> Bool
  intervalOverlap less_thans non_negatives (Interval lb1 ne1 st1) (Interval lb2 ne2 st2)
    | st1 == st2,
      AlgSimplify.lessThanish less_thans non_negatives lb1 lb2,
      AlgSimplify.lessThanish less_thans non_negatives (lb1 + ne1 - 1) lb2 =
        False
    | st1 == st2,
      AlgSimplify.lessThanish less_thans non_negatives lb2 lb1,
      AlgSimplify.lessThanish less_thans non_negatives (lb2 + ne2 - 1) lb1 =
        False
    | otherwise = True
  
  primBool :: TPrimExp Bool VName -> Maybe Bool
  primBool p
    | Just (BoolValue b) <- evalPrimExp (const Nothing) $ untyped p = Just b
    | otherwise = Nothing
  
  intervalPairs :: [Interval] -> [Interval] -> [(Interval, Interval)]
  intervalPairs = intervalPairs' []
    where
      intervalPairs' :: [(Interval, Interval)] -> [Interval] -> [Interval] -> [(Interval, Interval)]
      intervalPairs' acc [] [] = reverse acc
      intervalPairs' acc (i@(Interval lb _ st) : is) [] = intervalPairs' ((i, Interval lb 1 st) : acc) is []
      intervalPairs' acc [] (i@(Interval lb _ st) : is) = intervalPairs' ((Interval lb 1 st, i) : acc) [] is
      intervalPairs' acc (i1@(Interval lb1 _ st1) : is1) (i2@(Interval lb2 _ st2) : is2)
        | st1 == st2 = intervalPairs' ((i1, i2) : acc) is1 is2
        | otherwise =
            let res1 = intervalPairs' ((i1, Interval lb1 1 st1) : acc) is1 (i2 : is2)
                res2 = intervalPairs' ((Interval lb2 1 st2, i2) : acc) (i1 : is1) is2
             in if length res1 <= length res2
                  then res1
                  else res2
  
  -- | Returns true if the intervals are self-overlapping, meaning that for a
  -- given dimension d, the stride of d is larger than the aggregate spans of the
  -- lower dimensions.
  selfOverlap :: scope -> asserts -> [(VName, PrimExp VName)] -> [PrimExp VName] -> [Interval] -> Maybe Interval
  selfOverlap _ _ _ _ [_] = Nothing
  selfOverlap _ _ less_thans non_negatives' is
    | Just non_negatives <- namesFromList <$> mapM justLeafExp non_negatives' =
        -- TODO: Do we need to do something clever using some ranges of known values?
        let selfOverlap' acc (x : xs) =
              let interval_span = (lowerBound x + numElements x - 1) * stride x
                  res = AlgSimplify.lessThanish less_thans non_negatives (AlgSimplify.simplify' acc) (AlgSimplify.simplify' $ stride x)
               in if res then selfOverlap' (acc + interval_span) xs else Just x
            selfOverlap' _ [] = Nothing
         in selfOverlap' 0 $ reverse is
  selfOverlap _ _ _ _ (x : _) = Just x
  selfOverlap _ _ _ _ [] = Nothing
  
  justLeafExp :: PrimExp VName -> Maybe VName
  justLeafExp (LeafExp v _) = Just v
  justLeafExp _ = Nothing
  