sum-pyramid-0.0.1: src/UniqueLogic.hs

module UniqueLogic where

import Common

import qualified UniqueLogic.ST.TF.Rule as Rule
import qualified UniqueLogic.ST.TF.System.Simple as Sys
import qualified UniqueLogic.ST.TF.ZeroFractional as ZeroFrac

import Control.Monad.ST (ST, runST)

import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.Map as Map
import Data.Array.Comfort.Boxed ((!))
import Data.Map (Map)
import Data.Set (Set)
import Data.Foldable (for_, traverse_)
import Data.Maybe (isJust)


{- $setup
>>> import qualified Data.Array.Comfort.Boxed as BoxedArray
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import qualified Data.Map as Map
>>> import Common
-}


type Variable s a = Sys.Variable (ST s) a
type System s = Sys.T (ST s)

system ::
   BoxedArray.Array
      (Shape.LowerTriangular ShapeOp)
      (Variable s a -> Variable s a -> Variable s a -> System s ()) ->
   ST s
      (BoxedArray.Array (Shape.LowerTriangular ShapeInt) (Variable s a),
       System s ())
system ops = do
   let n = sizeFromOps ops
   vars <-
      Trav.sequence $
      BoxedArray.replicate
         (Shape.lowerTriangular $ Shape.ZeroBased n)
         Sys.globalVariable
   return
      (vars,
       traverse_
         (\((i,j),rule) -> rule (vars!(i+1,j)) (vars!(i+1,j+1)) (vars!(i,j)))
         (BoxedArray.toAssociations ops))


{- |
>>> solve 3 $ Map.fromList [((0,0),1), ((1,0),1), ((2,0),1::Integer)]
BoxedArray...Triangular... 3... [Just 1,Just 1,Just 0,Just 1,Just 0,Just 0]

>>> solve 3 $ Map.fromList [((0,0),1), ((2,0),1), ((2,2),1::Integer)]
BoxedArray...Triangular... 3... [Just 1,Nothing,Nothing,Just 1,Nothing,Just 1]
-}
solve ::
   (Num a) =>
   Int -> Map (Int,Int) a ->
   BoxedArray.Array (Shape.LowerTriangular ShapeInt) (Maybe a)
solve n xs =
   runST
      (do
         (vars, sys) <-
            system $
            BoxedArray.replicate (Shape.lowerTriangular $ ShapeOp n) Rule.add
         Sys.solve $ do
            sys
            for_ (Map.toList xs) $ \(ij,x) ->
               Rule.equ (vars!ij) =<< Sys.constant x
         traverse Sys.query vars)

{- |
>>> solveMixed (BoxedArray.fromList (Shape.lowerTriangular $ ShapeOp 2) [Mul]) (Map.fromList [((0,0),0), ((1,1),0::Rational)])
BoxedArray...Triangular... 2... [Just (0 % 1),Nothing,Just (0 % 1)]
>>> solveMixed (BoxedArray.fromList (Shape.lowerTriangular $ ShapeOp 2) [Mul]) (Map.fromList [((0,0),0), ((1,1),5::Rational)])
BoxedArray...Triangular... 2... [Just (0 % 1),Just (0 % 1),Just (5 % 1)]
-}
solveMixed ::
   (ZeroFrac.C a) =>
   BoxedArray.Array (Shape.LowerTriangular ShapeOp) Op ->
   Map (Int,Int) a ->
   BoxedArray.Array (Shape.LowerTriangular ShapeInt) (Maybe a)
solveMixed ops xs =
   runST
      (do
         let rule op =
               case op of
                  Add -> Rule.add
                  Mul -> Rule.mul
         (vars, sys) <- system $ fmap rule ops
         Sys.solve $ do
            sys
            for_ (Map.toList xs) $ \(ij,x) ->
               Rule.equ (vars!ij) =<< Sys.constant x
         traverse Sys.query vars)


solvable :: Int -> Set (Int, Int) -> Bool
solvable n = Fold.all isJust . solve n . Map.fromSet (const (0::Integer))

{- |
>>> map (length . solvables) [0..5]
[1,1,3,16,122,1188]
-}
solvables :: Int -> [Set (Int,Int)]
solvables n = filter (solvable n) $ allCellSelections n


solvableMixed :: SolutionCheck Integer
solvableMixed ops puzzle =
   Fold.all isJust $ solveMixed ops $ fmap toRational puzzle