{- |
This module improves "Nonogram.Encoding.Naive" in these ways:

* It provides distinguished elements for black and white squares.
  Thus the solver can conclude like so:
  \"If there is no possibility to make a square white, it must be black.\"

* Single squares can be set to black or white.
  This is done when synchronizing horizontal and vertical slices
  in 'Base.assignsFromPositions'.

* The search ranges for brick positions are narrowed
  to the positions that can really occur.

* The search ranges for white squares are narrowed accordingly.

* The left-most and the right-most brick in each strip
  is combined with the space to the left and right border, respectively.
-}
module Nonogram.Encoding.BlackWhite
         (assigns, assignsBW, bitAssigns, bitVectorAssigns) where

import qualified Nonogram.Base as Base
import Nonogram.Base
         (Strip(Strip), strip, BrickId(BrickId),
          Orientation(Horizontal, Vertical), Color(White, Black), noAssign)

import qualified Math.SetCover.BitSet as BitSet
import qualified Math.SetCover.Exact as ESC
import Data.Bits (bit)

import qualified Data.Map as Map; import Data.Map (Map)
import qualified Data.Set as Set; import Data.Set (Set)
import qualified Data.Monoid.HT as Mn
import qualified Data.List.Match as Match
import qualified Data.List as List
import Data.Foldable (foldMap, fold)
import Data.Monoid (Monoid)


data Item = Brick BrickId | Position Int Color | Reserve BrickId Int
   deriving (Eq, Ord, Show)

instance Base.Position Item where position = Position


type Assign map = ESC.Assign map (Map Strip (Set Item))

{-
For efficiency reasons combine the left-most and right-most brick
with the space to the left and right border, respectively.
-}
assignsFromBrick ::
   (Monoid map) =>
   Orientation -> Int -> Int -> BrickId -> BrickId ->
   Int -> Int -> Int -> [Assign map]
assignsFromBrick orient width line lastBrick brick
      leftBorder rightBorder size =
   flip map [leftBorder .. width-rightBorder] $ \col ->
   noAssign $ strip orient line $
   Brick brick
   :
   (map (flip Position White) $ take size [col ..])
   ++
   (if brick > BrickId 0
      then map (Reserve $ pred brick) (takeWhile (<width) [col .. ])
      else map (flip Position Black) (take col [0 ..]))
   ++
   (if brick < lastBrick
      then map (Reserve brick) (takeWhile (<width) [0 .. col+size])
           ++
           Mn.when (col+size<width) [Position (col+size) Black]
      else map (flip Position Black) (takeWhile (<width) [col+size ..]))

assignsFromLine ::
   (Monoid map) => Orientation -> Int -> Int -> [Int] -> [Assign map]
assignsFromLine orient width line xs =
   let bricks = Match.take xs [BrickId 0 ..]
   in  concat
         (List.zipWith4
            (assignsFromBrick orient width line (maximum bricks))
            bricks
            (scanl (+) 0 $ map succ xs)
            (scanr (+) (-1) $ map succ xs)
            xs)
       ++
       Mn.when (null xs)
         [noAssign . strip orient line .
          {- The Reserve item makes sure, that this line of squares is taken
             and that it is not assembled from the squares generated by
             Base.assignsFromPositions. -}
          (Reserve (BrickId 0) 0 :) .
          map (flip Position Black) $
          take width [0..]]
       ++
       do
         (_,brick,left,right) <-
            List.zip4 (drop 1 xs) bricks
               (drop 1 $ scanl (+) 0 $ map succ xs)
               (drop 1 $ scanr (+) 0 $ map succ xs)
         c <- [left .. width-right]
         return $ noAssign $
            strip orient line [Reserve brick c, Position c Black]

assignsGen ::
   (Monoid map) =>
   (Int -> Int -> Color -> map) -> [[Int]] -> [[Int]] -> [Assign map]
assignsGen square rows columns =
   concat (zipWith (assignsFromLine Horizontal (length columns)) [0..] rows)
   ++
   concat (zipWith (assignsFromLine Vertical (length rows)) [0..] columns)
   ++
   Base.assignsFromPositions square rows columns

assigns :: [[Int]] -> [[Int]] -> [Assign (Set (Int,Int))]
assigns = assignsGen Base.square

assignsBW :: [[Int]] -> [[Int]] -> [Assign (Map (Int,Int) Color)]
assignsBW = assignsGen Base.squareBW


type BitVector = BitSet.Set Integer

bitAssigns ::
   [[Int]] -> [[Int]] ->
   [ESC.Assign map (Map Strip (Set Item))] ->
   [ESC.Assign map (Map Strip BitVector)]
bitAssigns rows columns =
   let m =
         (fmap ((,) (length columns)) $ fmap length $
          Map.fromList $ zip (map (Strip Horizontal) [0..]) rows)
         `Map.union`
         (fmap ((,) (length rows)) $ fmap length $
          Map.fromList $ zip (map (Strip Vertical) [0..]) columns)
   in  map (fmap (Map.intersectionWith (foldMap . bitFromItem) m))

bitFromItem :: (Int,Int) -> Item -> BitVector
bitFromItem (width, numBricks) x =
   BitSet.Set $ bit $
   case x of
      Position k color -> 2*k + fromEnum color
      Reserve (BrickId brick) k -> (2+brick)*width + k
      Brick (BrickId brick) -> (2*numBricks)*width + brick

bitVectorAssigns ::
   [ESC.Assign map (Map Strip (Set Item))] -> [ESC.Assign map BitVector]
bitVectorAssigns =
   ESC.bitVectorFromSetAssigns .
   map (fmap (fold . Map.mapWithKey (\str -> Set.map ((,) str))))