module Main where

import Control.Monad (mapM)
import Control.Monad.Random
import Data.Ix
import Data.Maybe
import Debug.Trace
import Hittable
import Linear.V2
import Linear.V3
import Linear.Vector
import Material
import Maths
import Ray
import Text.Printf

type Color = (Double, Double, Double)

type Coord = (Int, Int)

type ImageDimensions = (Int, Int)

type ViewPortDimensions = (Double, Double)

-- TODO
-- antialiasing

backgroundColor = V3 0.0 0.0 0.0

scatter :: (RandomGen g) => Ray Double -> Hit -> Rand g (Ray Double)
scatter (Ray o d) (Hit t p n (Lambertian c)) = do
  r <- randV3
  let target = p ^+^ colorWithNormal (unitVector n) ^+^ r
  return $ Ray p (target ^-^ p)
scatter (Ray o d) (Hit t p n (Metal c)) = do
  let reflected = reflect (unitVector d) n
  return $ Ray p reflected
scatter (Ray o d) (Hit t p n (Glass c)) = do
  let refraction_ratio = if isFrontFace (Ray o d) n then 1.0 / 1.5 else 1.0
  let ud = unitVector d
  let ct = min (ud `dotP` (unitVector n)) 1.0
  let st = sqrt (1.0 - ct * ct)
  rVal <- rand

  if refraction_ratio * st > 1.0 || reflectance ct refraction_ratio < rVal
    then do
      let reflected = reflect ud n
      return $ Ray p reflected
    else do
      let refracted = refract ud n refraction_ratio
      return $ Ray p refracted

getRayColor :: (Shape s, RandomGen g) => [s] -> Int -> Ray Double -> Rand g (V3 Double)
getRayColor ws d r = do
  let h = getBestHit (mapMaybe (testHit r) ws)
  v <- getRayColor' r h ws d
  return v

getRayColor' :: (Shape s, RandomGen g) => Ray Double -> Maybe Hit -> [s] -> Int -> Rand g (V3 Double)
getRayColor' (Ray o d) Nothing _ _ = do return $ colorGradient (unitVector d)
getRayColor' (Ray o d) _ _ 0 = do return $ colorGradient (unitVector d)
getRayColor' (Ray o d) (Just (Hit t p n m)) ws depth = do
  let a = attenuation m
  ray <- scatter (Ray o d) (Hit t p n m)
  nextVal <- getRayColor ws (depth - 1) ray
  return (a * nextVal)

colorGradient :: V3 Double -> V3 Double
colorGradient (V3 x y z) = do
  let t = 0.5 * (y + 1.0)
  (1.0 - t) *^ V3 1.0 1.0 1.0 + t *^ V3 0.5 0.7 1.0

getBestHit [] = Nothing
getBestHit (h : hs) = Just (getBestHit' hs h)

getBestHit' :: [Hit] -> Hit -> Hit
getBestHit' (h : hs) bestHit =
  -- always pick closest hit to camera
  if root h < root bestHit
    then getBestHit' hs h
    else getBestHit' hs bestHit
getBestHit' [] hit = hit

colorWithNormal :: V3 Double -> V3 Double
colorWithNormal v = 0.5 *^ (v ^+^ V3 1.0 1.0 1.0)

-- rayTrace :: (Shape s) => (Int, Int) -> ImageDimensions -> ViewPortDimensions -> [s] -> IO (V3 Double)
rayTrace (x, y) (w, h) (vw, vh) world = do
  l <- replicateM 25 (evalRandIO (getRayColor world 50 $ buildRay (xf * (vw / wf)) (yf * (vh / hf))))
  return $ foldl (^+^) (V3 0.0 0.0 0.0) l
  where
    xf = fromIntegral x
    yf = fromIntegral y
    wf = fromIntegral (w - 1)
    hf = fromIntegral (h - 1)

colorToString :: Color -> IO ()
colorToString (x, y, z) = tupToString (xr, yr, zr)
  where
    -- scale 1
    scale = 1.0 / 25
    xr = truncate ((clamp (sqrt (x * scale)) 0.0 0.999) * 256.0)
    yr = truncate ((clamp (sqrt (y * scale)) 0.0 0.999) * 256.0)
    zr = truncate ((clamp (sqrt (z * scale)) 0.0 0.999) * 256.0)

tupToString :: (Show a) => (a, a, a) -> IO ()
tupToString (x, y, z) = putStrLn (show x ++ " " ++ show y ++ " " ++ show z)

v3ToColor :: V3 Double -> Color
v3ToColor (V3 x y z) = (x, y, z)

main :: IO ()
main = do
  let w = 400
  let h = 400
  let world =
        [ Circle (V3 (0.0) 0.0 (-1.0)) 0.25 (Glass (V3 1.0 1.0 1.0)),
          Circle (V3 (-0.5) 0.0 (-1.0)) 0.25 (Lambertian (V3 0.7 0.1 0.1)),
          Circle (V3 0.0 (-10.25) (-1.0)) 10 (Lambertian (V3 0.7 0.6 0.5))
        ]
  let coords = reverse $ range ((-1 * (w `div` 2), -1 * (h `div` 2)), ((w `div` 2) - 1, (h `div` 2) - 1)) -- "canvas"
  colors <- mapM (\x -> rayTrace x (w, h) (1.0, 1.0) world) coords
  let convertedColors = map v3ToColor colors
  putStrLn "P3"
  putStrLn (show w ++ " " ++ show h)
  putStrLn "255"
  mapM_ colorToString convertedColors
  mapM_ print convertedColors