{-|
A partition of time into contiguous spans, for defining reporting periods.
-}
module Hledger.Data.DayPartition
( DayPartition
, boundariesToDayPartition
, boundariesToMaybeDayPartition

, lookupDayPartition
, unionDayPartitions

, dayPartitionToNonEmpty
, dayPartitionToList
, dayPartitionSpans
, dayPartitionToDateSpans
, dayPartitionToPeriodData
, maybeDayPartitionToDateSpans

, splitSpan
, intervalBoundaryBefore

, tests_DayPartition
) where

import qualified Data.IntMap.Strict as IM
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NE
import Data.Time (Day(..), addDays, addGregorianMonthsClip, addGregorianYearsClip, fromGregorian)

import Hledger.Data.Dates
import Hledger.Data.PeriodData
import Hledger.Data.Types
import Hledger.Utils


-- | A partition of time into contiguous spans, along with a historical period
-- before any of the spans.
--
-- This is a newtype wrapper around 'PeriodData Day', where the start dates are
-- the keys and the end dates are the values. Spans are stored in inclusive format
-- [start, end]. Note that this differs from 'DateSpan' which uses [start, end)
-- format.
--
-- The constructor is not exported so that we can ensure the spans are valid
-- partitions of time.
newtype DayPartition = DayPartition { dayPartitionToPeriodData :: PeriodData Day } deriving (Eq, Ord, Show)

-- Developer's note. All constructors must guarantee that:
--   1. The value stored in pdperiods has at least one key.
--   2. The value stored in pdpre equals one day before the smallest key in pdperiods.
--   3. The value stored in each entry of pdperiods equals one day before the
--      next largest key, except for the value associated to the largest key.
isValidDayPartition :: DayPartition -> Bool
isValidDayPartition (DayPartition pd) = case ds of
    [] -> False  -- Must be at least one key in pdperiods
    xs -> and $ zipWith isContiguous ((nulldate, h) : xs) xs
  where
    (h, ds) = periodDataToList pd
    isContiguous (_, e) (s, _) = addDays 1 e == s


-- | Construct a 'DayPartition' from a non-empty list of boundary days.
boundariesToDayPartition :: NonEmpty Day -> DayPartition
boundariesToDayPartition xs = DayPartition . periodDataFromList (addDays (-1) b) $ case bs of
    []  -> [(b, b)]  -- If only one boundary is supplied, it ends on the same day
    _:_ -> zip (b:bs) $ map (addDays (-1)) bs  -- Guaranteed non-empty
  where b:|bs = NE.nub $ NE.sort xs

-- | Construct a 'DayPartition' from a list of boundary days, returning
-- 'Nothing' for the empty list.
boundariesToMaybeDayPartition :: [Day] -> Maybe DayPartition
boundariesToMaybeDayPartition = fmap boundariesToDayPartition . NE.nonEmpty


-- | Find the span of a 'DayPartition' which contains a given day.
lookupDayPartition :: Day -> DayPartition -> (Maybe Day, Day)
lookupDayPartition d (DayPartition xs) = lookupPeriodDataOrHistorical d xs

-- | Return the union of two 'DayPartition's if they are consistent, or 'Nothing' otherwise.
unionDayPartitions :: DayPartition -> DayPartition -> Maybe DayPartition
unionDayPartitions (DayPartition (PeriodData h as)) (DayPartition (PeriodData h' as')) =
    if equalIntersection as as' && isValidDayPartition union then Just union else Nothing
  where
    union = DayPartition . PeriodData (min h h') $ as <> as'
    equalIntersection x y = and $ IM.intersectionWith (==) x y


-- | Convert 'DayPartition' to a non-empty list of start and end dates for the periods.
--
-- Note that the end date of each period will be one day before the start date
-- of the next period.
dayPartitionToNonEmpty :: DayPartition -> NonEmpty (Day, Day)
dayPartitionToNonEmpty (DayPartition xs) = NE.fromList . snd $ periodDataToList xs  -- Constructors guarantee this is non-empty

-- | Convert 'DayPartition' to a list of start and end dates for the periods.
--
-- Note that the end date of each period will be one day before the start date
-- of the next period.
dayPartitionToList :: DayPartition -> [(Day, Day)]
dayPartitionToList = NE.toList . dayPartitionToNonEmpty

-- | Return the whole day range spanned by a `PeriodData Day`.
dayPartitionSpans :: DayPartition -> (Day, Day)
dayPartitionSpans (DayPartition (PeriodData _ ds)) =
    -- Guaranteed not to error because the IntMap in non-empty.
    (intToDay . fst $ IM.findMin ds, snd $ IM.findMax ds)

-- | Convert 'DayPartition' to a list of 'DateSpan's.
--
-- Note that the end date of each period will be equal to the start date of
-- the next period.
dayPartitionToDateSpans :: DayPartition -> [DateSpan]
dayPartitionToDateSpans = map toDateSpan . dayPartitionToList
  where
    toDateSpan (s, e) = DateSpan (toEFDay s) (toEFDay $ addDays 1 e)
    toEFDay = Just . Exact

-- Convert a periodic report 'Maybe DayPartition' to a list of 'DateSpans',
-- replacing the empty case with an appropriate placeholder.
--
-- Note that the end date of each period will be equal to the start date of
-- the next period.
maybeDayPartitionToDateSpans :: Maybe DayPartition -> [DateSpan]
maybeDayPartitionToDateSpans = maybe [DateSpan Nothing Nothing] dayPartitionToDateSpans


-- | Split a 'DateSpan' into a 'DayPartition' consisting of consecutive exact
-- spans of the specified Interval, or `Nothing` if the span is invalid.
-- If no interval is specified, the original span is returned.
-- If the original span is the null date span, ie unbounded, `Nothing` is returned.
-- If the original span is empty, eg if the end date is <= the start date, `Nothing` is returned.
--
-- ==== Date adjustment
-- Some intervals respect the "adjust" flag (years, quarters, months, weeks, every Nth weekday
-- of month seem to be the ones that need it). This will move the start date earlier, if needed,
-- to the previous natural interval boundary (first of year, first of quarter, first of month,
-- monday, previous Nth weekday of month). Related: #1982 #2218
--
-- The end date is always moved later if needed to the next natural interval boundary,
-- so that the last period is the same length as the others.
--
-- ==== Examples
-- >>> let t i y1 m1 d1 y2 m2 d2 = fmap dayPartitionToNonEmpty . splitSpan True i $ DateSpan (Just $ Flex $ fromGregorian y1 m1 d1) (Just $ Flex $ fromGregorian y2 m2 d2)
-- >>> t NoInterval 2008 01 01 2009 01 01
-- Just ((2008-01-01,2008-12-31) :| [])
-- >>> t (Quarters 1) 2008 01 01 2009 01 01
-- Just ((2008-01-01,2008-03-31) :| [(2008-04-01,2008-06-30),(2008-07-01,2008-09-30),(2008-10-01,2008-12-31)])
-- >>> splitSpan True (Quarters 1) nulldatespan
-- Nothing
-- >>> t (Days 1) 2008 01 01 2008 01 01  -- an empty datespan
-- Nothing
-- >>> t (Quarters 1) 2008 01 01 2008 01 01
-- Nothing
-- >>> t (Months 1) 2008 01 01 2008 04 01
-- Just ((2008-01-01,2008-01-31) :| [(2008-02-01,2008-02-29),(2008-03-01,2008-03-31)])
-- >>> t (Months 2) 2008 01 01 2008 04 01
-- Just ((2008-01-01,2008-02-29) :| [(2008-03-01,2008-04-30)])
-- >>> t (Weeks 1) 2008 01 01 2008 01 15
-- Just ((2007-12-31,2008-01-06) :| [(2008-01-07,2008-01-13),(2008-01-14,2008-01-20)])
-- >>> t (Weeks 2) 2008 01 01 2008 01 15
-- Just ((2007-12-31,2008-01-13) :| [(2008-01-14,2008-01-27)])
-- >>> t (MonthDay 2) 2008 01 01 2008 04 01
-- Just ((2008-01-02,2008-02-01) :| [(2008-02-02,2008-03-01),(2008-03-02,2008-04-01)])
-- >>> t (NthWeekdayOfMonth 2 4) 2011 01 01 2011 02 15
-- Just ((2010-12-09,2011-01-12) :| [(2011-01-13,2011-02-09),(2011-02-10,2011-03-09)])
-- >>> t (DaysOfWeek [2]) 2011 01 01 2011 01 15
-- Just ((2010-12-28,2011-01-03) :| [(2011-01-04,2011-01-10),(2011-01-11,2011-01-17)])
-- >>> t (MonthAndDay 11 29) 2012 10 01 2013 10 15
-- Just ((2012-11-29,2013-11-28) :| [])
splitSpan :: Bool -> Interval -> DateSpan -> Maybe DayPartition
splitSpan _      _                        (DateSpan Nothing Nothing) = Nothing
splitSpan _      _                        ds | isEmptySpan ds = Nothing
splitSpan _      NoInterval               (DateSpan (Just s) (Just e)) = Just $ boundariesToDayPartition (fromEFDay s :| [fromEFDay e])
splitSpan _      NoInterval               _  = Nothing
splitSpan _      (Days n)                 ds = splitspan id addDays n ds
splitSpan adjust (Weeks n)                ds = splitspan (if adjust then startofweek    else id) addDays                 (7*n) ds
splitSpan adjust (Months n)               ds = splitspan (if adjust then startofmonth   else id) addGregorianMonthsClip  n     ds
splitSpan adjust (Quarters n)             ds = splitspan (if adjust then startofquarter else id) addGregorianMonthsClip  (3*n) ds
splitSpan adjust (Years n)                ds = splitspan (if adjust then startofyear    else id) addGregorianYearsClip   n     ds
splitSpan adjust (NthWeekdayOfMonth n wd) ds = splitspan (startWeekdayOfMonth n wd)              advancemonths           1     ds
  where
    startWeekdayOfMonth = if adjust then prevNthWeekdayOfMonth else nextNthWeekdayOfMonth
    advancemonths 0 = id
    advancemonths m = advanceToNthWeekday n wd . startofmonth . addGregorianMonthsClip m
splitSpan _      (MonthDay dom)           ds = splitspan (nextnthdayofmonth dom) (addGregorianMonthsToMonthday dom) 1 ds
splitSpan _      (MonthAndDay m d)        ds = splitspan (nextmonthandday m d)   addGregorianYearsClip              1 ds
splitSpan _      (DaysOfWeek [])          _  = Nothing
splitSpan _      (DaysOfWeek days@(n:_))  ds = do
    (s, e) <- dateSpanSplitLimits (nthdayofweekcontaining n) nextday ds
    let -- can't show this when debugging, it'll hang:
        bdrys = concatMap (\d -> map (addDays d) starts) [0,7..]
        -- The first representative of each weekday
        starts = map (\d -> addDays (toInteger $ d - n) $ nthdayofweekcontaining n s) days
    spansFromBoundaries e bdrys


-- | Fill in missing start/end dates for calculating 'splitSpan'.
dateSpanSplitLimits :: (Day -> Day) -> (Day -> Day) -> DateSpan -> Maybe (Day, Day)
dateSpanSplitLimits _     _    (DateSpan Nothing   Nothing) = Nothing
dateSpanSplitLimits _     _    ds | isEmptySpan ds          = Nothing
dateSpanSplitLimits start _    (DateSpan (Just s) (Just e)) = Just (start $ fromEFDay s, fromEFDay e)
dateSpanSplitLimits start next (DateSpan (Just s) Nothing)  = Just (start $ fromEFDay s, next $ start $ fromEFDay s)
dateSpanSplitLimits start next (DateSpan Nothing  (Just e)) = Just (start $ fromEFDay e, next $ start $ fromEFDay e)

-- Split the given span into exact spans using the provided helper functions:
--
-- 1. The start function is used to adjust the provided span's start date to get the first sub-span's start date.
--
-- 2. The next function is used to calculate subsequent sub-spans' start dates, possibly with stride increased by a multiplier.
--    It should handle spans of varying length, eg when splitting on "every 31st of month",
--    it adjusts to 28/29/30 in short months but returns to 31 in the long months.
splitspan :: (Day -> Day) -> (Integer -> Day -> Day) -> Int -> DateSpan -> Maybe DayPartition
splitspan start next mult ds = do
    (s, e) <- dateSpanSplitLimits start (next (toInteger mult)) ds
    let bdrys = mapM (next . toInteger) [0,mult..] $ start s
    spansFromBoundaries e bdrys

-- | Construct a list of exact 'DateSpan's from a list of boundaries, which fit within a given range.
spansFromBoundaries :: Day -> [Day] -> Maybe DayPartition
spansFromBoundaries _ []     = Nothing
spansFromBoundaries e (x:_)  | x >= e = Nothing
spansFromBoundaries e (x:xs) = Just . boundariesToDayPartition $ takeUntilFailsNE (<e) (x:|xs)


-- Get the natural start for the given interval that falls on or before the given day,
-- when applicable. Works for Weeks, Months, Quarters, Years, eg.
intervalBoundaryBefore :: Interval -> Day -> Day
intervalBoundaryBefore i d =
  case dayPartitionToNonEmpty <$> splitSpan True i (DateSpan (Just $ Exact d) (Just . Exact $ addDays 1 d)) of
    Just ((start, _) :| _ ) -> start
    _ -> d


intToDay = ModifiedJulianDay . toInteger


tests_DayPartition =
  testGroup "splitSpan" [
      testCase "weekday" $ do
        fmap dayPartitionToNonEmpty (splitSpan False (DaysOfWeek [1..5]) (DateSpan (Just $ Exact $ fromGregorian 2021 07 01) (Just $ Exact $ fromGregorian 2021 07 08)))
          @?= Just ( (fromGregorian 2021 06 28, fromGregorian 2021 06 28) :|
                   [ (fromGregorian 2021 06 29, fromGregorian 2021 06 29)
                   , (fromGregorian 2021 06 30, fromGregorian 2021 06 30)
                   , (fromGregorian 2021 07 01, fromGregorian 2021 07 01)
                   , (fromGregorian 2021 07 02, fromGregorian 2021 07 04)
                   -- next week
                   , (fromGregorian 2021 07 05, fromGregorian 2021 07 05)
                   , (fromGregorian 2021 07 06, fromGregorian 2021 07 06)
                   , (fromGregorian 2021 07 07, fromGregorian 2021 07 07)
                   ])

        fmap dayPartitionToNonEmpty (splitSpan False (DaysOfWeek [1, 5]) (DateSpan (Just $ Exact $ fromGregorian 2021 07 01) (Just $ Exact $ fromGregorian 2021 07 08)))
          @?= Just ( (fromGregorian 2021 06 28, fromGregorian 2021 07 01) :|
                   [ (fromGregorian 2021 07 02, fromGregorian 2021 07 04)
                   -- next week
                   , (fromGregorian 2021 07 05, fromGregorian 2021 07 08)
                   ])

    , testCase "match dayOfWeek" $ do
        let dayofweek n = splitspan (nthdayofweekcontaining n) (\w -> (if w == 0 then id else applyN (n-1) nextday . applyN (fromInteger w) nextweek)) 1
            matchdow ds day = splitSpan False (DaysOfWeek [day]) ds @?= dayofweek day ds
            ys2021 = fromGregorian 2021 01 01
            ye2021 = fromGregorian 2021 12 31
            ys2022 = fromGregorian 2022 01 01
        mapM_ (matchdow (DateSpan (Just $ Exact ys2021) (Just $ Exact ye2021))) [1..7]
        mapM_ (matchdow (DateSpan (Just $ Exact ys2021) (Just $ Exact ys2022))) [1..7]
        mapM_ (matchdow (DateSpan (Just $ Exact ye2021) (Just $ Exact ys2022))) [1..7]

        mapM_ (matchdow (DateSpan (Just $ Exact ye2021) Nothing)) [1..7]
        mapM_ (matchdow (DateSpan (Just $ Exact ys2022) Nothing)) [1..7]

        mapM_ (matchdow (DateSpan Nothing (Just $ Exact ye2021))) [1..7]
        mapM_ (matchdow (DateSpan Nothing (Just $ Exact ys2022))) [1..7]

    ]