{-
  Copyright 2012 Ken Takusagawa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

  -}
{-# LANGUAGE ScopedTypeVariables #-}
module Main  where {
import Control.Monad(replicateM,foldM);
import Data.List(intersperse);
import Control.Exception(assert);
import System.Environment(getArgs);
import Data.Array.IO(newListArray,range,getBounds,readArray,getElems,writeArray,readArray,IOArray);
import System.Random(getStdRandom,randomR,StdGen,mkStdGen,setStdGen);

type Fdouble  = Float;

type Coord  = (Int, Int);

type Image  = (IOArray Coord Fdouble);

image_bounds :: Int -> (Coord, Coord);
image_bounds size = ((0, 0), ((pred size), (pred size)));

{- |image with all pixels defined. -}
full :: Rand_function -> Int -> (IO Image);
full rand size = (((replicateM ((*) size size))(getStdRandom(rand))) >>= (newListArray (image_bounds size)));

type Rand_function  = (StdGen -> (Fdouble, StdGen));

{- |returns real number uniformly between 0 and 1 -}
gray :: Rand_function;
gray = (randomR (0, 1));

{- |returns 0 or 1 -}
black_white :: Rand_function;
black_white = (let {
inner :: StdGen -> (Int, StdGen);
inner g = (randomR (0, 1) g);
convert :: StdGen -> (Fdouble, StdGen);
convert g = (case (inner g) of {
  (i, g2) -> ((fromIntegral i), g2)
})
} in convert);

{- |sets half the pixels in a checkerboard pattern -}
checkerboard :: Rand_function -> Int -> (IO Image);
checkerboard rand size = (((mapM (init_checkerboard rand))(range((image_bounds size)))) >>= (newListArray (image_bounds size)));

init_checkerboard :: Rand_function -> Coord -> (IO Fdouble);
init_checkerboard rand c = (case (odd_index c) of {
  False -> (getStdRandom rand);
  True -> (return (error "checkerboard not filled"))
});

getImageSize :: Image -> (IO Int);
getImageSize image = (do {
  (_, (x, _)) <- (getBounds image);
  (return (succ x));
});

odd_index :: Coord -> Bool;
odd_index (x, y) = (case (mod ((+) x y) 2) of {
  1 -> True;
  _ -> False
});

wrapped_read :: Image -> Coord -> (IO Fdouble);
wrapped_read image (x, y) = (do {
  (_z, (xmax, ymax)) <- (getBounds image);
  (readArray image ((mod x ((+) 1 xmax)), (mod y ((+) 1 ymax))));
});

{- |orthogonal neighbors -}
coords_checkerboard :: Coord -> ([] Coord);
coords_checkerboard (x, y) = [((pred x), y), ((succ x), y), (x, (pred y)), (x, (succ y))];

{- |diagonal neighbors, forming an X -}
coords_double :: Coord -> ([] Coord);
coords_double (x, y) = (do {
  [f, g] <- (replicateM 2 [pred, succ]);
  (return ((f x), (g y)));
});

average_read :: Image -> ([] Coord) -> (IO Fdouble);
average_read image c = (do {
  (vals :: ([] Fdouble)) <- (mapM (wrapped_read image) c);
  (return ((/) (sum vals) (fromIntegral (length c))));
});

{- |sets a pixel with a value averaged from given neighbors -}
one_pixel :: Image -> (Coord -> ([] Coord)) -> Coord -> (IO ());
one_pixel image which_neighbors c = (((average_read image)(which_neighbors(c))) >>= (writeArray image c));

fill_checkerboard :: Image -> (IO ());
fill_checkerboard image = ((getBounds image) >>= (return . (filter odd_index) . range) >>= (mapM_ (one_pixel image coords_checkerboard)));

pgm_image :: Image -> (IO ());
pgm_image image = (do {
  (putStrLn "P2");
  (size :: Int) <- (getImageSize image);
  (putStrLn ((show size) ++ " " ++ (show size)));
  (putStrLn "255");
  ((getElems image) >>= (mapM_ (print . clip_round)));
});

clip_round :: Fdouble -> Int;
clip_round x = (case (floor ((*) 256 x)) of {
  256 -> 255;
  i -> i
});

{- |doubles image, leaving 3/4 pixels undefined -}
double_image :: Image -> (IO Image);
double_image image = (do {
  (oldsize :: Int) <- (getImageSize image);
  let {
newsize :: Int;
newsize = ((*) 2 oldsize)
};
  (((mapM (init_double image))(range((image_bounds newsize)))) >>= (newListArray (image_bounds newsize)));
});

init_double :: Image -> Coord -> (IO Fdouble);
init_double old (x, y) = (case ((divMod x 2), (divMod y 2)) of {
  ((sx, 0), (sy, 0)) -> (readArray old (sx, sy));
  _ -> (return (error "undefined double pixel"))
});

double_fillable_index :: Coord -> Bool;
double_fillable_index (x, y) = ((&&) ((==) 1 (mod x 2)) ((==) 1 (mod y 2)));

{- |fills in some undefined pixels from a recently doubled image -}
fill_double :: Image -> (IO ());
fill_double image = ((getBounds image) >>= (return . (filter double_fillable_index) . range) >>= (mapM_ (one_pixel image coords_double)));

{- |doubles the dimensions of the image, applying smoothing -}
auto_double :: Image -> (IO Image);
auto_double small = (do {
  image <- (double_image small);
  (fill_double image);
  (fill_checkerboard image);
  (return image);
});

auto_checkerboard :: Rand_function -> Int -> (IO Image);
auto_checkerboard rand size = (do {
  image <- (checkerboard rand size);
  (fill_checkerboard image);
  (return image);
});

multiple_smooth1 :: Integer -> Image -> (IO Image);
multiple_smooth1 n small = (case (compare n 0) of {
  EQ -> (return small);
  GT -> ((auto_double small) >>= (multiple_smooth1 (pred n)))
});

multiple_smooth :: Int -> Image -> (IO Image);
multiple_smooth n small = (iterateM auto_double small n);

iterateM :: forall a m . (Monad m) => (a -> (m a)) -> a -> Int -> (m a);
iterateM f start n = (let {
inner :: a -> () -> (m a);
inner x _ = (f x)
} in (foldM inner start (replicate n ())));

read_rand :: String -> Rand_function;
read_rand s = (case s of {
  "binary" -> black_white;
  "gray" -> gray;
  _ -> (error "rand = {binary | gray}")
});

read_starting :: String -> (Rand_function -> Int -> (IO Image));
read_starting s = (case s of {
  "full" -> full;
  "checkerboard" -> auto_checkerboard;
  _ -> (error "starting = {full | checkerboard}")
});

main :: (IO ());
main = (do {
  (getArgs >>= (\lambda_case_var ->case lambda_case_var of {
  [starting, rand, size, iterations] -> (((read_starting starting) (read_rand rand) (read size)) >>= (multiple_smooth (read iterations)) >>= pgm_image);
  _ -> (putStrLn "starting rand size iterations")
}));
})
}