-- An ugly raytracer in Haskell
-- ****************************
--
-- Written as an excuse to learn Haskell, and as a refresher for 3d geometry concepts.
--
-- The MIT License (MIT)
--
-- Copyright (c) 2013 Matthias P. Braendli
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy
-- of this software and associated documentation files (the "Software"), to deal
-- in the Software without restriction, including without limitation the rights
-- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-- copies of the Software, and to permit persons to whom the Software is
-- furnished to do so, subject to the following conditions:
--
-- The above copyright notice and this permission notice shall be included in all
-- copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-- SOFTWARE.
import System.IO
import Data.Char
import Debug.Trace

-- ppm image file format:
-- P3 <width> <height> <maxcolorval> r g b\nr g b\nr g b\n...
-- max line length: 70


type Angle = Double
type Direction = (Angle, Angle) -- azimut, inclination
type Color = (Int, Int, Int)    -- red green blue
type Coord = (Double, Double, Double) -- x, y, z (x and y on floor, z up)
data Sphere = Sphere Coord Double Color deriving (Show, Eq)

degrees = pi / 180

eye :: Coord
eye = (-40, -60, 20)

x_of (x, _, _) = x
y_of (_, y, _) = y
z_of (_, _, z) = z

sphere1 = Sphere (0, 80, 25) 60 (55,255,0)
sphere2 = Sphere (80, 0, 35) 20 (255,60,200)
--sphere3 = Sphere (0, -80, 5) 20 (5,60,200)
--sphere4 = Sphere (-80, 0, 5) 20 (0,255,255)
--spheres = [sphere1, sphere2, sphere3, sphere4]


filename num = "foo/foo" ++ show num ++ ".ppm"

--spherepos = take 1 [0,20..]
--spherepos = take 10 [0,36..]
spherepos = [60]

spheres num = [ trace ("Sphere at " ++
                    show (round (80 * sin(num * degrees))) ++ "," ++
                    show (round (80 * cos(num * degrees))) ++ ",5" )
                Sphere (80 * sin(num * degrees), 80 * cos(num * degrees), 45) 10 (255,60,0),
                sphere1, sphere2 ]

writenum :: Double -> IO ()
writenum num = trace ("Rendering " ++ show (filename $ round num))
               writeFile (filename num) (image $ spheres num)

main = mapM writenum spherepos


alpha1 = 100 * degrees
alpha2 = 0 * degrees

beta1 = 65 * degrees
beta2 = 120 * degrees

floorscale = 4

w :: Int
w = 192 * 6
h :: Int
h = 108 * 6

oversampling = 1 -- each pixel is oversampling^2 rays

black :: Color
black = (0,0,0)

ov_alphaoffset = ((alpha2 - alpha1) / (fromIntegral w-1)) / oversampling
ov_betaoffset  = ((beta2 - beta1) / (fromIntegral h-1)) / oversampling

ov_alphaoffsets = take (round oversampling) [0,ov_alphaoffset..]
ov_betaoffsets  = take (round oversampling) [0,ov_betaoffset..]

imgheader = "P3 " ++ show w ++ " " ++ show h ++ " 255\n"

alphas = take w [alpha1,(alpha1 + ((alpha2 - alpha1) / (fromIntegral w-1)))..]
betas  = take h [beta1,(beta1 + ((beta2 - beta1) / (fromIntegral h-1)))..]

attenuate_color :: Double -> Color -> Color
attenuate_color factor (r,g,b)  = ( round $ fromIntegral r * factor,
                                    round $ fromIntegral g * factor,
                                    round $ fromIntegral b * factor)

colormix :: Double -> Color -> Color -> Color
colormix amount c1 c2 = (round $ p*fromIntegral r1 + (1-p)*fromIntegral r2,
                    round $ p*fromIntegral g1 + (1-p)*fromIntegral g2,
                    round $ p*fromIntegral b1 + (1-p)*fromIntegral b2)
                    where   (r1, g1, b1) = c1
                            (r2, g2, b2) = c2
                            p = max 0 $ min 1 amount

(a, b, c) `dot` (d, e, f) = a*d + b*e + c*f
(a, b, c) `plus` (d, e, f) = (a+d, b+e, c+f)
(a, b, c) `minus` (d, e, f) = (a-d, b-e, c-f)
s `scale` (a,b,c) = (s*a, s*b, s*c)
magnitude (a,b,c) = sqrt(a^2 + b^2 + c^2)
normalise v = (1/(magnitude v)) `scale` v

-- spherical projection,
-- return coordinates from a given coordinate, extended by given
-- angles to some distance
spherical_proj :: Coord -> Angle -> Angle -> Double -> Coord
spherical_proj (x,y,z) alpha beta dist = (x + dist*(sin beta * cos alpha),
                                          y + dist*(sin beta * sin alpha),
                                          z + dist*cos beta)

cart2spher :: Coord -> Direction
cart2spher coord = ( atan2 y x, acos $ z/r)
                    where   (x, y, z) = coord
                            r = magnitude coord

spher2cart :: Direction -> Coord
spher2cart direction = (sin beta * cos alpha,
                        sin beta * sin alpha,
                        cos beta)
                    where (alpha, beta) = direction

-- reflection on sphere
reflect :: Coord -> Direction -> Direction
reflect normal incoming = (alph, beta)
                    where   a = (b `dot` n) `scale` n
                            b = spher2cart incoming
                            n = normalise normal
                            (alpha, beta) = cart2spher $ ((-2) `scale` a) `plus` b
                    --      (alpha, beta) = cart2spher normal
                            alph = degrees*(fromIntegral (alpha_i `mod` 360) + alpha_f)
                            (alpha_i, alpha_f) = properFraction (alpha/degrees)

-- intersect sphere

-- discr = 4(( A u + B v + C w )^2 - (A^2 + B^2 + C^2)(u^2 + v^2 + w^2))
discr :: Coord -> Direction -> Sphere -> Double
discr source (alpha, beta) (Sphere centre radius _) = 4*(( aa * u + bb * v + cc * w )^2 -
                                        (aa*aa + bb*bb + cc*cc)*(u*u + v*v + w*w - radius^2))
      where u = (x_of source) - (x_of centre)
            v = (y_of source) - (y_of centre)
            w = (z_of source) - (z_of centre)
            aa = sin beta * cos alpha
            bb = sin beta * sin alpha
            cc = cos beta

--debug
simplecolor (alpha, beta)
            | 0   <= alph && alph < 90    =  (a,0,0)
            | 90  <= alph && alph < 180   =  (0,a,0)
            | 180 <= alph && alph < 270   =  (0,0,a)
            | 270 <= alph && alph < 360   =  (a,a,0)
            | otherwise                           =  (0,255,255)
            where   a    = trace ("alpha " ++ show (alph)) (if beta > (90*degrees) then 28 else 255)
                    alph = fromIntegral (alpha_i `mod` 360) + alpha_f
                    (alpha_i, alpha_f) = properFraction (alpha/degrees)

-- the intersect functions return (Coord, Distance, Color)
-- distance = 0 means no intersection
intersect_sphere :: Coord -> [Sphere] -> Direction -> Sphere -> (Double, Color)
intersect_sphere source spheres (alpha, beta) (Sphere centre radius color)
        | delta > 0 =    (t,
                         attenuate_color 0.7 $
                         pixel_color intersect_point spheres reflection_angle
                         --simplecolor reflection_angle
                        )
        | otherwise = (0, black)
            where   t = min ((-b - sqrt(delta)) / (2*a)) ((-b + sqrt(delta)) / (2*a))
                    delta = discr source (alpha, beta) (Sphere centre radius color)
                    a = aa^2 + bb^2 + cc^2
                    b = 2 * (aa*u + bb*v + cc*w)
                    u = (x_of source) - (x_of centre)
                    v = (y_of source) - (y_of centre)
                    w = (z_of source) - (z_of centre)
                    aa = sin beta * cos alpha
                    bb = sin beta * sin alpha
                    cc = cos beta
                    intersect_point = spherical_proj source alpha beta t
                    reflection_angle = reflect normal (alpha, beta)
                    normal = intersect_point `minus` centre
                    --subspheres = filter ((/=) (Sphere centre radius color)) spheres --limit recursion


intersect_point_floor :: Coord -> Direction -> (Coord, Double)
intersect_point_floor (x, y, z) (alpha, beta) =
                      (  (x - z * sin beta * cos alpha / cos beta,
                          y - z * sin beta * sin alpha / cos beta,
                          0),
                      -z / (cos beta) )

direction_color :: Double -> Double -> Int -> Color
direction_color x y attn
        | x > 0 && y > 0  = (attn, 0, 0) -- red
        | x <= 0 && y > 0 = (0, attn, 0) -- green
        | x > 0 && y <= 0 = (attn, attn, 0) -- yellow
        | otherwise       = (0, 0, attn) -- blue

checkerboard_pattern :: Double -> Double -> Int -> Color
checkerboard_pattern x y attn
        | (round (x/floorscale) `mod` 2) == (round (y/floorscale) `mod` 2) = direction_color x y attn
        | otherwise              = (attn, attn, attn)

intersect_floor :: Coord -> Direction -> (Coord, Double, Color)
intersect_floor source (alpha, beta)
        | beta <= 90*degrees     = ((x, y, z), 0, (128, 128, 128))--checkerboard_pattern x y 128)
        | x > (-0.5) && x < 0.5  = ((x, y, z), t, (0, attn, attn)) -- x near 0 : cyan
        | y > (-0.5) && y < 0.5  = ((x, y, z), t, (attn, attn, 0)) -- y near 0 : yellow
        | otherwise              = ((x, y, z), t, checkerboard_pattern x y attn)
        where   attn = max 0 (round (255 - 4*(sqrt $ abs t)))
                ((x, y, z), t) = intersect_point_floor source (alpha, beta)

-- blue is beautiful, but a green tint is nice too. with stars it's even better

is_star_center crit (x, y) = (((x^2 + y^3 + (x+y+2)^3) `mod` (131*90)) < (90 `div` crit))

stardist :: Direction -> (Int, Int) -> Double
stardist p1 p2 = sqrt ( (x1-fromIntegral x2)^2 + (y1-fromIntegral y2)^2)
                    where   (x1, y1) = p1
                            (x2, y2) = p2

roughness = 3
nearest_potential_star_centre :: Direction -> (Int, Int)
nearest_potential_star_centre (x, y) = (((round x) `div` roughness) * roughness + (roughness `div` 2),
                                        ((round y) `div` roughness) * roughness + (roughness `div` 2))

star_color crit darkness (x, y) = if is_star_center crit nearest
                    then colormix ((stardist (x,y) nearest) / 2 ) darkness starlight
                    else darkness
                where   nearest = nearest_potential_star_centre (x,y)

starlight = (250, 250, 180)

skycolor :: Coord -> Direction -> Color
skycolor source (alpha, beta) = star_color 2 darkness (10*alpha/degrees, 10*beta/degrees)
            where   r = 60
                    g = max 0 $ round $ (sqrt (greenness/2)) * 64
                    b = max 0 $ round $ (sqrt (-beta+90*degrees)) / (sqrt (90 * degrees)) * 255
                    greenness = cos (alpha) + 1
                    darkness = (r,g,b)

data SphereIntersect = SphereIntersect Double Color deriving (Eq, Show) -- distance color
instance Ord SphereIntersect where
  (SphereIntersect d1 _) `compare` (SphereIntersect d2 _)
                        | d2 <= 0 = LT
                        | d1 <= 0 = GT
                        | otherwise = d1 `compare` d2

nearest_sphere :: Coord -> Direction -> [Sphere] -> SphereIntersect
nearest_sphere source scoord spheres =
                        minimum [(SphereIntersect distance color) | (distance, color) <- intersections]
                        where intersections = map (intersect_sphere source spheres scoord) spheres

-- also include floor in objects
nearest_obj :: Coord -> Direction -> [Sphere] -> (Double, Color)
nearest_obj source scoord [] = (floordist, floorcolor)
            where   (_, floordist, floorcolor) = intersect_floor source scoord
nearest_obj source scoord spheres
            | floordist == 0 && spheredist > 0           = (spheredist, spherecolor)
            | floordist > spheredist && spheredist > 0   = (spheredist, spherecolor)
            | otherwise                                  = (floordist, floorcolor)
            where   (SphereIntersect spheredist spherecolor)  = nearest_sphere source scoord spheres
                    (_, floordist, floorcolor) = intersect_floor source scoord

-- First iteration
pixel_color :: Coord -> [Sphere] -> Direction -> Color
pixel_color_only_floor source spheres scoord =  floorcolor
                    where   ((x,y,z), floordist, floorcolor) = intersect_floor source scoord
                            (alpha, beta) = scoord

pixel_color source spheres scoord
            | nearest_object_dist > 0 = objcolor
            | beta == 90 * degrees    = (0, 255, 0)
            | otherwise               = skycolor source scoord
            where   (_, beta) = scoord
                    (nearest_object_dist, objcolor) = nearest_obj source scoord spheres


cartProdTranspose xs ys = [(y,x) | x <- xs, y <- ys]
cartProd xs ys = [(x,y) | x <- xs, y <- ys]

pixel_to_ppm (r,g,b) = show r ++ " " ++ show g ++ " " ++ show b ++ "\n"

tuple2sum x y = (a1 + b1, a2 + b2)
               where (a1, a2) = x
                     (b1, b2) = y

-- from one pixel (alpha, beta), get a list of oversampled pixels
oversample :: Direction -> [Direction]
oversample (a,b) = map (tuple2sum (a,b)) (cartProd ov_alphaoffsets ov_betaoffsets)

tuple3sum x y = (a1 + b1, a2 + b2, a3 + b3)
               where (a1, a2, a3) = x
                     (b1, b2, b3) = y

coloraverage :: [Color] -> Color
coloraverage xs = ( round (fromIntegral s1/l),
                    round (fromIntegral s2/l),
                    round (fromIntegral s3/l) )
                    where (s1, s2, s3) = foldr tuple3sum (0,0,0) xs
                          l = fromIntegral (length xs)

-- calculate color of oversampled pixels
ov_color :: [Sphere] -> [Direction] -> Color
ov_color spheres coords  = coloraverage (map (pixel_color eye spheres) coords)

-- list of list of (alpha, beta)-tuples
ov_pixels = map oversample (cartProdTranspose betas alphas)

allpixels spheres = map (ov_color spheres) ov_pixels

image spheres = imgheader ++ (foldr (++) "" (map pixel_to_ppm (allpixels spheres)))