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 qualified Combinatorics

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 Data.Array.Comfort.Boxed ((!))
import Data.Foldable (for_, traverse_)
import Data.Maybe (isJust)
import Data.Tuple.HT (mapSnd)


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


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

system ::
   BoxedArray.Array
      (Shape.LowerTriangular ShapeInt)
      (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 [((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 [((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 -> [((Int,Int),a)] ->
   BoxedArray.Array (Shape.LowerTriangular ShapeInt) (Maybe a)
solve n xs =
   runST
      (do
         (vars, sys) <-
            system $
            BoxedArray.replicate
               (Shape.lowerTriangular $ Shape.ZeroBased $ n-1)
               Rule.add
         Sys.solve $ do
            sys
            for_ xs $ \(ij,x) ->
               Rule.equ (vars!ij) =<< Sys.constant x
         traverse Sys.query vars)

{- |
>>> solveMixed (BoxedArray.fromList (Shape.lowerTriangular $ Shape.ZeroBased 1) [Mul]) [((0,0),0), ((1,1),0::Rational)]
BoxedArray...Triangular... 2... [Just (0 % 1),Nothing,Just (0 % 1)]
>>> solveMixed (BoxedArray.fromList (Shape.lowerTriangular $ Shape.ZeroBased 1) [Mul]) [((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 ShapeInt) Op ->
   [((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_ xs $ \(ij,x) ->
               Rule.equ (vars!ij) =<< Sys.constant x
         traverse Sys.query vars)


solvable :: Int -> [(Int, Int)] -> Bool
solvable n = Fold.all isJust . solve n . map (\ij -> (ij, 0::Integer))

solvables :: Int -> [[(Int,Int)]]
solvables n =
   filter (solvable n) $
   Combinatorics.tuples n $ Shape.indices $
   Shape.lowerTriangular $ Shape.ZeroBased n


solvableMixed :: SolutionCheck Integer
solvableMixed ops puzzle =
   Fold.all isJust $ solveMixed ops $ map (mapSnd toRational) puzzle