-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.SquareBirthday
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- As of January 2026, to access the careers link at <http://math.inc>, you need to solve the following
-- puzzle:
--
-- @
-- Suppose that today is June 1, 2025. We call a date "square" if all of its components (day, month, and year) are
-- perfect squares. I was born in the last millennium, and my next birthday (relative to that date) will be the last
-- square date in my life. If you sum the square roots of the components of that upcoming square birthday
-- (day, month, year), you obtain my age on June 1, 2025. My mother would have been born on a square date if the month
-- were a square number; in reality it is not a square date, but both the month and day are perfect cubes. When was
-- I born, and when was my mother born?
-- @
--
-- So, let's solve it using SBV.
-----------------------------------------------------------------------------

{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE OverloadedRecordDot #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.SquareBirthday where

import Prelude hiding (fromEnum, toEnum)

import Data.SBV
import Data.SBV.Control

import qualified Data.SBV.List  as SL
import qualified Data.SBV.Tuple as ST

-- | Months in a year.
data Month = Jan | Feb | Mar | Apr | May | Jun
           | Jul | Aug | Sep | Oct | Nov | Dec
           deriving Int -> Month -> ShowS
[Month] -> ShowS
Month -> String
(Int -> Month -> ShowS)
-> (Month -> String) -> ([Month] -> ShowS) -> Show Month
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Month -> ShowS
showsPrec :: Int -> Month -> ShowS
$cshow :: Month -> String
show :: Month -> String
$cshowList :: [Month] -> ShowS
showList :: [Month] -> ShowS
Show

-- | A date. We use unbounded integers for day and year, which simplifies coding,
-- though one can also enumerate the possible values from the problem itself.
data Date = MkDate { Date -> Integer
day   :: Integer
                   , Date -> Month
month :: Month
                   , Date -> Integer
year  :: Integer
                   }

-- | Make 'Month' and 'Date' usable in symbolic contexts.
mkSymbolic [''Month, ''Date]

-- | Show instance for date, for pretty-printing.
instance Show Date where
  show :: Date -> String
show (MkDate Integer
d Month
m Integer
y) = Month -> String
forall a. Show a => a -> String
show Month
m String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
pad String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
d String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
y
   where pad :: String
pad | Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
10 = String
" "
             | Bool
True   = String
""

-- | Get a symbolic date with the given name. Since we used
-- integers for the day and year fields, we constrain them
-- appropriately. Note that one can further constrain days
-- based on the year and month; but that level detail isn't
-- necessary for the current problem.
symDate :: String -> Symbolic SDate
symDate :: String -> Symbolic (SBV Date)
symDate String
nm = do SBV Date
dt <- String -> Symbolic (SBV Date)
forall a. SymVal a => String -> Symbolic (SBV a)
free String
nm

                SBool -> SymbolicT IO ()
forall a. QuantifiedBool a => a -> SymbolicT IO ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain [sCase|Date dt of
                              MkDate d _ y -> sAnd [ 1 .<= d, d .<= 31
                                                   , 0 .<= y
                                                   ]
                          |]

                SBV Date -> Symbolic (SBV Date)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SBV Date
dt

-- | Encode today as a symbolic value. The puzzle says today is June 1st, 2025.
today :: SDate
today :: SBV Date
today = Date -> SBV Date
forall a. SymVal a => a -> SBV a
literal (Date -> SBV Date) -> Date -> SBV Date
forall a b. (a -> b) -> a -> b
$ MkDate { day :: Integer
day   =    Integer
1
                         , month :: Month
month =  Month
Jun
                         , year :: Integer
year  = Integer
2025
                         }

-- | A date is on or after another, if the month-day combo is
-- lexicographically later. Note that we ignore the year for this
-- comparison, as we're interested if the anniversary of a date is after or not.
onOrAfter :: SDate -> SDate -> SBool
SBV Date
d1 onOrAfter :: SBV Date -> SBV Date -> SBool
`onOrAfter` SBV Date
d2 = (SBV Date -> SBV Month
smonth SBV Date
d1, SBV Date -> SBV Integer
sday SBV Date
d1) (SBV Month, SBV Integer) -> (SBV Month, SBV Integer) -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= (SBV Date -> SBV Month
smonth SBV Date
d2, SBV Date -> SBV Integer
sday SBV Date
d2)

-- | Similar to 'onOrAfter', except we require strictly later.
after :: SDate -> SDate -> SBool
SBV Date
d1 after :: SBV Date -> SBV Date -> SBool
`after` SBV Date
d2 = (SBV Date -> SBV Month
smonth SBV Date
d1, SBV Date -> SBV Integer
sday SBV Date
d1) (SBV Month, SBV Integer) -> (SBV Month, SBV Integer) -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>  (SBV Date -> SBV Month
smonth SBV Date
d2, SBV Date -> SBV Integer
sday SBV Date
d2)

-- | The age based on a given date is the difference between years less than one.
-- We have to adjust by 1 if today happens to be after the given date.
age :: SDate -> SInteger
age :: SBV Date -> SBV Integer
age SBV Date
d = SBV Date -> SBV Integer
syear SBV Date
today SBV Integer -> SBV Integer -> SBV Integer
forall a. Num a => a -> a -> a
- SBV Date -> SBV Integer
syear SBV Date
d SBV Integer -> SBV Integer -> SBV Integer
forall a. Num a => a -> a -> a
- SBV Integer
1 SBV Integer -> SBV Integer -> SBV Integer
forall a. Num a => a -> a -> a
+ SBool -> SBV Integer
forall a. (Ord a, Num (SBV a), SymVal a) => SBool -> SBV a
oneIf (SBV Date
today SBV Date -> SBV Date -> SBool
`after` SBV Date
d)

-- | We can let years to range over arbitrary integers. But that complicates the
-- job of the solver. So, based on what we know from the problem, we restrict
-- our attention to years betweek 1900 and 2100. Note that there are only
-- two years that satisfy this in that range: 1936 and 2025. (Any other square
-- year makes no sense for the setting of the problem.) To simplify the square-root
-- computation, we also store the square root in this list as the second component:
--
-- >>> squareYears
-- [(1936,44),(2025,45)]
squareYears :: [(Integer, Integer)]
squareYears :: [(Integer, Integer)]
squareYears = ((Integer, Integer) -> Bool)
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\(Integer
y, Integer
_) -> Integer
y Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2100)
            ([(Integer, Integer)] -> [(Integer, Integer)])
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ ((Integer, Integer) -> Bool)
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (\(Integer
y, Integer
_) -> Integer
y Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1900)
            ([(Integer, Integer)] -> [(Integer, Integer)])
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ [(Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
i, Integer
i) | Integer
i <- [Integer
1::Integer ..]]

-- | A date is square if all its components are.
squareDate :: SDate -> SBool
squareDate :: SBV Date -> SBool
squareDate SBV Date
dt = [sCase|Date dt of
                   MkDate d m y -> squareDay d .&& squareMonth m .&& squareYear y
                |]
  where squareDay :: a -> SBool
squareDay   a
d = a
d a -> [a] -> SBool
forall a. EqSymbolic a => a -> [a] -> SBool
`sElem` [a
1, a
4, a
9, a
16, a
25]
        squareMonth :: SBV Month -> SBool
squareMonth SBV Month
m = SBV Month
m SBV Month -> [SBV Month] -> SBool
forall a. EqSymbolic a => a -> [a] -> SBool
`sElem` [SBV Month
sJan, SBV Month
sApr, SBV Month
sSep]
        squareYear :: SBV Integer -> SBool
squareYear  SBV Integer
y = SBV Integer
y SBV Integer -> [SBV Integer] -> SBool
forall a. EqSymbolic a => a -> [a] -> SBool
`sElem` ((Integer, Integer) -> SBV Integer)
-> [(Integer, Integer)] -> [SBV Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> SBV Integer
forall a. SymVal a => a -> SBV a
literal (Integer -> SBV Integer)
-> ((Integer, Integer) -> Integer)
-> (Integer, Integer)
-> SBV Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst) [(Integer, Integer)]
squareYears


-- | Summing the square-roots of the components of a date.
sqrSum :: SDate -> SInteger
sqrSum :: SBV Date -> SBV Integer
sqrSum SBV Date
dt = [sCase|Date dt of
               MkDate d m y -> r d + mr m + r y
            |]
 where r :: SBV Integer -> SBV Integer
r SBV Integer
v  = SBV Integer
v SBV Integer -> SList (Integer, Integer) -> SBV Integer
forall k v. (SymVal k, SymVal v) => SBV k -> SList (k, v) -> SBV v
`SL.lookup` [(Integer, Integer)] -> SList (Integer, Integer)
forall a. SymVal a => a -> SBV a
literal ([(Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
i, Integer
i) | Integer
i <- [Integer
1, Integer
2, Integer
3, Integer
4, Integer
5]] [(Integer, Integer)]
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a. [a] -> [a] -> [a]
++ [(Integer, Integer)]
squareYears)

       mr :: SMonth -> SInteger
       mr :: SBV Month -> SBV Integer
mr SBV Month
m = [sCase|Month m of
                  Jan -> 1
                  Apr -> 2
                  Sep -> 3
                  _   -> some "Non-Square Month" (const sTrue)
              |]

-- | Formalizing the puzzle. We literally write down the description in
-- SBV notation. As with any formalization, this step is subjective; there
-- could be many different ways to express the same problem. The description
-- below is quite faithful to the problem description given. We have:
--
-- >>> puzzle
-- Me : Sep 25, 1971
-- Mom: Aug  1, 1936
puzzle :: IO ()
puzzle :: IO ()
puzzle = SymbolicT IO () -> IO ()
forall a. Symbolic a -> IO a
runSMT (SymbolicT IO () -> IO ()) -> SymbolicT IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do

    -----------------------------------
    -- Constraints about my birthday
    -----------------------------------
    SBV Date
myBirthday <- String -> Symbolic (SBV Date)
symDate String
"My Birthday"

    -- I was born in the last millenium
    SBool -> SymbolicT IO ()
forall a. QuantifiedBool a => a -> SymbolicT IO ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> SymbolicT IO ()) -> SBool -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ SBV Date -> SBV Integer
syear SBV Date
myBirthday SBV Integer -> SBV Integer -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.< SBV Integer
2000 SBool -> SBool -> SBool
.&& SBV Date -> SBV Integer
syear SBV Date
myBirthday SBV Integer -> SBV Integer -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= SBV Integer
1900

    -- My next birthday will be a square
    let next :: SBV Date
next = [sCase|Date myBirthday of
                  MkDate d m _ -> sMkDate d m (syear today + oneIf (today `onOrAfter` myBirthday))
               |]

    SBool -> SymbolicT IO ()
forall a. QuantifiedBool a => a -> SymbolicT IO ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> SymbolicT IO ()) -> SBool -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ SBV Date -> SBool
squareDate SBV Date
next

    -- And it'll be the last square day of my life, so we maximize the metric corresponding to the
    -- date. We turn it into a 3-tuple of year, month, date over integers, which preserves the
    -- order of the dates.
    String -> SBV (Integer, Integer, Integer) -> SymbolicT IO ()
forall a. Metric a => String -> SBV a -> SymbolicT IO ()
maximize String
"Next Birthday Latest" (SBV (Integer, Integer, Integer) -> SymbolicT IO ())
-> SBV (Integer, Integer, Integer) -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ (SBV Integer, SBV Integer, SBV Integer)
-> SBV (Integer, Integer, Integer)
forall tup a. Tuple tup a => a -> SBV tup
ST.tuple (SBV Date -> SBV Integer
syear SBV Date
next, SBV Month -> SBV Integer
forall a. EnumSymbolic a => SBV a -> SBV Integer
fromEnum (SBV Date -> SBV Month
smonth SBV Date
next), SBV Date -> SBV Integer
sday SBV Date
next)

    -- If you square the components of my next birthday, it gives me my current age on Jun 1, 2025
    SBool -> SymbolicT IO ()
forall a. QuantifiedBool a => a -> SymbolicT IO ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain (SBool -> SymbolicT IO ()) -> SBool -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ SBV Date -> SBV Integer
sqrSum SBV Date
next SBV Integer -> SBV Integer -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== SBV Date -> SBV Integer
age SBV Date
myBirthday

    -----------------------------------
    -- Constraints about mom's birthday
    -----------------------------------
    SBV Date
momBirthday <- String -> Symbolic (SBV Date)
symDate String
"Mom's Birthday"

    -- Mom has a square birth-date, except for the month:
    SBool -> SymbolicT IO ()
forall a. QuantifiedBool a => a -> SymbolicT IO ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain [sCase|Date momBirthday of
                 MkDate d _ y -> squareDate (sMkDate d sJan y)
              |]

    -- Mom's day and month are perfect cubes
    SBool -> SymbolicT IO ()
forall a. QuantifiedBool a => a -> SymbolicT IO ()
forall (m :: * -> *) a.
(SolverContext m, QuantifiedBool a) =>
a -> m ()
constrain [sCase|Date momBirthday of
                 MkDate d m _ -> sAnd [ d `sElem` [1, 8, 27]
                                      , m `sElem` [sJan, sAug]
                                      ]
              |]

    -- Extract the results:
    Query () -> SymbolicT IO ()
forall a. Query a -> Symbolic a
query (Query () -> SymbolicT IO ()) -> Query () -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ do CheckSatResult
cs <- Query CheckSatResult
checkSat
               case CheckSatResult
cs of
                 CheckSatResult
Sat -> do Date
me  <- SBV Date -> Query Date
forall a. SymVal a => SBV a -> Query a
getValue SBV Date
myBirthday
                           Date
mom <- SBV Date -> Query Date
forall a. SymVal a => SBV a -> Query a
getValue SBV Date
momBirthday

                           IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ do String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Me : " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Date -> String
forall a. Show a => a -> String
show Date
me
                                   String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Mom: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Date -> String
forall a. Show a => a -> String
show Date
mom

                 CheckSatResult
_   -> String -> Query ()
forall a. HasCallStack => String -> a
error (String -> Query ()) -> String -> Query ()
forall a b. (a -> b) -> a -> b
$ String
"Unexpected result: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ CheckSatResult -> String
forall a. Show a => a -> String
show CheckSatResult
cs