{-# OPTIONS -O2 -optc-O3 #-}
{-# LANGUAGE BangPatterns #-}

-- The Computer Language Benchmarks Game
--   http://shootout.alioth.debian.org/
--
--   Sterling Clover's translation of Tim Hochberg's Clean implementation

module Main where

import System.Environment (getArgs)
import Data.Bits (setBit, testBit, shiftL, shiftR, (.&.), (.|.))
import Data.List (partition, intersperse, find, delete)
import Data.Array.Unboxed (UArray, Array, array, (!))
import Control.Arrow (first, (&&&))
import Prelude hiding (flip)


--- The Board ---
n_elem, n_col, n_row :: Int
n_elem = 5
n_col = 5
n_row = 10

m_top :: Mask
m_top = 0x1F

cells :: [Cell]
cells = [0..49]

colors :: [Color]
colors = [0..9]

cellAt :: Int -> Int -> Int
cellAt x y = x + n_col * y

coordOf :: Int -> (Int, Int)
coordOf i = snd &&& fst $ i `quotRem` n_col

isValid :: Int -> Col -> Bool
isValid x y = 0 <= x && x < n_col && 0 <= y && y < n_row

--- Piece Operations ---
data Direction = E | SE | SW | W | NW | NE deriving (Enum, Eq, Ord)
type Piece = [Direction]
type CellCoord = (Int, Int)
type Mask = Int; type Color = Int; type Row = Int;
type Col = Int; type Tag = Int; type Cell = Int
type Solution = [Mask]

pieces :: Array Int Piece
pieces = array (0,9) $ zip [0..9] $
         [[E,  E,  E,  SE],
          [SE, SW, W,  SW],
          [W,  W,  SW, SE],
          [E,  E,  SW, SE],
          [NW, W,  NW, SE, SW],
          [E,  E,  NE, W],
          [NW, NE, NE, W],
          [NE, SE, E,  NE],
          [SE, SE, E,  SE],
          [E,  NW, NW, NW]]

permutations :: Piece -> [Piece]
permutations = take 12 . perms
    where
      perms p = p : flip p : perms (rotate p)
      rotate piece = map r piece
          where r E  = NE
                r NE = NW
                r NW = W
                r W  = SW
                r SW = SE
                r SE = E
      flip piece = map f piece
          where f E  = W
                f NE = NW
                f NW = NE
                f W  = E
                f SW = SE
                f SE = SW

--- Mask Operations ----
untag :: Mask -> Mask
untag mask   = mask .&. 0x1ffffff

retag :: Mask -> Tag -> Mask
retag mask n = untag mask .|. n `shiftL` 25

tagof :: Mask -> Tag
tagof mask   = mask `shiftR` 25

tag :: Mask -> Tag -> Mask
tag   mask n = mask .|. n `shiftL` 25

count1s :: Mask -> Int
count1s i
    | i == 0 = 0
    | i .&. 1 == 1 = 1 + count1s (i `shiftR` 1)
    | otherwise = count1s (i `shiftR` 1)

first0 :: Mask -> Int
first0 i
    | i .&. 1 == 0 = 0
    | otherwise = 1 + first0 (i `shiftR` 1)

--- Making the Bitmasks ---
mod2 :: Int -> Int
mod2 x = x .&. 1

packSize :: Num a => a -> a -> a
packSize a b = a*5+b

unpackSize :: Integral a => a -> (a, a)
unpackSize n = quotRem n 5

move :: Direction -> CellCoord -> CellCoord
move E  (x, y) = (x+1, y)
move W  (x, y) = (x-1, y)
move NE (x, y) = (x+(mod2 y),   y-1)
move NW (x, y) = (x+(mod2 y)-1, y-1)
move SE (x, y) = (x+(mod2 y),   y+1)
move SW (x, y) = (x+(mod2 y)-1, y+1)

pieceBounds :: Piece -> Bool -> (Int, Int, Int, Int)
pieceBounds piece isodd = bnds piece 0 y0 0 y0 0 y0
  where
    y0 | isodd = 1 | otherwise = 0
    bnds [] _ _ xmin ymin xmax ymax = (xmin, ymin, xmax, ymax)
    bnds (d:rest) x y xmin ymin xmax ymax =
        bnds rest x' y' (min x' xmin) (min y' ymin) (max x' xmax) (max y' ymax)
            where (x', y') = move d (x, y)

pieceMask :: Piece -> (Mask, Mask)
pieceMask piece
    | odd y1    = (tag (msk piece x2 y2 0) (packSize w2 h2),
                   tag (msk piece x1 (y1+1) 0 `shiftR` n_col) (packSize w1 h1))
    | otherwise = (tag (msk piece x1 y1 0) (packSize w1 h1),
                   tag (msk piece x2 (y2+1) 0 `shiftR` n_col) (packSize w2 h2))
    where
      (xmin, ymin, xmax, ymax) = pieceBounds piece False
      (x1, y1) = (-xmin, -ymin)
      w1 = xmax - xmin
      h1 = ymax - ymin
      (xmin', ymin', xmax', ymax') = pieceBounds piece True
      (x2, y2) = (-xmin', (-ymin')+1)
      w2 = xmax' - xmin'
      h2 = ymax' - ymin'
      msk :: Piece -> Col -> Row -> Mask -> Mask
      msk [] x y m = m `setBit` cellAt x y
      msk (d:rest) x y m = msk rest x' y' (m `setBit` cellAt x y)
          where (x', y') = move d (x, y)

templatesForColor :: Color -> ([Mask], [Mask])
templatesForColor c = (unzip . map pieceMask) perms
    where perms | c == 5 = take 6 ps | otherwise = ps
          ps = Main.permutations $ pieces ! c

--- Looking for Islands ---
noLineIslands :: Mask -> Cell -> Cell -> Int -> Bool
noLineIslands mask start stop step
    | (fnd testBit . fnd ((not .) . testBit) . fnd testBit)  start > stop  = True
    | otherwise = False
  where
    fnd test !x
        | x >= 25     = 25
        | test mask x = x
        | otherwise   = fnd test (x+step)

noLeftIslands, noRightIslands :: Mask -> Bool
noLeftIslands  mask  = noLineIslands mask 0 20 5
noRightIslands mask  = noLineIslands mask 4 24 5

noIslands :: Mask -> Bool
noIslands board = noisles board (count1s board)

noisles :: Mask -> Int -> Bool
noisles _ 30 = True
noisles board ones
    | (ones' - ones) `rem` n_elem /= 0 = False
    | otherwise = noisles board' ones'
    where board' = fill board (coordOf (first0 board))
          ones' = count1s board'

fill :: Mask -> CellCoord -> Mask
fill m cc@(x, y)
    | x < 0 || x >= n_col = m
    | y < 0 || y >= 6     = m
    | testBit m i = m
    | otherwise = foldl (\mi d -> fill mi (move d cc)) (setBit m i)
                  [E, NE, NW, W, SW, SE]
    where i = cellAt x y

--- More Mask Generation ---
masksForColor :: Color -> [(Row, Mask)]
masksForColor c = concatMap atCell cells
  where
    (evens, odds) = templatesForColor c
    atCell n =
        if even y
          then [(y, retag (m `shiftL` x) c) | m <- evens , isok m x y]
          else [(y, retag (m `shiftL` x) c) | m <- odds  , isok m x y]
        where (x, y) = coordOf n

isok :: Mask -> Row -> Col -> Bool
isok mask x y =
    isValid (x+width) (y+height) &&
            case (y == 0, y+height==9) of
              (False, False) -> noLeftIslands mask' && noRightIslands mask'
              (False, True)  -> noIslands (mask' `shiftL` (n_col * (y - 4)))
              (True, _ ) -> noIslands mask'
    where (width, height) = unpackSize (tagof mask)
          mask' = untag mask `shiftL` x

masksAtCell :: Array (Row,Col) (Array Color [Mask])
masksAtCell = trps $ map (masksAt cells . masksForColor) colors

masksAt :: [Int] -> [(Row,Mask)]-> [[Mask]]
masksAt [] _ = []
masksAt (n:ns) !masks = map snd t : masksAt ns f
    where
      (t, f) = partition test masks
      test (r, m) = n' >= 0 && n' < 25 &&  m `testBit` n'
          where n' = n - (n_col * r)

trps :: [[masks]] -> Array (Row, Col) (Array Color masks)
trps !a = array ((0,0),(9,4)) $ concatMap (uncurry (map . first . (,))) $
          zip [0..9] [copy !! y | y <- [1,0,1,0,1,2,3,4,5,6]]
    where
      copy = [ [(x,copy' (cellAt x y)) | x <- [0..n_col-1]] |
               y <- [1,2,5,6,7,8,9]]
      copy' cell = array (0,9) $ map (\clr -> (clr,a !! clr !! cell)) colors

--- Formatting ---
format :: Bool -> String -> String
format _ [] = ""
format isodd chars | isodd = " " ++ str | otherwise = str
        where
          (cur, rest) = splitAt 5 chars
          str =  intersperse ' ' cur ++ " \n" ++ format (not isodd) rest

toString :: Solution -> String
toString !masks = map color cells
    where
      masksWithRows = withRows 0 0 (reverse masks)
      withRows _ _ [] = []
      withRows board r (m:rest) = (r', m) : withRows board' r' rest
          where delta = first0 board `quot` n_col
                board' = board `shiftR`  (delta * n_col) .|. untag m
                r' = r+delta
      color n = maybe '.' (("0123456789" !!) . tagof . snd)
                (find matches masksWithRows)
          where
            matches (r, m)
              | n' < 0 || n' > 30  = False
              | otherwise  = (untag m) `testBit` n'
              where n' = n - (n_col * r)

--- Generate the solutions ---
firstZero :: UArray Int Int
firstZero = array (0,31) $ zip [0..31]
            [0,1,0,2,0,1,0,3,0,1,0,2,0,1,0,4,0,1,0,2,0,1,0,3,0,1,0,2,0,1,0,5]

solutions :: [String]
solutions = solveCell 0 colors 0 [] []

solveCell :: Row -> [Color] -> Mask -> Solution -> [String] -> [String]
solveCell _ [] _board soln results = let s = toString soln
                                    in  s:(reverse s):results
solveCell !row !todo !board !soln results
    | top/=m_top = foldr solveMask results
                   [(m, c) | c <- todo, m  <- masks ! c,  board .&. m == 0]
    | otherwise  = solveCell (row+1) todo (board `shiftR` n_col) soln results
    where top = board .&. m_top
          masks = masksAtCell ! (row, (firstZero ! top) )
          solveMask (!m,!c) ress =
              solveCell row (delete c todo) (untag m .|. board) (m:soln) ress

main :: IO ()
main = do
    n <- return.read.head =<< getArgs
    let nsolutions = take n solutions
    putStrLn $ (show $ length nsolutions) ++ " solutions found\n"
    putStrLn . format False . minimum $ nsolutions
    putStrLn . format False . maximum $ nsolutions