-----------------------------
-- Nonogram Solving Module --
-----------------------------

module Nonogram
(
    Grid,
    Line,
    Cell(..),

    PuzzleDefinition,
    LineDefinition,
    Strip,

    solvePuzzle
) where

--
-- Imports
--

import Data.List
import Control.Exception

--
-- Types
--

-- A puzzle definition is consisted of a list of row definitions and column definitions.
-- Note: The length of the rows should be equal to the amount of columns and vice versa.
type PuzzleDefinition = ([LineDefinition], [LineDefinition])

-- A line defintion is consited of a list of strips of black cells.
type LineDefinition = [Strip]

-- A strip defines the length of the sequence of consequtive black cells in a line.
-- Note: Strips should be positive numbers.
--       There shouldn't be strips of zero length.
type Strip = Int

-- A grid is consisted of a list of rows.
type Grid = [Line]

-- A line is consisted of a list of cells.
type Line = [Cell]

-- Cell is a data type that defines the state of a single cell.
-- Note: The Unknown state could be represented differently,
--       perhaps as a non-deterministic computation (i.e. [White, Black]).
--       This way maybe some of the functions could be expressed more elegantly.
data Cell = Unknown
          | Black
          | White
          deriving (Eq, Show)

--
-- Functions
--

-- This function takes a puzzle definition and returns the solution grid to that puzzle.
-- If the puzzle cannot be solved without guesses, the resulting grid will contain Unknown cells.
solvePuzzle :: PuzzleDefinition -> Grid
solvePuzzle puzzleDef =
    let (rows, cols) = puzzleDef
        width = length cols
        hight = length rows
        grid = unknownGrid width hight
    in solveGrid puzzleDef grid

-- This function takes a grid size (width and hight), and returns a grid filled with Unknown cells.
unknownGrid :: Int -> Int -> Grid
unknownGrid width hight = replicate hight (unknownLine width)

-- This function takes a line length, and returns a line filled with Unknown cells.
unknownLine :: Int -> Line
unknownLine lineLength = replicate lineLength Unknown

--TODO: refactor?
-- This function takes a puzzle definition and a grid, and returns the solved grid.
-- If the grid cannot be solved without guesses, some cells in the solution will be Unknown.
solveGrid :: PuzzleDefinition -> Grid -> Grid
solveGrid puzzleDef rows =
    let (rowDefs, colDefs) = puzzleDef
        -- Solve grid using row definitions
        solvedRows = zipWith solveLine rowDefs rows
        -- Transpose grid to get the columns
        cols = transpose solvedRows
        -- Solve grid using column definitions
        solvedCols = zipWith solveLine colDefs cols
        -- Transpose grid again to get the rows
        newRows = transpose solvedCols
    in -- If solving the columns didn't improve the solution,
       -- Then we solved the grid as much as possible (without guesses).
       -- Else, continue solving the grid based on the information we obtained
       -- at this step.
       if newRows == solvedRows
       then newRows
       else solveGrid puzzleDef newRows

-- This function takes a line definition and a line, and returns the solved line.
-- If the line cannot be completely solved, some cells will remain Unknown.
-- If the line has not possible solution, an error is produced.
solveLine :: LineDefinition -> Line -> Line
solveLine lineDef line =
    let solutions = (lineSolutions lineDef line)
        mergedSolution = foldr1 mergeSolutions solutions
    in if null solutions
       then error "Line has no possible solution"
       else mergedSolution

-- This function takes a line definition and a line, and returns all the possible
-- solutions to that line.
lineSolutions :: LineDefinition -> Line -> [Line]
lineSolutions lineDef line = filter (isSolution lineDef) (linePermutaions line)

-- This function takes a line definition and a line, and checks whether the line
-- is a solution to the line definition.
isSolution :: LineDefinition -> Line -> Bool
isSolution strips cells =
    let blocks = group cells
        blackBlocks = filter ((==Black) . head) blocks
        blackBlocksLengths = map length blackBlocks
    in blackBlocksLengths == strips

-- This function takes a line, and returns all possible permutations of that line.
-- I.E. for every Unknown cell, it produces a line with that cell White, and a
-- line with that cell Black.
linePermutaions :: Line -> [Line]
linePermutaions (Unknown:cells) =
    let otherLinePermutations = linePermutaions cells
        unknownIsBlackPremutations = map (Black:) otherLinePermutations
        unknownIsWhitePremutations = map (White:) otherLinePermutations
    in unknownIsBlackPremutations ++ unknownIsWhitePremutations
linePermutaions (knownCell:cells) = map (knownCell:) (linePermutaions cells)
linePermutaions [] = [[]]

-- This function takes two lines (that are possible solutions to a line definition),
-- and returns a line that is a merge of these solutions.
-- Merging solutions keeps cells that are equal in both solutions.
-- Unknown cells are placed instead of cells that are not equal.
mergeSolutions :: Line -> Line -> Line
mergeSolutions a b = zipWith keepCellIfEqual a b

-- This function takes two cells, and returns the first if they are equal,
-- otherwise it returns an Unknown cell.
keepCellIfEqual :: Cell -> Cell -> Cell
keepCellIfEqual a b =
    if a == b
    then a
    else Unknown