{- |
In this approach I construct the board row by row from the bottom to the top.
In every step I maintain the necessary information
in order to know, what ships and positions and orientations
are allowed in the next row.
This information is stored in the Frontier.

possible optimization:
   "meet in the middle"
   compute counts for 5x10 boards and put them together,
   problem:
      for a given frontier there are many other half boards that may match
-}
module Combinatorics.Battleship.Count.ShortenShip where

import qualified Combinatorics.Battleship.Count.CountMap as CountMap
import qualified Combinatorics.Battleship.Count.Frontier as Frontier
import qualified Combinatorics.Battleship.Fleet as Fleet

import Combinatorics.Battleship.Count.CountMap (Count, )

import Data.Map (Map, )
import qualified Data.Map as Map

import qualified Data.Foldable as Fold
import Control.Monad (guard, zipWithM_, forM_, )
import Data.Monoid (mappend, mconcat, )
import Data.Foldable (foldMap, )

import Data.Function.HT (nest, )
import Data.List (zipWith4, intercalate, )
import Text.Printf (printf, )


type CountAll = CountMap.T
-- type Count = Integer
-- type CountAll = Map (Frontier.T, Fleet.T) Count


-- * count all possible fleets on a board with given width

baseCase :: Int -> CountAll
baseCase _size =
   CountMap.singleton (Frontier.empty, Fleet.empty) 1

{- |
In this approach, the fleet contains all ships
also the ones at the frontier.
-}
nextFrontier ::
   Int -> CountAll -> CountAll
nextFrontier size =
   mconcat .
   concatMap
      (\((frontier,fleet), cnt) ->
         map (flip CountMap.singleton cnt) $
         processFrontier size frontier 0 Frontier.empty fleet) .
   CountMap.toAscList

processFrontier ::
   Int -> Frontier.T ->
   Int -> Frontier.T ->
   Fleet.T ->
   [(Frontier.T, Fleet.T)]
processFrontier size oldFrontier =
   let go pos newFrontier newFleet =
          if pos>=size
            then [(newFrontier,newFleet)]
            else
              case Frontier.lookup oldFrontier pos of
                 Frontier.Blocked ->
                    go (pos+1) newFrontier newFleet
                 Frontier.Vertical n ->
                    go (pos+2) newFrontier newFleet
                    ++
                    (go (pos+2)
                        (Frontier.blockBounded size (pos+1) $
                         Frontier.blockBounded size (pos-1) $
                         Frontier.insertNew pos (Frontier.Vertical (n+1))
                         newFrontier)
                        (Fleet.inc (n+1) $ Fleet.dec n newFleet))
                 Frontier.Free ->
                    go (pos+1) newFrontier newFleet
                    ++
                    (go (pos+2)
                        (Frontier.blockBounded size (pos+1) $
                         Frontier.blockBounded size (pos-1) $
                         Frontier.insertNew pos (Frontier.Vertical 1)
                         newFrontier)
                        (Fleet.inc 1 newFleet))
                    ++
                    (concat $
                     zipWith4
                        (\newPos newFrontierUpdate shipSize _ ->
                            go newPos newFrontierUpdate
                               (Fleet.inc shipSize newFleet))
                        [pos+2 ..]
                        (tail $ tail $ tail $
                         scanl
                            (flip (Frontier.blockBounded size))
                            newFrontier [pos-1 ..])
                        [1 .. Fleet.maxSize] $
                     takeWhile id $
                     map (Frontier.isFree oldFrontier) [pos .. size-1])
   in  go


count :: (Int, Int) -> Fleet.T -> Count
count (height,width) reqFleet =
   sum $
   map snd $
   filter (\((_front,fleet), _) -> fleet == reqFleet) $
   CountMap.toAscList $
   nest height (nextFrontier width) $ baseCase width


-- * count fleets with an upper bound

{- |
Here we save memory and speed up the computation in the following way:
We stop searching deeper if

1. the fleet becomes larger than the requested fleet
    ("larger" means, that for at least one ship size
     the number of ships is larger than in the requested fleet)

2. the cumulated fleet becomes larger than the cumulated requested fleet
     This is necessary, since we do not know the final length
     of the vertical ships at the frontier.

In this approach,
the fleet does not contain the vertical ships at the frontier.
-}
nextFrontierBounded ::
   Int -> Fleet.T -> CountAll -> CountAll
nextFrontierBounded size maxFleet =
--   foldMap is not efficient enough
--   foldl mappend mempty .    -- not efficient enough
   mconcat .
   map
      (\((frontier,fleet), cnt) ->
         CountMap.fromMap $
         Map.fromListWith (+) $
         map
            (\(fr,fl) ->
               -- merge symmetric cases
               ((min fr (Frontier.reverse size fr), fl), cnt)) $
         processFrontierBounded size frontier maxFleet 0 Frontier.empty fleet) .
   CountMap.toAscList

nextFrontierBoundedExt ::
   Int -> Fleet.T -> CountAll -> IO CountAll
nextFrontierBoundedExt size maxFleet =
   CountMap.fromListExt .
   concatMap
      (\((frontier,fleet), cnt) ->
         map
            (\(fr,fl) ->
               -- merge symmetric cases
               ((min fr (Frontier.reverse size fr), fl), cnt)) $
         processFrontierBounded size frontier maxFleet 0 Frontier.empty fleet) .
   CountMap.toAscList

processFrontierBounded ::
   Int -> Frontier.T -> Fleet.T ->
   Int -> Frontier.T -> Fleet.T ->
   [(Frontier.T, Fleet.T)]
processFrontierBounded size oldFrontier maxFleet =
   let cumMaxFleet = Fleet.cumulate maxFleet
       guardCumulativeSubset frontier fleet =
          guard $ Fleet.subset
             (Fleet.cumulate $ addFrontierFleet frontier fleet)
             cumMaxFleet
       newShip shipSize frontier fleet =
          let newFleetUpdate = Fleet.inc shipSize fleet
          in  guard (Fleet.subset newFleetUpdate maxFleet)
              >>
              guardCumulativeSubset frontier newFleetUpdate
              >>
              return newFleetUpdate
       go pos newFrontier newFleet =
          if pos>=size
            then [(newFrontier,newFleet)]
            else
              let insertVertical n =
                     let newFrontierUpdate =
                            Frontier.blockBounded size (pos+1) $
                            Frontier.blockBounded size (pos-1) $
                            Frontier.insertNew pos (Frontier.Vertical n) newFrontier
                     in  guardCumulativeSubset newFrontierUpdate newFleet
                         >>
                         go (pos+2) newFrontierUpdate newFleet
              in  case Frontier.lookup oldFrontier pos of
                     Frontier.Blocked ->
                        go (pos+1) newFrontier newFleet
                     Frontier.Vertical n ->
                        (newShip n newFrontier newFleet
                         >>=
                         go (pos+2) newFrontier)
                        ++
                        insertVertical (n+1)
                     Frontier.Free ->
                        go (pos+1) newFrontier newFleet
                        ++
                        insertVertical 1
                        ++
                        (concat $
                         zipWith4
                            (\newPos shipSize newFrontierUpdate _ ->
                               newShip shipSize newFrontierUpdate newFleet
                               >>=
                               go newPos newFrontierUpdate)
                            [pos+2 ..]
                            [1 .. Fleet.maxSize]
                            (tail $ tail $ tail $
                             scanl
                                (flip (Frontier.blockBounded size))
                                newFrontier [pos-1 ..]) $
                         takeWhile id $
                         map (Frontier.isFree oldFrontier) [pos .. size-1])
   in  go


countBounded :: (Int, Int) -> Fleet.T -> Count
countBounded (height,width) reqFleet =
   countBoundedFromMap reqFleet $
   nest height (nextFrontierBounded width reqFleet) $ baseCase width


{- |
This solves a different problem.
In this variant the ships are allowed to touch each other.
-}
nextFrontierTouching ::
   Int -> Fleet.T -> CountAll -> CountAll
nextFrontierTouching size maxFleet =
   mconcat .
   map
      (\((frontier,fleet), cnt) ->
         CountMap.fromMap $
         Map.fromListWith (+) $
         map
--            (\key -> (key, cnt)) $
            (\(fr,fl) ->
               -- merge symmetric cases
               ((min fr (Frontier.reverse size fr), fl), cnt)) $
         processFrontierTouching size frontier maxFleet 0 Frontier.empty fleet) .
   CountMap.toAscList

nextFrontierTouchingExt ::
   Int -> Fleet.T -> CountAll -> IO CountAll
nextFrontierTouchingExt size maxFleet =
   CountMap.fromListExt .
   concatMap
      (\((frontier,fleet), cnt) ->
         map
            (\(fr,fl) ->
               -- merge symmetric cases
               ((min fr (Frontier.reverse size fr), fl), cnt)) $
         processFrontierTouching size frontier maxFleet 0 Frontier.empty fleet) .
   CountMap.toAscList

processFrontierTouching ::
   Int -> Frontier.T -> Fleet.T ->
   Int -> Frontier.T -> Fleet.T ->
   [(Frontier.T, Fleet.T)]
processFrontierTouching size oldFrontier maxFleet =
   let cumMaxFleet = Fleet.cumulate maxFleet
       guardCumulativeSubset frontier fleet =
          guard $ Fleet.subset
             (Fleet.cumulate $ addFrontierFleet frontier fleet)
             cumMaxFleet
       newShip shipSize frontier fleet =
          let newFleetUpdate = Fleet.inc shipSize fleet
          in  guard (Fleet.subset newFleetUpdate maxFleet)
              >>
              guardCumulativeSubset frontier newFleetUpdate
              >>
              return newFleetUpdate
       insertVertical cont n pos frontier fleet =
          let newFrontier =
                 Frontier.insertNew pos (Frontier.Vertical n) frontier
          in  guardCumulativeSubset newFrontier fleet
              >>
              cont (pos+1) newFrontier fleet

       finishVerticals pos newFrontier newFleet =
          if pos>=size
            then startNewShips 0 newFrontier newFleet
            else
              case Frontier.lookup oldFrontier pos of
                 Frontier.Blocked ->
                    error "in touching mode there must be no blocked fields"
                 Frontier.Vertical n ->
                    insertVertical finishVerticals
                       (n+1) pos newFrontier newFleet
                    ++
                    (newShip n newFrontier newFleet
                     >>=
                     finishVerticals (pos+1) newFrontier)
                 Frontier.Free ->
                    finishVerticals (pos+1) newFrontier newFleet

       startNewShips pos newFrontier newFleet =
          if pos>=size
            then [(newFrontier,newFleet)]
            else
              case Frontier.lookup newFrontier pos of
                 Frontier.Blocked ->
                    error "finishVerticals must not block fields"
                 Frontier.Vertical _ ->
                    startNewShips (pos+1) newFrontier newFleet
                 Frontier.Free ->
                    startNewShips (pos+1) newFrontier newFleet
                    ++
                    insertVertical startNewShips 1 pos newFrontier newFleet
                    ++
                    (concat $
                     zipWith
                        (\shipSize _ ->
                           newShip shipSize newFrontier newFleet
                           >>=
                           startNewShips (pos+shipSize) newFrontier)
                        [1 .. Fleet.maxSize] $
                     takeWhile id $
                     map (Frontier.isFree newFrontier) [pos .. size-1])
   in  finishVerticals

countTouching :: (Int, Int) -> Fleet.T -> Count
countTouching (height,width) reqFleet =
   countBoundedFromMap reqFleet $
   nest height (nextFrontierTouching width reqFleet) $ baseCase width


fleetAtFrontier :: Frontier.T -> Fleet.T
fleetAtFrontier =
   foldMap
      (\use ->
         case use of
            Frontier.Vertical n -> Fleet.singleton n 1
            _ -> Fleet.empty) .
   Frontier.toList


addFrontierFleet :: Frontier.T -> Fleet.T -> Fleet.T
addFrontierFleet frontier fleet =
   mappend fleet $ fleetAtFrontier frontier


-- * retrieve counts from count maps

countBoundedFromMap :: Fleet.T -> CountAll -> Count
countBoundedFromMap reqFleet =
   sum .
   map snd .
   filter (\((front,fleet), _) ->
             addFrontierFleet front fleet == reqFleet) .
   CountMap.toAscList

countBoundedFleetsFromMap :: CountAll -> Map Fleet.T Integer
countBoundedFleetsFromMap =
   Map.fromListWith (+) .
   map (\((front,fleet), cnt) ->
             (addFrontierFleet front fleet,
              fromIntegral cnt)) .
   CountMap.toAscList

{-
maybe this is not lazy enough and thus requires to much memory at once
-}
countBoundedFleetsFromMap_ :: CountAll -> Map Fleet.T Integer
countBoundedFleetsFromMap_ =
   Map.mapKeysWith (+) (uncurry addFrontierFleet) .
   fmap fromIntegral .
   CountMap.toMap


standardFleet :: Fleet.T
standardFleet = Fleet.fromList [(5,1), (4,2), (3,3), (2,4)]


{-
*ShortenShip> let height=3::Int; width=10::Int; reqFleet = Fleet.fromList [(2,3),(3,1)]
(0.01 secs, 524480 bytes)

*ShortenShip> let counts = nest height (nextFrontier width) $ baseCase width in (Map.size counts, Fold.sum counts, Fold.maximum counts)
(658486,37986080,16640)
(77.32 secs, 9147062872 bytes)

*ShortenShip> let counts = nest height (nextFrontierBounded width reqFleet) $ baseCase width in (Map.size counts, Fold.sum counts, Fold.maximum counts)
(59485,870317,2295)
(41.05 secs, 4961028184 bytes)

This was computed, where we marked horizontal ships
instead of blocked columns.

*ShortenShip> let width=10::Int; reqFleet = Fleet.fromList [(5,1), (4,2), (3,3), (2,4)]
*ShortenShip> map Map.size $ iterate (nextFrontierBounded width reqFleet) $ baseCase width
[1,976,9441,129247,727781,Interrupted.

Here we switched to blocked columns and thus could merge some cases.
*ShortenShip> map Map.size $ iterate (nextFrontierBounded width reqFleet) $ baseCase width
[1,762,8712,110276,671283,Heap exhausted

Now merge symmetric cases.
*ShortenShip> map Map.size $ iterate (nextFrontierBounded width reqFleet) $ baseCase width
[1,400,4209,53897,331185,Heap exhausted

Now correctly stop searching, when we exceed the requested fleet
in a cumulative way.
*ShortenShip> map Map.size $ iterate (nextFrontierBounded width reqFleet) $ baseCase width
[1,400,2780,33861,156962,596354,1078596,
-}


countSingleKind :: IO ()
countSingleKind =
   mapM_
      (print . countBounded (10,10) . Fleet.fromList . (:[]))
      [(5,1), (4,2), (3,3), (2,4)]

{- | <http://math.stackexchange.com/questions/58769/how-many-ways-can-we-place-these-ships-on-this-board>
-}
count8x8 :: IO ()
count8x8 =
{-
   print $ countTouching (8,8) $
   Fleet.fromList [(2,1), (3,2), (4,1), (5,1)]
-}
   let reqFleet = Fleet.fromList [(2,1), (3,2), (4,1), (5,1)]
       width = 8
       height = 8
   in  do CountMap.writeFile (printf tmpPath (0::Int)) $
             baseCase width
          forM_ [0..height::Int] $ \n -> do
             print . countBoundedFromMap reqFleet
                =<< CountMap.readFile (printf tmpPath n)
             CountMap.writeFile (printf tmpPath (n+1))
                =<< nextFrontierTouchingExt width reqFleet
                =<< CountMap.readFile (printf tmpPath n)
{-
0
0
0
24348
712180
8705828
50637316
193553688
571126760
-}

{- |
http://mathoverflow.net/questions/8374/battleship-permutations
-}
count10x10 :: IO ()
count10x10 =
   print $ countBounded (10,10) $
   Fleet.fromList [(2,1), (3,2), (4,1), (5,1)]

{-
width = 10
reqFleet = Fleet.fromList [(5,1), (4,1), (3,2), (2,1)]

0 (height 0)
0
0
28
3216
665992
7459236
49267288
212572080
703662748
1925751392 (height 10)
4558265312
9655606528
-}


countStandard :: IO ()
countStandard =
   let -- reqFleet = Fleet.fromList [(5,1), (4,2), (3,3), (2,4)]
       reqFleet = Fleet.fromList [(2,1), (3,2), (4,1), (5,1)]
       -- reqFleet = Fleet.fromList [(5,3), (3,3), (2,4)]
       -- reqFleet = Fleet.fromList [(5,1), (4,5), (2,4)]
       -- reqFleet = Fleet.fromList [(5,1), (4,2), (3,7)]
       -- reqFleet = Fleet.fromList [(5,1), (4,2), (3,3)]
       width = 10
       height = 12
   in  mapM_ (print . countBoundedFromMap reqFleet) $
       take (height+1) $
       iterate (nextFrontierBounded width reqFleet) $
       baseCase width

{-
width = 8

0 (height 0)
0
0
0
0
0
0
0
0
41590204
7638426604 (height 10)
362492015926
7519320122520
-}

{-
width = 9

0 (height 0)
0
0
0
0
0
0
3436
41590204 (height 8)
14057667720
810429191552
19372254431062
259204457356150 (height 12)
-}

tmpPath :: FilePath
tmpPath = "/tmp/battleship%02d"


writeTmps :: IO ()
writeTmps =
   let width = 10
   in  zipWithM_
          (\n -> CountMap.writeFile (printf tmpPath n)) [0::Int ..] $
       iterate (nextFrontierBounded width standardFleet) $
       baseCase width


countExt :: IO ()
countExt =
   let width = 10
       height = 12
   in  do CountMap.writeFile (printf tmpPath (0::Int)) $
             baseCase width
          forM_ [0..height::Int] $ \n -> do
             print . countBoundedFromMap standardFleet
                =<< CountMap.readFile (printf tmpPath n)
             CountMap.writeFile (printf tmpPath (n+1))
                =<< nextFrontierBoundedExt width standardFleet
                =<< CountMap.readFile (printf tmpPath n)

{-
width = 10

0 (height 6)
13662566
7638426604
810429191552
26509655816984
430299058359872
4354180199012068
-}

countFleets :: IO ()
countFleets =
   Fold.mapM_ putStrLn .
   Map.mapWithKey
      (\fleet cnt ->
         "|-\n| " ++
         intercalate " || "
            (map ((\n -> if n==0 then "" else show n) . Fleet.lookup fleet) [2..5]
              ++ [show cnt])) .
   Map.filterWithKey (\fleet _cnt -> Fleet.subset fleet standardFleet) .
   countBoundedFleetsFromMap =<<
   CountMap.readFile (printf tmpPath (10::Int))


printMapSizes :: IO ()
printMapSizes =
   mapM_ (print . CountMap.size) $
   iterate (nextFrontierBounded 10 standardFleet) $
   baseCase 10


main :: IO ()
main = count8x8