{-# LANGUAGE BangPatterns #-}

module PathTracer where

import Elem3D ( Luz(..), RGB(..), divRGB, scale, modRGB, (.*),
       Point3D (Point3D), normal,nanRGB )
import Figuras ( Obj(..), Shape, getShapeID )
import System.Random (StdGen, split)
import Funciones
    ( objAleatorio, objEspejo, objCristal, formula, 
    colision, brdf, ruletaRusa, addNiebla, dirEspejo)
import Debug.Trace (trace)
import qualified Data.Set as Set

-- | Función principal del path tracer, calcula toda la suma de la luz indirecta.
{-# INLINE pathTracer #-}
pathTracer :: Float -> [Luz] -> Set.Set Shape -> Int -> Obj -> StdGen -> RGB
pathTracer :: Float -> [Luz] -> Set Shape -> Int -> Obj -> StdGen -> RGB
pathTracer Float
rFl [Luz]
luz !Set Shape
figuras Int
ppp Obj
obj StdGen
gen
  | Obj -> Float
mindObj Obj
obj Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
< Float
0 = RGB -> RGB
scale (RGB -> RGB) -> RGB -> RGB
forall a b. (a -> b) -> a -> b
$ Float -> Float -> Float -> RGB
RGB Float
20 Float
50 Float
30
  | Int
ppp Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = RGB
colorIndirecto
  | Bool
otherwise = RGB
colorIndirecto RGB -> RGB -> RGB
forall a. Num a => a -> a -> a
+ Float -> [Luz] -> Set Shape -> Int -> Obj -> StdGen -> RGB
pathTracer Float
rFl [Luz]
luz Set Shape
figuras (Int
ppp Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Obj
obj StdGen
gen' -- Luz directa + indirecta
  where
    colorIndirecto :: RGB
colorIndirecto = Obj -> [Luz] -> Set Shape -> StdGen -> RGB
luzIndirecta Obj
obj [Luz]
luz Set Shape
figuras StdGen
gen''
    (StdGen
gen',StdGen
gen'') = StdGen -> (StdGen, StdGen)
forall g. RandomGen g => g -> (g, g)
split StdGen
gen


-- | Función principal que calcula la luz directa de una escena.
{-# INLINE luzDirecta #-}
luzDirecta :: [Luz] -> Set.Set Shape -> Obj -> RGB
luzDirecta :: [Luz] -> Set Shape -> Obj -> RGB
luzDirecta [Luz]
luces Set Shape
figuras Obj
obj
  | [Luz] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Luz]
luces = Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0
  | Obj -> Float
mindObj Obj
obj Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
< Float
0 = RGB -> RGB
scale (RGB -> RGB) -> RGB -> RGB
forall a b. (a -> b) -> a -> b
$ Float -> Float -> Float -> RGB
RGB Float
20 Float
40 Float
50
  | [Luz] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Luz]
luces Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = {- addNiebla (head luces) obj 0.6 figuras $ -} Obj -> Luz -> Set Shape -> RGB
luzMono Obj
obj ([Luz] -> Luz
forall a. [a] -> a
head [Luz]
luces) Set Shape
figuras
  | [Luz] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length  [Luz]
luces Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = (Obj -> Luz -> Set Shape -> RGB
luzMono Obj
obj ([Luz] -> Luz
forall a. [a] -> a
head [Luz]
luces) Set Shape
figuras RGB -> RGB -> RGB
forall a. Num a => a -> a -> a
+ [Luz] -> Set Shape -> Obj -> RGB
luzDirecta ([Luz] -> [Luz]
forall a. [a] -> [a]
tail [Luz]
luces) Set Shape
figuras Obj
obj) RGB -> Float -> RGB
`divRGB` Float
2 -- Si todas pesaran igual que no es asi

-- | Función auxiliar que calcula la luz directa de una escena solo con una luz.
{-# INLINE luzMono #-}
luzMono :: Obj -> Luz -> Set.Set Shape -> RGB
luzMono :: Obj -> Luz -> Set Shape -> RGB
luzMono Obj
obj (Luz Point3D
pointLuz RGB
rgbLuz Float
intLuz) Set Shape
figuras = RGB
difuso RGB -> RGB -> RGB
forall a. Num a => a -> a -> a
+ RGB
espejo RGB -> RGB -> RGB
forall a. Num a => a -> a -> a
+RGB
cristal
  where
    (Float
kd,Float
ke,Float
kr) = Obj -> (Float, Float, Float)
trObj Obj
obj

    newRGB :: RGB
newRGB = RGB -> Float -> Point3D -> Point3D -> Direction -> RGB -> RGB
formula (RGB -> RGB
scale RGB
rgbLuz) Float
intLuz Point3D
pointLuz (Obj -> Point3D
colObj Obj
obj) (Obj -> Direction
normObj Obj
obj) (Obj -> Set Shape -> RGB
brdf Obj
obj Set Shape
figuras)

    figuras' :: Set Shape
figuras' = (Shape -> Bool) -> Set Shape -> Set Shape
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\Shape
shape -> Obj -> Int
idObj Obj
obj Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Shape -> Int
getShapeID Shape
shape) Set Shape
figuras
    objEsp :: Obj
objEsp = Set Shape -> Direction -> Direction -> Point3D -> Obj
objEspejo Set Shape
figuras' (Obj -> Direction
w0Obj Obj
obj) (Obj -> Direction
normObj Obj
obj) (Obj -> Point3D
colObj Obj
obj)
    (Obj
objCr, Float
_) = Set Shape
-> Direction
-> Direction
-> Float
-> Float
-> Point3D
-> (Obj, Float)
objCristal Set Shape
figuras' (Obj -> Direction
w0Obj Obj
obj) (Obj -> Direction
normObj Obj
obj) Float
1 (Obj -> Float
reflObj Obj
obj) (Obj -> Point3D
colObj Obj
obj)

    difuso :: RGB
difuso = if Point3D -> Point3D -> Set Shape -> Bool
colision (Obj -> Point3D
colObj Obj
obj) Point3D
pointLuz Set Shape
figuras Bool -> Bool -> Bool
&& Float
kd Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
> Float
0 then RGB
newRGB RGB -> Float -> RGB
`modRGB` Float
kd else Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0
    espejo :: RGB
espejo = if Float
kr Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
> Float
0 then Obj -> Luz -> Set Shape -> RGB
luzMono Obj
objEsp (Point3D -> RGB -> Float -> Luz
Luz Point3D
pointLuz RGB
rgbLuz Float
intLuz) Set Shape
figuras RGB -> Float -> RGB
`modRGB` Float
kr else Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0
    cristal :: RGB
cristal = if Float
ke Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
> Float
0 then Obj -> Luz -> Set Shape -> RGB
luzMono Obj
objCr (Point3D -> RGB -> Float -> Luz
Luz Point3D
pointLuz RGB
rgbLuz Float
intLuz) Set Shape
figuras RGB -> Float -> RGB
`modRGB` Float
ke else Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0

    --iridiscencia = iris obj

    -- wH = normal (dirEspejo (w0Obj obj) (normObj obj) - w0Obj obj)
    -- fres = fresnell obj 1
    -- micro = microfacet wH (normObj obj) 0.4
    -- shadow = shadowing (w0Obj obj) (normObj obj) 0.4

-- | Función que calcula la luz indirecta de una escena.
{-# INLINE luzIndirecta #-}
luzIndirecta :: Obj -> [Luz] -> Set.Set Shape -> StdGen -> RGB
luzIndirecta :: Obj -> [Luz] -> Set Shape -> StdGen -> RGB
luzIndirecta Obj
obj [Luz]
luz Set Shape
figuras StdGen
gen = RGB
result where
    result :: RGB
result = case Int
caso of
        Int
0 -> Obj -> RGB -> RGB
rgbNew Obj
rndObj (Obj -> Set Shape -> RGB
brdf Obj
obj Set Shape
figuras) RGB -> Float -> RGB
`modRGB` Float
por
        Int
1 -> Obj -> RGB
colorIndirecto Obj
objCr RGB -> Float -> RGB
`modRGB` Float
por
        Int
2 -> Obj -> RGB
colorIndirecto Obj
objEsp RGB -> Float -> RGB
`modRGB` Float
por
        Int
_ -> Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0

    (Int
caso, Float
por) = (Float, Float, Float) -> StdGen -> (Int, Float)
ruletaRusa (Obj -> (Float, Float, Float)
trObj Obj
obj) StdGen
gen

    colorDirecto :: Obj -> RGB
colorDirecto Obj
nxtObj = [Luz] -> Set Shape -> Obj -> RGB
luzDirecta [Luz]
luz Set Shape
figuras Obj
nxtObj RGB -> Float -> RGB
`modRGB` Float -> Float
forall a. Num a => a -> a
abs ((Obj -> Direction
w0Obj Obj
nxtObj Direction -> Direction -> Float
.* Obj -> Direction
normObj Obj
obj) Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
forall a. Floating a => a
pi)

    colorIndirecto :: Obj -> RGB
colorIndirecto Obj
nxtObj = Obj -> [Luz] -> Set Shape -> StdGen -> RGB
luzIndirecta Obj
nxtObj [Luz]
luz Set Shape
figuras StdGen
gen'

    rgbNew :: Obj -> RGB -> RGB
rgbNew Obj
nxtObj = RGB -> Float -> Point3D -> Point3D -> Direction -> RGB -> RGB
formula (Obj -> RGB
colorDirecto Obj
nxtObj RGB -> RGB -> RGB
forall a. Num a => a -> a -> a
+ Obj -> RGB
colorIndirecto Obj
nxtObj) Float
1 (Obj -> Point3D
colObj Obj
obj) (Obj -> Point3D
colObj Obj
nxtObj) (Obj -> Direction
normObj Obj
nxtObj)

    figuras' :: Set Shape
figuras' = (Shape -> Bool) -> Set Shape -> Set Shape
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\Shape
shape -> Obj -> Int
idObj Obj
obj Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Shape -> Int
getShapeID Shape
shape) Set Shape
figuras
    rndObj :: Obj
rndObj = Set Shape -> Obj -> StdGen -> Obj
objAleatorio Set Shape
figuras' Obj
obj StdGen
gen

    gen' :: StdGen
gen' = (StdGen, StdGen) -> StdGen
forall a b. (a, b) -> b
snd (StdGen -> (StdGen, StdGen)
forall g. RandomGen g => g -> (g, g)
split StdGen
gen)

    (Obj
objCr, Float
rFlNew) = Set Shape
-> Direction
-> Direction
-> Float
-> Float
-> Point3D
-> (Obj, Float)
objCristal Set Shape
figuras' (Obj -> Direction
w0Obj Obj
obj) (Obj -> Direction
normObj Obj
obj) Float
1 (Obj -> Float
reflObj Obj
obj) (Obj -> Point3D
colObj Obj
obj)
    objEsp :: Obj
objEsp = Set Shape -> Direction -> Direction -> Point3D -> Obj
objEspejo Set Shape
figuras' (Obj -> Direction
w0Obj Obj
obj) (Obj -> Direction
normObj Obj
obj) (Obj -> Point3D
colObj Obj
obj) -- Multiplicarlo por el coseno de donde sale * 2pi por el monteCarlo

-- | Función principal que calcula la luz de una escena con una figura como luz de Area.
{-# INLINE luzArea #-}
luzArea :: Set.Set Shape -> Int -> Obj -> StdGen -> RGB
luzArea :: Set Shape -> Int -> Obj -> StdGen -> RGB
luzArea Set Shape
figuras Int
0 Obj
obj StdGen
gen = Set Shape -> Obj -> StdGen -> RGB
luzAreaRec Set Shape
figuras Obj
obj StdGen
gen
luzArea Set Shape
figuras Int
p Obj
obj StdGen
gen = Set Shape -> Int -> Obj -> StdGen -> RGB
luzArea Set Shape
figuras (Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Obj
obj StdGen
gen' RGB -> RGB -> RGB
forall a. Num a => a -> a -> a
+ Set Shape -> Obj -> StdGen -> RGB
luzAreaRec Set Shape
figuras Obj
obj StdGen
gen''
  where
    (StdGen
gen',StdGen
gen'') =StdGen -> (StdGen, StdGen)
forall g. RandomGen g => g -> (g, g)
split StdGen
gen

-- | Función que calcula la luz indirecta de una escena con una figura como luz de Area de manera recursiva.
{-# INLINE luzAreaRec #-}
luzAreaRec :: Set.Set Shape -> Obj -> StdGen -> RGB
luzAreaRec :: Set Shape -> Obj -> StdGen -> RGB
luzAreaRec Set Shape
figuras Obj
obj StdGen
gen = RGB
rgbFin
  where
    rgbFin :: RGB
rgbFin = if Obj -> Int
idObj Obj
obj Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4 then Float -> Float -> Float -> RGB
RGB Float
1 Float
1 Float
1 else if Obj -> RGB
rgbObj Obj
rndObj RGB -> RGB -> Bool
forall a. Eq a => a -> a -> Bool
== Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0 then Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0 else if RGB -> Bool
nanRGB RGB
result then Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0 else RGB
result
    result :: RGB
result = case Int
caso of
        Int
0 -> Obj -> RGB -> RGB
rgbNew Obj
rndObj (Obj -> Set Shape -> RGB
brdf Obj
obj Set Shape
figuras) RGB -> Float -> RGB
`modRGB`(Float
forall a. Floating a => a
pi Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
por)
        Int
1 -> Set Shape -> Obj -> StdGen -> RGB
luzAreaRec Set Shape
figuras Obj
objCr StdGen
gen' RGB -> Float -> RGB
`modRGB` Float
por
        Int
2 -> Set Shape -> Obj -> StdGen -> RGB
luzAreaRec Set Shape
figuras Obj
objEsp StdGen
gen' RGB -> Float -> RGB
`modRGB` Float
por
        Int
_ -> Float -> Float -> Float -> RGB
RGB Float
0 Float
0 Float
0
    
    (Int
caso, Float
por) = (Float, Float, Float) -> StdGen -> (Int, Float)
ruletaRusa (Obj -> (Float, Float, Float)
trObj Obj
obj) StdGen
gen

    rgbNew :: Obj -> RGB -> RGB
rgbNew Obj
nxtObj = RGB -> Float -> Point3D -> Point3D -> Direction -> RGB -> RGB
formula (Set Shape -> Obj -> StdGen -> RGB
luzAreaRec Set Shape
figuras Obj
nxtObj StdGen
gen') Float
1 (Obj -> Point3D
colObj Obj
obj) (Obj -> Point3D
colObj Obj
nxtObj) (Obj -> Direction
normObj Obj
nxtObj)

    figuras' :: Set Shape
figuras' = (Shape -> Bool) -> Set Shape -> Set Shape
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (\Shape
shape -> Obj -> Int
idObj Obj
obj Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Shape -> Int
getShapeID Shape
shape) Set Shape
figuras
    rndObj :: Obj
rndObj = Set Shape -> Obj -> StdGen -> Obj
objAleatorio Set Shape
figuras' Obj
obj StdGen
gen

    gen' :: StdGen
gen' = (StdGen, StdGen) -> StdGen
forall a b. (a, b) -> b
snd (StdGen -> (StdGen, StdGen)
forall g. RandomGen g => g -> (g, g)
split StdGen
gen)

    (Obj
objCr, Float
rFlNew) = Set Shape
-> Direction
-> Direction
-> Float
-> Float
-> Point3D
-> (Obj, Float)
objCristal Set Shape
figuras' (Obj -> Direction
w0Obj Obj
obj) (Obj -> Direction
normObj Obj
obj) Float
1 (Obj -> Float
reflObj Obj
obj) (Obj -> Point3D
colObj Obj
obj)
    objEsp :: Obj
objEsp = Set Shape -> Direction -> Direction -> Point3D -> Obj
objEspejo Set Shape
figuras' (Obj -> Direction
w0Obj Obj
obj) (Obj -> Direction
normObj Obj
obj) (Obj -> Point3D
colObj Obj
obj) -- Multiplicarlo por el coseno de donde sale * 2pi por el monteCarlo