content: {-# LANGUAGE FlexibleContexts #-}
  {-# OPTIONS_GHC -Wno-name-shadowing #-}
  
  module Main where
  
  import Bvh (buildBvh)
  import Camera (Camera, CameraTrait (..))
  import Codec.Picture
  import Data.Time.Clock
  import HittableList (mkHittableList)
  import Materials.Dielectric (Dielectric (..))
  import Materials.DiffuseLight (DiffuseLight (..))
  import Materials.Lambertian (Lambertian (..))
  import Materials.Metal (mkMetal)
  import Object (SomeObject, mkSomeObject)
  import Random (uniformVec3M)
  import Scene (Scene (..), SomeWorld (MkSomeWorld), mkScene, renderSceneIO)
  import Shapes.Ellipse (mkEllipse)
  import Shapes.Quad (mkBox, mkQuad)
  import Shapes.Sphere (mkMovingSphere, mkSphere)
  import Shapes.Tri (mkTri)
  import System.Random (mkStdGen)
  import System.Random.Stateful (StatefulGen, newIOGenM, uniformRM)
  import Text.Printf
  import Textures.Checker (mkCheckerTexture)
  import Textures.Image (mkImageTexture)
  import Textures.Perlin (mkNoiseTexture)
  import Textures.SolidColor (SolidColor (SolidColor))
  import Vec3 (Vec3 (..))
  
  width :: Int
  width = 512
  
  main :: IO ()
  main = do
    t1 <- getCurrentTime
    printf "Program started at %s\n" (show t1)
    let selectedScene = 4
        aspectRatio = 16.0 / 9.0
        samplesPerPixel = 200
        maxDepth = 50
    gen <- newIOGenM (mkStdGen 2024)
    scene <- case selectedScene of
      0 -> finalScene width aspectRatio samplesPerPixel maxDepth 22 gen
      1 -> perlinSpheresScene width aspectRatio samplesPerPixel maxDepth gen
      2 -> quadsScene width aspectRatio samplesPerPixel maxDepth
      3 -> simpleLight width aspectRatio samplesPerPixel maxDepth gen
      4 -> cornellBox width aspectRatio samplesPerPixel maxDepth gen
      _ -> earthScene width aspectRatio samplesPerPixel maxDepth
    printf "SamplesPerPixel: %s\nMaxDepth: %s\nImage Width: %s\n" (show samplesPerPixel) (show maxDepth) (show width)
    image <- renderSceneIO scene gen
    saveJpgImage 100 "test.jpg" (ImageRGB8 image)
    t2 <- getCurrentTime
    printf "Time taken: %s\n" (show $ diffUTCTime t2 t1)
    printf "Program finished %s\n" (show t2)
  
  mapPairs :: (StatefulGen g m) => g -> (Int, Int) -> m SomeObject
  mapPairs gen (a, b) = do
    chooseMatP :: Double <- uniformRM (0, 1) gen
    offsetX <- uniformRM (0, 1) gen
    offsetZ <- uniformRM (0, 1) gen
    albedo <- SolidColor . (\v -> v <.> v) <$> uniformVec3M (0, 1) gen
    metalColor <- uniformVec3M (0.5, 1) gen
    let center = fromXYZ (fromIntegral a + 0.9 * offsetX, 0.2, fromIntegral b + 0.9 * offsetZ)
    center2 <- (<+> center) . fromXYZ . (0,,0) <$> uniformRM (0, 1) gen
    let sphere
          | chooseMatP < 0.8 = mkSomeObject (mkMovingSphere center center2 0.2) (Lambertian albedo)
          | chooseMatP < 0.95 = mkSomeObject (mkSphere center 0.2) (mkMetal metalColor 0.5)
          | otherwise = mkSomeObject (mkSphere center 0.2) (Dielectric 1.5)
    return sphere
  
  finalScene :: (StatefulGen g IO) => Int -> Double -> Int -> Int -> Int -> g -> IO Scene
  finalScene w aspectRatio samplesPerPixel maxDepth n gen = do
    let lookFrom = fromXYZ (13, 2, 3)
        lookAt = fromXYZ (0, 0, 0)
        vUp = fromXYZ (0, 1, 0)
        vfov = 20
        defocusAngle = 0.6
        focusDist = 10
        camera :: Camera
        camera =
          createCamera
            w
            aspectRatio
            vfov
            lookFrom
            lookAt
            vUp
            defocusAngle
            focusDist
            samplesPerPixel
            maxDepth
            (fromXYZ (0.7, 0.8, 1.0))
        checker =
          mkCheckerTexture
            (SolidColor $ fromXYZ (0.2, 0.3, 0.1))
            (SolidColor $ fromXYZ (0.9, 0.9, 0.9))
            0.32
        materialGround = Lambertian checker
        sphereGround = mkSphere (fromXYZ (0, -1000, -1)) 1000
        material1 = Dielectric 1.5
        sphere1 = mkSphere (fromXYZ (0, 1, 0)) 1.0
        material2 = Lambertian $ SolidColor (fromXYZ (0.4, 0.2, 0.1))
        sphere2 = mkSphere (fromXYZ (-4, 1, 0)) 1.0
        material3 = mkMetal (fromXYZ (0.7, 0.6, 0.5)) 0
        sphere3 = mkSphere (fromXYZ (4, 1, 0)) 1.0
        actualMapFn = mapPairs gen
        halfRange = floor (fromIntegral n / 2 :: Double)
        (leftN, rightN) = (-halfRange, halfRange)
        pairs = [(a, b) | a <- [leftN .. rightN], b <- [leftN .. rightN]]
    spheres <- mapM actualMapFn pairs
    let world = MkSomeWorld $ buildBvh ([mkSomeObject sphereGround materialGround, mkSomeObject sphere3 material3, mkSomeObject sphere1 material1, mkSomeObject sphere2 material2] ++ spheres) 0
    pure $ mkScene camera world
  
  earthScene :: Int -> Double -> Int -> Int -> IO Scene
  earthScene w aspectRatio samplesPerPixel maxDepth = do
    let lookFrom = fromXYZ (0, 0, 12)
        lookAt = fromXYZ (0, 0, 0)
        vUp = fromXYZ (0, 1, 0)
        vfov = 20
        defocusAngle = 0
        focusDist = 10
        camera :: Camera
        camera =
          createCamera
            w
            aspectRatio
            vfov
            lookFrom
            lookAt
            vUp
            defocusAngle
            focusDist
            samplesPerPixel
            maxDepth
            (fromXYZ (0.7, 0.8, 1.0))
    earthTexture <- mkImageTexture "assets/earthmap.jpg"
    let earthSurface = Lambertian earthTexture
        globe = mkSomeObject (mkSphere (fromXYZ (0, 0, 0)) 2) earthSurface
        world = MkSomeWorld $ mkHittableList [globe]
    pure $ mkScene camera world
  
  perlinSpheresScene :: (StatefulGen g IO) => Int -> Double -> Int -> Int -> g -> IO Scene
  perlinSpheresScene w aspectRatio samplesPerPixel maxDepth gen = do
    let lookFrom = fromXYZ (13, 2, 3)
        lookAt = fromXYZ (0, 0, 0)
        vUp = fromXYZ (0, 1, 0)
        vfov = 20
        defocusAngle = 0
        focusDist = 10
        camera :: Camera
        camera =
          createCamera
            w
            aspectRatio
            vfov
            lookFrom
            lookAt
            vUp
            defocusAngle
            focusDist
            samplesPerPixel
            maxDepth
            (fromXYZ (0.7, 0.8, 1.0))
    perlinTexture <- mkNoiseTexture 4 gen
    let perlinSurface = Lambertian perlinTexture
        sphereGround = mkSomeObject (mkSphere (fromXYZ (0, -1000, 0)) 1000) perlinSurface
        sphere1 = mkSomeObject (mkSphere (fromXYZ (0, 2, 0)) 2.0) perlinSurface
        world = MkSomeWorld $ mkHittableList [sphereGround, sphere1]
    pure $ mkScene camera world
  
  quadsScene :: Int -> Double -> Int -> Int -> IO Scene
  quadsScene w aspectRatio samplesPerPixel maxDepth = do
    let lookFrom = fromXYZ (0, 0, 9)
        lookAt = fromXYZ (0, 0, 0)
        vUp = fromXYZ (0, 1, 0)
        vfov = 80
        defocusAngle = 0
        focusDist = 10
        camera :: Camera
        camera =
          createCamera
            w
            aspectRatio
            vfov
            lookFrom
            lookAt
            vUp
            defocusAngle
            focusDist
            samplesPerPixel
            maxDepth
            (fromXYZ (0.7, 0.8, 1.0))
    let leftRed = Lambertian . SolidColor $ fromXYZ (1, 0.2, 0.2)
        backGreen = Lambertian . SolidColor $ fromXYZ (0.2, 1, 0.2)
        rightBlue = Lambertian . SolidColor $ fromXYZ (0.2, 0.2, 1)
        upperOrange = Lambertian . SolidColor $ fromXYZ (1, 0.5, 0.0)
        lowerTeal = Lambertian . SolidColor $ fromXYZ (0.2, 0.8, 0.8)
        quads =
          map
            (uncurry mkSomeObject)
            [ (mkQuad (fromXYZ (-3, -2, 5)) (fromXYZ (0, 0, -4)) (fromXYZ (0, 4, 0)), leftRed)
            -- (mkQuad (fromXYZ (-2, -2, 0)) (fromXYZ (4, 0, 0)) (fromXYZ (0, 4, 0)), backGreen)
            -- (mkQuad (fromXYZ (3, -2, 1)) (fromXYZ (0, 0, 4)) (fromXYZ (0, 4, 0)), rightBlue)
            -- (mkQuad (fromXYZ (-2, 3, 1)) (fromXYZ (4, 0, 0)) (fromXYZ (0, 0, 4)), upperOrange),
            -- (mkQuad (fromXYZ (-2, -3, 5)) (fromXYZ (4, 0, 0)) (fromXYZ (0, 0, -4)), lowerTeal)
            ]
        tris =
          map
            (uncurry mkSomeObject)
            [ (mkTri (fromXYZ (-2, 3, 1)) (fromXYZ (4, 0, 0)) (fromXYZ (0, 0, 4)), upperOrange),
              (mkTri (fromXYZ (-2, -3, 5)) (fromXYZ (4, 0, 0)) (fromXYZ (0, 0, -4)), lowerTeal)
            ]
        ellipses =
          [ mkSomeObject (mkEllipse (fromXYZ (-2, -2, 0)) (fromXYZ (4, 0, 0)) (fromXYZ (0, 4, 0))) backGreen,
            mkSomeObject (mkEllipse (fromXYZ (3, -2, 1)) (fromXYZ (0, 0, 4)) (fromXYZ (0, 4, 0))) rightBlue,
            mkSomeObject (mkEllipse (fromXYZ (-2, -2, 3)) (fromXYZ (2, 0, 0)) (fromXYZ (0, 2, 0))) rightBlue
          ]
        world = MkSomeWorld . mkHittableList $ quads ++ tris ++ ellipses
    pure $ mkScene camera world
  
  simpleLight :: (StatefulGen g IO) => Int -> Double -> Int -> Int -> g -> IO Scene
  simpleLight w aspectRatio samplesPerPixel maxDepth gen = do
    let lookFrom = fromXYZ (26, 3, 6)
        lookAt = fromXYZ (0, 2, 0)
        vUp = fromXYZ (0, 1, 0)
        vfov = 20
        defocusAngle = 0
        focusDist = 10
        camera :: Camera
        camera =
          createCamera
            w
            aspectRatio
            vfov
            lookFrom
            lookAt
            vUp
            defocusAngle
            focusDist
            samplesPerPixel
            maxDepth
            (fromXYZ (0.0, 0.0, 0.0))
    perlinTexture <- mkNoiseTexture 4 gen
    let perlinSurface = Lambertian perlinTexture
        sphereGround = mkSomeObject (mkSphere (fromXYZ (0, -1000, 0)) 1000) perlinSurface
        sphere1 = mkSomeObject (mkSphere (fromXYZ (0, 2, 0)) 2.0) perlinSurface
        diffLight = DiffuseLight $ SolidColor (fromXYZ (4, 4, 4))
        lightSource1 = mkSomeObject (mkTri (fromXYZ (3, 1, -2)) (fromXYZ (2, 0, 0)) (fromXYZ (0, 2, 0))) diffLight
        lightSource2 = mkSomeObject (mkSphere (fromXYZ (0, 7, 0)) 2) diffLight
        world = MkSomeWorld $ mkHittableList [sphereGround, sphere1, lightSource1, lightSource2]
    pure $ mkScene camera world
  
  cornellBox :: (StatefulGen g IO) => Int -> Double -> Int -> Int -> g -> IO Scene
  cornellBox w aspectRatio samplesPerPixel maxDepth gen = do
    let lookFrom = fromXYZ (278, 278, -800)
        lookAt = fromXYZ (278, 278, 0)
        vUp = fromXYZ (0, 1, 0)
        vfov = 40
        defocusAngle = 0
        focusDist = 10
        camera :: Camera
        camera =
          createCamera
            w
            aspectRatio
            vfov
            lookFrom
            lookAt
            vUp
            defocusAngle
            focusDist
            samplesPerPixel
            maxDepth
            (fromXYZ (0.0, 0.0, 0.0))
    let red = Lambertian . SolidColor $ fromXYZ (0.65, 0.05, 0.05)
        white = Lambertian . SolidColor $ fromXYZ (0.73, 0.73, 0.73)
        green = Lambertian . SolidColor $ fromXYZ (0.12, 0.45, 0.15)
        light = DiffuseLight . SolidColor $ fromXYZ (15, 15, 15)
        room =
          map
            (uncurry mkSomeObject)
            [ (mkQuad (fromXYZ (555, 0, 0)) (fromXYZ (0, 555, 0)) (fromXYZ (0, 0, 555)), green),
              (mkQuad (fromXYZ (0, 0, 0)) (fromXYZ (0, 555, 0)) (fromXYZ (0, 0, 555)), red),
              (mkQuad (fromXYZ (0, 0, 0)) (fromXYZ (555, 0, 0)) (fromXYZ (0, 0, 555)), white),
              (mkQuad (fromXYZ (555, 555, 555)) (fromXYZ (-555, 0, 0)) (fromXYZ (0, 0, -555)), white),
              (mkQuad (fromXYZ (0, 0, 555)) (fromXYZ (555, 0, 0)) (fromXYZ (0, 555, 0)), white)
            ]
        lights = [mkSomeObject (mkQuad (fromXYZ (343, 554, 332)) (fromXYZ (-130, 0, 0)) (fromXYZ (0, 0, -105))) light]
        box1 = map (uncurry mkSomeObject . (,white)) (mkBox (fromXYZ (130, 0, 65)) (fromXYZ (295, 165, 230)))
        box2 = map (uncurry mkSomeObject . (,green)) (mkBox (fromXYZ (265, 0, 295)) (fromXYZ (430, 330, 460)))
        world = MkSomeWorld $ buildBvh (room ++ lights ++ box1 ++ box2) 0
    pure $ mkScene camera world
  