ImageLib

fun load _ = (); (* for SmlNJ *)

(* for MOSML *)
load "Vector";
load "Int";
load "Real";
load "String";
load "Math";
load "Random";

type grayimg = int vector vector
type img = (int*int*int) vector vector



(* File Management *)

fun toBytes 0 i = nil
  | toBytes n i = (i mod 256) :: toBytes (n-1) (i div 256)

fun toByteString n = implode o (map Char.chr) o (toBytes n)
fun apply f = Vector.map (Vector.app f)
fun applyVerbose f = Vector.mapi (fn (y,a) => (print(Int.toString(y)^"\n");Vector.app f a))

fun pixelToString (a,b,c) = toByteString 1 c ^ toByteString 1 b ^ toByteString 1 a
fun grayPixelToString a = toByteString 1 a ^ toByteString 1 a ^ toByteString 1 a

fun saveBMP filename img =
    let
        val h = Vector.length img
        val w = Vector.length (Vector.sub(img,0))
        val fd = TextIO.openOut filename
    in
        if (w*3) mod 4 <> 0 then raise Domain else
        (
            TextIO.output (fd, "BM");(* BM *)
            TextIO.output (fd, toByteString 4 (w*h*3+54));(* Size *)
            TextIO.output (fd, "0000"); (* 0000 *)
            TextIO.output (fd, toByteString 4 54);(* 54 *)
            TextIO.output (fd, toByteString 4 40);(* 40 *)
            TextIO.output (fd, toByteString 4 w);(* width *)
            TextIO.output (fd, toByteString 4 (~h));(* height *)
            TextIO.output (fd, toByteString 2 1); (* 00 *)
            TextIO.output (fd, toByteString 2 24);(* colordepth *)
            TextIO.output (fd, toByteString 4 0);(* compression *)
            TextIO.output (fd, toByteString 4 0);(* size *)
            TextIO.output (fd, toByteString 4 0);(* dpi x *)
            TextIO.output (fd, toByteString 4 0);(* dpi y *)
            TextIO.output (fd, toByteString 4 0);(* colortable *)
            TextIO.output (fd, toByteString 4 0);(* count *)
            apply (fn x => TextIO.output (fd, pixelToString x)) img; (* image *)
            TextIO.closeOut fd
        )
    end

exception InvalidOS
fun savePNG filename img = let val bmpFilename = "tmp.bmp" in (saveBMP bmpFilename img;
    if OS.Process.system(String.concatWith " " ["convert",bmpFilename,filename]) = 0 then
     (OS.Process.system(String.concat["rm ",bmpFilename]);())
    else raise InvalidOS
    ) end



(* usual Helper Procedure *)

fun iter n s f =
    if n<1 then s else iter (n-1) (f s) f
fun newton (a:real, x:real, n:int) : real =
    if n<1 then a else newton (0.5*(a+x/a), x, n-1)
fun sqrt (x:real) = newton (x/2.0, x , 5)
fun sign x = if x>0.0 then 1.0 else (if x<0.0 then ~1.0 else 0.0)
fun first (s:int) (p:int->bool) : int =
    if p s then s else first (s+1) p
fun first' s a p f =
    if p (s,a) then (s,a) else first' (s+1) (f(s,a)) p f



(* usual Helper Math *)

fun abs x = if x<0.0 then ~x else x
val rnd = Real.floor

fun toReal (x,y) = (Real.fromInt x, Real.fromInt y)
fun fromReal (x,y) = (rnd x, rnd y)



(* Random *)

val rand=Random.rand(1,0)
fun randomReal () = Random.randReal(rand)



(* Image Manipulation *)

fun createImg w h f = Vector.tabulate(h,fn y => Vector.tabulate(w, fn x => f(x,y)))

fun getPixel img x y = Vector.sub(Vector.sub(img,y),x)

fun setPixel img f x y =
    let
        val row = Vector.sub(img,y)
    in
        Vector.update(img,y,Vector.update(row,x,f (Vector.sub(row,x))))
    end

fun setPixel2 img f (x,y) =
    let
        val y = Int.min(Vector.length img-1,Int.max(0,y))
        val row = Vector.sub(img,y)
        val x = Int.min(Vector.length row-1,Int.max(0,x))
    in
        Vector.update(img,y,Vector.update(row,x,f (Vector.sub(row,x))))
    end

fun imgMap f = Vector.mapi (fn (y,v) => Vector.mapi (
        fn (x,a) => f(x,y,a)) v)

fun imageMap f = Vector.map (fn (v) => Vector.map (
        fn (a) => f a) v)

fun stretch (xs,ys) img = imgMap (fn (x,y,a) => getPixel img (rnd(Real.fromInt(x)*xs)) (rnd(Real.fromInt(y)*ys))) img

fun crop (w,h) img = createImg w h (fn (x,y) => getPixel img x y)

fun zipWith f img img2 = imgMap (fn (x,y,a) => f(a,getPixel img2 x y)) img



(* ImageConversion *)

val toGrayImage = Vector.map (Vector.map (fn (a,b,c) => (a+b+c) div 3))
val toColorImage = Vector.map (Vector.map (fn g:int => (g,g,g)))
fun toColoredImage (rm,gm,bm) = Vector.map (Vector.map (fn g:int => (g*rm,g*gm,g*bm)))



(* Image Math *)

val findMax = Vector.foldl (fn (v,a) => Real.max(a,Vector.foldl Real.max Real.negInf v)) Real.negInf
fun normalize img = case findMax img of max => imgMap (fn (_,_,a) => rnd(255.0*a/max)) img
fun findMaxInt x = Vector.foldl (fn (v,a) => Int.max(a,Vector.foldl Int.max (Vector.sub(v,0)) v)) (Vector.sub(Vector.sub(x,0),0)) x
val findMin = Vector.foldl (fn (v,a) => Real.min(a,Vector.foldl Real.min Real.posInf v)) Real.posInf
fun normalize2 img = let
        val max = findMax img
        val min = findMin img
    in
        imgMap (fn (_,_,a) => rnd(255.0*(a-min)/(max-min))) img
    end

val toRealImage = imgMap (Real.fromInt o #3)



(* Color Functions *)

fun hsvToRgb (h,s,v) =
    let
        val h'=Real.rem(h,360.0)/60.0
        val c=s*v
        val x=c*(1.0-abs(Real.rem(h', 2.0) -1.0))
        val m=v-c
    in
        case (case rnd(h') of
              0 => (c,x,0.0)
            | 1 => (x,c,0.0)
            | 2 => (0.0,c,x)
            | 3 => (x,0.0,c)
            | 4 => (x,0.0,c)
            | _ => (c,0.0,x)) of
        (r,g,b) => (rnd(255.0*(r+m)),
            rnd(255.0*(g+m)),
            rnd(255.0*(b+m)))
    end

fun toColorHSVImage s v = Vector.map (Vector.map (fn h => hsvToRgb(h,s,v)))
fun toColorHSVImage2 s v = Vector.map (Vector.map (fn h => if h < 0.01 then (0,0,0) else (hsvToRgb(h,s,v))))

fun lerp (a,b) p = a*p+b*(1.0-p)
fun lerpInt a b p = rnd(lerp(Real.fromInt a,Real.fromInt b) p)
fun lerpColor ((r1,g1,b1),(r2,g2,b2),p) = (lerpInt r1 r2 p,lerpInt g1 g2 p,lerpInt b1 b2 p)

fun colorGradient v =
            let
                val len = Vector.length v
            in
                imgMap (fn (_,_,a) => let
                    val a=Int.min(a,255)
                    val p = Real.rem(Real.fromInt(len-1)*Real.fromInt(a)/256.0,1.0)
                    val ind = (len-1)*a div 256
                in
                    lerpColor(Vector.sub(v,ind),Vector.sub(v,ind+1),1.0-p)
                end)
            end

val greenGradient = #[(0,0,0),(255,255,255),(0,100,0),(0,255,0),(0,255,100),(255,255,255)]
val blueGradient = #[(0,0,0),(0,100,100),(0,100,255),(0,0,255),(200,200,255)]
val flameGradient = #[(0,0,0),(255,0,0),(255,255,0)]
val rgbGradient = #[(0,0,0),(255,0,0),(0,255,0),(0,0,255)]
val landGradient = #[(255,255,255),(0,255,0),(255,255,0),(0,0,255),(0,0,0)]
val seaGradient = #[(0,0,100),(0,80,200),(0,150,200),(150,200,200)]



(* TextOutput *)

fun printVec pf = Vector.map (fn v => (Vector.map (fn g => (print o pf) g) v;print "\n"))
val printGrayImg = printVec ((fn a => a^"\t ") o Int.toString)
val printImg = printVec (fn (a,b,c) => Int.toString a ^ Int.toString b ^ Int.toString c)
fun toStringVec f = (String.concatWith "\n") o (Vector.foldr (fn (x, l) => ((Vector.foldr (fn (x,a) => x^a) "") o (Vector.map f)) x::l) nil)
val printGrayImg = printVec ((fn a => a^"\t ") o Int.toString)
val toStringGrayImg = toStringVec ((fn a => a^"\t ") o Int.toString)
val printImg = printVec (fn (a,b,c) => Int.toString a ^ Int.toString b ^ Int.toString c)



(* Complex Functions *)

fun angle (r,i) = Math.atan2(i,r)
val arg=angle
fun radius (r,i) = Math.sqrt(r*r+i*i)
fun fromPolar (ang,rad) = (rad*Math.cos(ang),rad*Math.sin(ang))
fun toPolar (c as (r,i)) = (angle c,radius c)

fun restrict (a,b) (c as (r,i)) = fromPolar(Real.min(b,Real.max(a,angle c)),radius c)
fun toComplex (w,h) (x,y) = let val w2=Real.fromInt(w)/2.0 val h2=Real.fromInt(h)/2.0 in (x/w2-1.0,y/h2-1.0) end
fun fromComplex (w,h) (x,y) = let val w2=Real.fromInt(w)/2.0 val h2=Real.fromInt(h)/2.0 in (x*w2+w2,y*h2+h2) end

fun complexToHsv (r,i) = let
        val r=Real.max(Real.min(2.0,r),~2.0)
        val i=Real.max(Real.min(2.0,i),~2.0)
        val rad = Real.max(Real.min(0.0,2.0/Math.sqrt(r*r+i*i)),1.0)
        val phi = Math.atan2(r,i)
    in
        (phi*180.0/Math.pi+180.0,rad,rad)
    end

fun modifyAngleReal (bounds as (w,h)) f pos =
    let
        val c = toComplex bounds pos
        val ang = angle c
        val rad = radius c
    in
        ((fromComplex bounds) o fromPolar) (f ang,rad)
    end

fun modifyAngle (bounds as (w,h)) f = fromReal o (modifyAngleReal bounds f) o (toReal)



(* Advance Pixeldraw *)

fun drawTimes img _ 0 _ _ = img
  | drawTimes img f n cf pos = drawTimes (setPixel2 img f pos) f (n-1) cf (cf pos)

fun rotatePixel bounds n img f pos =
    let
        val c = toComplex bounds (toReal pos)
        val ang = angle c
        val rad = radius c
    in
        drawTimes img f n (modifyAngle bounds (fn ang => ang+2.0*Math.pi/Real.fromInt(n))) pos
    end

fun setMirrorPixel bounds draw img f (p as (x,y)) =
    let
        val (x2,y2) = modifyAngle bounds op~ p
        val img = draw img f (x,y)
        val img = draw img f (x2,y2) (* no reference => ugly *)
    in
        img
    end



(* for Newton (Expression+Derivation) *)

exception Unbound of string
val empty = fn x => raise Unbound x
fun update env id v var =
  if var = id then v else env var
datatype exp = C of real*real | V of string | A of exp * exp | M of exp * exp | D of exp * exp | F of string * exp
fun S (a,b) = A(a,M(C(~1.0,0.0),b))
fun P (a,0) = C(1.0,0.0)
  | P (a,1) = a
  | P (a,n) = M(a,P(a,n-1))
fun I(x) = C(x,0.0)
fun diff x (C _)        = C (0.0,0.0)
  | diff x (V v)        = if x = v then C (1.0,0.0) else C (0.0,0.0)
  | diff x (A (a,b))    = A(diff x a, diff x b)
  | diff x (M (a,b))    = A(M(diff x a, b),M(a,diff x b))
  | diff x (D (a,b))    = D(S(M(diff x a, b),M(a,diff x b)),M(b,b))
  | diff x (F (s,a))    = case diff x a of e =>
                            M(e,case s of
                                  "abs" => D(a,F("abs",a))
                                | "exp" => F("exp",a)
                                | "ln" => D(C(1.0,0.0),a)
                                | "cos" => M(C(~1.0,0.0),F("sin",a))
                                | "sin" => F("cos",a)
                             )

fun eval f (C c)        = c
  | eval f (V v)        = f v
  | eval f (A (a,b))    = (case (eval f a, eval f b) of
                            ((r1,i1),(r2,i2)) => (r1+r2,i1+i2))
  | eval f (M (a,b))    = (case (eval f a, eval f b) of
                            ((r1,i1),(r2,i2)) => (r1*r2-i1*i2,r1*i2+r2*i1))
  | eval f (D (a,b))    = let
                            val (fx,fy)=eval f a
                            val (fx',fy')=eval f b
                            val rez = 1.0/(fx'*fx'+fy'*fy')
                          in
                            (rez*(fx*fx'+fy*fy'),
                                    rez*(fy*fx'-fx*fy'))
                          end
  | eval f (F(s,a))     = (case eval f a of (r,i) => case s of
                              "abs" => (Math.sqrt(r*r+i*i),0.0)
                            | "exp" => (Math.exp(r)*Math.cos(i),Math.exp(r)*Math.sin(i))
                            | "ln" => (Math.ln(Math.sqrt(r*r+i*i)), arg(r,i))
                            | "sin" => (Math.sin(r)*Math.cosh(i),Math.cos(r)*Math.sinh(i))
                            | "cos" => (Math.cos(r)*Math.cosh(i),~ (Math.sin r)*Math.sinh(i))
                          )



(* Herausforderungen: *)
(*
Programmieren Sie folgende Fraktale:
Hopalong
Mandelbrot
Burning Ship
Newton Fractal
Flame Fractal
Perlin Noise (+ landscape, clouds, fire, wood, marble)
DLA with mirror axis and 6-fold rotation
IFS Fractals (koch, sierpinski, far)n

none of these is longer than 20 lines of code
*)