{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
module WaveFunctionCollapse where

import Data.Function ((&))
import Data.Hashable (Hashable)
import Data.Holmes
import Data.JoinSemilattice.Intersect (fromList, toList)
import Data.List (transpose)
import Data.List.Split (chunksOf)
import Data.Maybe (isJust, mapMaybe)
import Data.Propagator (lift)
import GHC.Generics (Generic)
import Relude ((!!?))
import Test.Hspec (Spec, it, shouldBe)

-- Wave function collapse* is an algorithm that works by placing constraints
-- between each cell and their neighbours. A cell is randomly specialised to a
-- particular value, and the effects ripple out via the constraints. Then,
-- another cell is specialised, and the process repeats until all cells are
-- specialised.
--
-- It turns out that this is actually just a special case of the propagator
-- idea, and specifically the `Intersect` strategy. While we're not going to
-- implement the full algorithm here, we'll demonstrate the idea with a
-- simplified version in order to draw some desert island maps!
--
-- * https://github.com/mxgmn/WaveFunctionCollapse

--------------------------------------------------

-- First, we'll start with a type to specify the possible terrain types in our
-- map:
data Tile = Water | Sand | Grass | Tree
  deriving stock (Eq, Ord, Bounded, Enum, Generic)
  deriving anyclass (Hashable)

instance Show Tile where
  show = \case
    Water -> "💦"
    Sand  -> "🔅"
    Grass -> "🍀"
    Tree  -> "🌲"

-- Now, we'll specify some constraints on our neighbours. Again, this is a very
-- simplified version of the WaveFunctionCollapse concept - typically, we'd
-- have far more "tiles", and neighbours would be chosen by properties attached
-- to each edge of each tile.

surroundings :: Tile -> Intersect Tile
surroundings = fromList . \case

  -- A tree must be entirely surrounded by grass. Two trees cannot touch, and
  -- trees cannot be on beaches or in water.
  Tree -> [ Grass ]

  -- The only thing that can neighbour water is more water or sand. This means
  -- that every island has a beach, and we might even get some small islands
  -- out in water, too!
  Water -> [ Sand, Water ]

  -- Sand must sit between water and grass. Note that this simple system
  -- doesn't prevent random sand tiles amid grass; we'd need to specify the
  -- constraints in a more comprehensive way to mitigate this.
  Sand -> [ Sand, Water, Grass ]

  -- Grass can neighbour sand, more grass, or trees!
  Grass -> [ Sand, Tree, Grass ]

-- Get the neighbours of a cell at a given index.
neighbours :: Int -> [ x ] -> [ x ]
neighbours index board = mapMaybe (board !!?)
  [ index - 21, index - 20, index - 19
  , index -  1, {- HOME! -} index +  1
  , index + 19, index + 20, index + 21
  ]

-- The 20 × 20 board makes up 400 tiles.
tiles :: Config Holmes (Intersect Tile)
tiles = shuffle (400 `from` [ Water, Sand, Grass, Tree ])

--------------------------------------------------

maps :: IO (Maybe [ Intersect Tile ])
maps = do
  tiles `satisfying` \board@(chunksOf 20 -> rows) -> do
    let columns = transpose rows

    and'
      [ -- As we're trying to draw an island, we'll surround the whole map with
        -- water:
        all' (.== lift Water) (head rows)
      , all' (.== lift Water) (last rows)

      , all' (.== lift Water) (head columns)
      , all' (.== lift Water) (last columns)

        -- To generate more interesting maps, we'll require that every valid
        -- map contains at least one tree (and thus has at least one 5 × 5
        -- island).
      , any' (.== lift Tree) board

        -- For each tile, find the valid surrounding tiles, then constraint its
        -- neighbours to those possibilities.
      , board & allWithIndex' \index tile -> do
          let candidates = tile .>>= surroundings
          all' (.== candidates) (neighbours index board)
      ]

-- If you want to see some of the generated maps, run `cabal new-repl examples`
-- and use the following function to print out a result:
--
-- > import WaveFunctionCollapse
-- > Just example <- maps
-- > printMap example

printMap :: [ Intersect Tile ] -> IO ()
printMap (chunksOf 20 -> rows) = mapM_ printRow rows
  where printRow = putStrLn . foldMap (show . head . toList)

-- Use `cabal new-test examples` to run these tests and check for correct
-- solutions.

spec_wfc :: Spec
spec_wfc = it "generates a map" do
  maps >>= \result -> isJust result `shouldBe` True