packages feed

sum-pyramid-0.0: src/Main.hs

module Main where

import qualified LinearAlgebra as LinAlg
import qualified UniqueLogic as Logic
import Common

import qualified Combinatorics

import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS
import Control.Monad (replicateM, join)
import Control.Applicative (pure, (<*>), (<|>))

import qualified System.Random as Random

import Text.Printf (printf)

import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.List.HT as ListHT
import qualified Data.Set as Set
import Data.Array.Comfort.Boxed (Array, (!))
import Data.Foldable (for_)
import Data.Set (Set)
import Data.Tuple.HT (mapPair, mapFst)

import qualified Options.Applicative as OP
import Shell.Utility.ParseArgument (parseNumber)



randomR :: (Random.RandomGen g, Random.Random a) => (a,a) -> MS.State g a
randomR rng = MS.state $ Random.randomR rng

pick :: (Random.RandomGen g) => MS.StateT (Set a) (MS.State g) a
pick = do
   set <- MS.get
   k <- MT.lift $ randomR (0, Set.size set - 1)
   MS.put $ Set.deleteAt k set
   return $ Set.elemAt k set


data Allowed = Allowed {allowedAdd, allowedMul :: Bool}

pyramid ::
   (Random.RandomGen g) =>
   Allowed ->
   Array ShapeInt Integer ->
   MS.State g
      (Array (Shape.LowerTriangular ShapeInt) Op,
       Array (Shape.LowerTriangular ShapeInt) Integer)
pyramid allowed base = do
   let nextRow xs =
         sequence $
         flip ListHT.mapAdjacent xs $
            \x0 x1 -> do
               op <-
                  case allowed of
                     Allowed True  True -> fmap toEnum $ randomR (0,1)
                     Allowed False True -> return Mul
                     Allowed _ False -> return Add
               return $ (,) op $
                  case op of
                     Add -> x0 + x1
                     Mul -> x0 * x1
   let go xs = do
         (ops0,ys0) <- fmap unzip $ nextRow xs
         fmap (mapPair ((ops0:),(ys0:))) $
            if null ys0
               then return ([],[])
               else go ys0
   let xs0 = BoxedArray.toList base
   let shape@(Shape.ZeroBased n) = BoxedArray.shape base
   let shape1 = Shape.ZeroBased (n-1)
   (ops,xs) <- go xs0
   return
      (BoxedArray.fromList (Shape.lowerTriangular shape1) $
         concat $ reverse ops,
       BoxedArray.fromList (Shape.lowerTriangular shape) $
         concat $ reverse (xs0:xs))

construct ::
   (Random.RandomGen g) =>
   SolutionCheck Integer ->
   Allowed -> Int -> Integer ->
   MS.State g
      (Array (Shape.LowerTriangular ShapeInt) Op,
       Array (Shape.LowerTriangular ShapeInt) (Integer,Bool))
construct check allowed n maxV = do
   let shape = Shape.ZeroBased n
   let triShape = Shape.lowerTriangular shape
   xs <- replicateM n $ randomR (0,maxV)
   (ops,pyr) <- pyramid allowed $ BoxedArray.fromList shape xs
   let go = do
         selected <-
            MS.evalStateT (replicateM n pick) $
            Set.fromList $ Shape.indices triShape
         let puzzle = map (\ij -> (ij, pyr!ij)) selected
         if check ops puzzle
            then return selected
            else go
   selected <- go
   return (ops,
      BoxedArray.zipWith (,) pyr $
      BoxedArray.fromAssociations False triShape $ map (flip (,) True) selected)


latexFromPuzzle ::
   String ->
   Either Int (Array (Shape.LowerTriangular ShapeInt) Op) ->
   Array (Shape.LowerTriangular ShapeInt) (Integer,Bool) ->
   [String]
latexFromPuzzle hidden mops xs =
   let n = either id sizeFromOps mops in
   printf "\\begin{picture}(%d,%d)" (2*n) n :
   map (\(i,j) -> printf "\\put(%d,%d){\\framebox(2,1){}}" (n-1-i + 2*j) (n-i))
      (Shape.indices $ Shape.lowerTriangular $ Shape.ZeroBased n) ++
   (BoxedArray.toAssociations xs >>= \((i,j),(x,display)) ->
      if null hidden && not display
         then []
         else
            let cell :: String
                cell =
                   if display
                      then printf "%d" x
                      else printf "\\%s{%d}" hidden x
            in [printf "\\put(%d,%d){\\makebox(2,1)[c]{%s}}"
                  (n-1-i + 2*j) (n-i) cell]) ++
   (case mops of
      Left _ -> []
      Right ops ->
         let half = 0.5 :: Double in
         BoxedArray.toAssociations ops >>= \((i,j),op) ->
            [printf "\\put(%d,%d){\\textcolor{white}{\\circle*{0.5}}}"
               (n-i + 2*j) (n-i),
             printf "\\put(%d,%d){\\circle{0.5}}" (n-i + 2*j) (n-i),
             printf "\\put(%.1f,%.1f){\\makebox(1,1)[c]{$%s$}}"
               (fromIntegral (n-i + 2*j) - half) (fromIntegral (n-i) - half)
               (case op of Add -> "+"; Mul -> "\\times{}")]) ++
   "\\end{picture}" :
   []

mainCreate ::
   (SolutionCheck Integer, (Allowed,Bool)) ->
   Int -> Int -> Integer -> String -> String -> IO ()
mainCreate (check,(allowed,displayOps)) n number maxV env hidden =
   putStr . unlines .
      concatMap
         ((if null env
            then id
            else (\pic ->
                     printf "\\begin{%s}" env : pic ++
                     printf "\\end{%s}" env : []))
            . uncurry (latexFromPuzzle hidden)
            . mapFst (if displayOps then Right else const (Left n))) .
      MS.evalState (replicateM number $ construct check allowed n maxV)
         =<< Random.initStdGen

commandCreate :: OP.Mod OP.CommandFields (IO ())
commandCreate =
   let parser =
         pure mainCreate
         <*>
            (
               (OP.flag'
                  (\ops xs -> LinAlg.solvable (sizeFromOps ops) (map fst xs),
                     (Allowed {allowedAdd = True, allowedMul = False}, False)) $
                  OP.long "allow-gaps" <>
                  OP.help "Employ both addition and multiplication")
               <|>
               (fmap ((,) Logic.solvableMixed) $
                  (OP.flag'
                        (Allowed {allowedAdd = True, allowedMul = True}, True) $
                     OP.long "mixed" <>
                     OP.help "Employ both addition and multiplication")
                  <|>
                  (OP.flag
                        (Allowed {allowedAdd = True, allowedMul = False}, False)
                        (Allowed {allowedAdd = False, allowedMul = True}, True) $
                     OP.long "multiplication" <>
                     OP.help "Employ multiplication only")
               )
            )
         <*>
            (OP.option
               (OP.eitherReader $
                  parseNumber "size" (\n -> 0<n && n<=1000)
                     "positive, below 1000") $
               OP.long "size" <>
               OP.metavar "NATURAL" <>
               OP.help "Width of the pyramid")
         <*>
            (OP.option
               (OP.eitherReader $
                  parseNumber "number" (\n -> 0<n && n<=1000000)
                     "positive, below 1000000") $
               OP.long "number" <>
               OP.value 1 <>
               OP.metavar "NATURAL" <>
               OP.help "Number of puzzles")
         <*>
            (OP.option
               (OP.eitherReader $
                  parseNumber "number" (\n -> 0<n && n<=1000000)
                     "positive, below 1000000") $
               OP.long "max-value" <>
               OP.value 10 <>
               OP.metavar "NATURAL" <>
               OP.help "Upper bound for values in the base line")
         <*>
            (OP.strOption $
               OP.long "environment" <>
               OP.value "" <>
               OP.metavar "NAME" <>
               OP.help "Custom LaTeX environment around pictures")
         <*>
            (OP.strOption $
               OP.long "hidden" <>
               OP.value "" <>
               OP.metavar "NAME" <>
               OP.help "Custom LaTeX command for hidden figures")
   in OP.command "create" $
      OP.info
         (OP.helper <*> parser)
         (OP.progDesc "create puzzle")


{-
solvable step-by-step (unique-logic):

1
1
3
16
122
1188
13844
185448
2781348

uniquely solvable (linear algebra):
1
1
3
17
149
1824
29001
573549
13604001
-}
mainCount :: (Int -> [(Int,Int)] -> Bool) -> IO ()
mainCount check =
   for_ [0..] $ \n ->
      print $ length $ filter (check n) $
      Combinatorics.tuples n $ Shape.indices $
      Shape.lowerTriangular $ Shape.ZeroBased n

commandCount :: OP.Mod OP.CommandFields (IO ())
commandCount =
   let parser =
         pure mainCount
         <*>
            (OP.flag Logic.solvable LinAlg.solvable $
               OP.long "allow-gaps" <>
               OP.help "Count puzzles that are uniquely solvable, but not stepwise")
   in OP.command "count" $
      OP.info
         (OP.helper <*> parser)
         (OP.progDesc "count solvable puzzles")


info :: OP.Parser a -> OP.ParserInfo a
info parser =
   OP.info
      (OP.helper <*> parser)
      (OP.fullDesc <> OP.progDesc "Sum pyramid aka Additionstreppe")

main :: IO ()
main =
   join $ OP.execParser $ info $
      OP.subparser $ commandCreate <> commandCount