{-
Draw images of the Ulam spiral.

The numbers along the spiral can be any arithmetic sequence a*n+b,
where n starts from 0.  Command-line arguments are image width, a, b.
Output is PNG.  For example, a 320x320 image of the classic Ulam
spiral with 1 at the center is generated with:

./ulam-spiral 320 1 1 > output.png

There is a space leak, growing faster than linear with the number of
pixels (which is quadratic).

547 bytes per pixel at 1000x1000
637 at 2000
650 at 2500.

Copyright 2019 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, LambdaCase, PackageImports #-}
module Main where {
import System.Environment(getArgs);
import Control.Exception(assert);
import Debug.Trace(trace);
import Data.Function((&));
import Control.Category((>>>));
import Prelude hiding((.),(>>));
--import System.IO(hSetBuffering,stdout,BufferMode(LineBuffering));
--import System.IO(stderr,hPutStrLn);
import qualified Data.List as List;
--import qualified Control.Monad as Monad;
--import qualified Data.Maybe as Maybe;
--import qualified Data.Map as Map; import Data.Map(Map);
--import qualified Data.Set as Set; import Data.Set(Set);
--import qualified Data.Bifunctor as Bifunctor;
--import qualified Data.Tuple as Tuple;
import "extra" Data.Tuple.Extra(both);
import Primality2(isPrime);
import qualified Codec.Picture.Types as JuicyT;
import qualified Codec.Picture as Juicy;
import Control.Monad.Primitive(PrimMonad,PrimState);
import qualified Data.ByteString.Lazy as Lazy;
import qualified Data.Biapplicative;

type Ii = Integer;

main :: IO();
main = getArgs >>= \case{
[n,a,b] -> go (Size $ read n) (Arithmeticsequence (read a) (read b));
_ -> undefined;
};

go :: Size -> Arithmeticsequence -> IO();
go n ab = do {
-- picplot n ab >>= (Juicy.encodeBitmap >>> Lazy.putStr);
picplot n ab >>= (Juicy.encodePng >>> Lazy.putStr);
};

type Pair = (Ii,Ii);

newtype Size = Size Ii deriving (Show);

-- counterclockwise
rotate90 :: Pair -> Pair;
rotate90 (x,y) = (negate y, x);

-- one L of a spiral
l1 :: Ii -> Pair -> [Pair];
l1 n dir = List.genericReplicate n dir ++ List.genericReplicate n (rotate90 dir);

-- two Ls go all the way around once

-- normal ulam spiral.  numbers densely fill the plane.
l2normal :: Pair -> Ii -> [Pair];
l2normal dir n = l1 (2*n+1) dir ++ l1 (2*n+2) (dir & rotate90 & rotate90) ;

-- a different spiral with empty space between windings
l2alternate :: Pair -> Ii -> [Pair];
l2alternate dir n = List.genericReplicate (4*n+1) dir
++ List.genericReplicate (4*n+2) (dir & rotate90)
++ List.genericReplicate (4*n+3) (dir & rotate90 & rotate90)
++ List.genericReplicate (4*n+4) (dir & rotate90 & rotate90 & rotate90)
;

-- direction facing after going around once stays constant, so we can hardcode it.
fullturn :: Ii -> [Pair];
fullturn = l2normal (1,0);

spiralsteps :: [Pair];
spiralsteps = concatMap fullturn [0..];

addcoords :: Pair -> Pair -> Pair;
-- addcoords (a,b) (x,y) = (a+x, b+y);
addcoords = Data.Biapplicative.biliftA2 (+) (+);

coordrange :: [Pair] -> (Pair,Pair);
coordrange l = let {
t :: ([Ii],[Ii]);
t = both (\f -> map f l) (fst,snd)
}
in (both minimum t, both maximum t);

allcoords :: [Pair];
allcoords = scanl addcoords (0,0) spiralsteps;

squarecoords :: Size -> [Pair];
squarecoords (Size n) = allcoords & List.genericTake (n*n);
-- for l2alternate, n*n does not end at a nice point

offsetsquare :: Size -> [Pair];
offsetsquare n = let {
c :: [Pair];
c = squarecoords n;
getmin :: Pair;
getmin = coordrange c & fst & both negate;
--subtract off the minimum so the minimum becomes zero
} in map (addcoords getmin) c;

getsize :: Size -> (Ii,Ii);
-- image dimensions is max + 1 because coordinates start at 0
getsize = offsetsquare >>> coordrange >>> snd >>> both succ;

-- a * t + b, first is a, second is b.
data Arithmeticsequence = Arithmeticsequence Ii Ii deriving (Show);

arithmeticsequence :: Arithmeticsequence -> [Ii];
arithmeticsequence (Arithmeticsequence a b) = assert (a/=0) $ map (\t -> a*t + b) [0..];

-- center dot is 1
dots1 :: Size -> [(Pair,Bool)];
dots1 n = zip (offsetsquare n) $ arithmeticsequence (Arithmeticsequence 1 1) & map isPrime;

-- (abs >>> isPrime) is also a possibility, but we avoid it to follow the principle of least surprise
dots :: Size -> Arithmeticsequence -> [(Pair,Bool)];
dots n ab = zip (offsetsquare n) $ arithmeticsequence ab & map isPrime;

-- becomes visible if we do l2alternate
backgroundcolor :: Juicy.Pixel8;
backgroundcolor = 0xf0;

colorprime :: Juicy.Pixel8;
colorprime = 0;

colorcomposite :: Juicy.Pixel8;
colorcomposite = 255;

-- draw additional points (ink bleed) in hopes of making horizontal and vertical lines more visible.
-- do be done someday as a postprocessing step.
bleedcolor :: Juicy.Pixel8;
bleedcolor = 64;

rangecheckedplot :: (PrimMonad m) => JuicyT.MutableImage (PrimState m) Juicy.Pixel8 -> Integer -> Integer -> Juicy.Pixel8 -> m ();
rangecheckedplot image a b c = let {
i :: Int;
i= fromInteger a;
j :: Int;
j = (JuicyT.mutableImageHeight image - 1) - fromInteger b;  -- flip top-bottom because JuicyPixels canvas has positive Y going downward
} in if (0 <= i) && (i < JuicyT.mutableImageWidth image) && (0 <= j) && (j < JuicyT.mutableImageHeight image)
then Juicy.writePixel image i j c
else trace ("out of range " ++ show a ++ " " ++ show b) $ return();

plot :: forall m . (PrimMonad m) => JuicyT.MutableImage (PrimState m) Juicy.Pixel8 -> (Pair,Bool) -> m();
plot image ((x,y),v) = let {
thecolor :: Juicy.Pixel8;
thecolor = if v then colorprime else colorcomposite;
} in do {
  rangecheckedplot image x y thecolor;
};


picplot :: Size -> Arithmeticsequence -> IO (Juicy.Image Juicy.Pixel8);
picplot n ab = do {
let {size :: (Int,Int) ; size = getsize n & both fromInteger};
-- PrimState IO seems more elegant than GHC.Prim.RealWorld
-- ScopedTypeVariables
image :: JuicyT.MutableImage (PrimState IO) Juicy.Pixel8 <- JuicyT.createMutableImage (fst size) (snd size) backgroundcolor;
dots n ab & mapM_ (plot image);
JuicyT.freezeImage image
};


} --end