module Shape.Sphere where

import Core.Constants
import Core.Types
import Core.Transform
import Core.Geometry
import Core.Intersection
import Core.DifferentialGeometry

import Data.Maybe (isJust)

type Radius = DFloat

data Sphere = Sphere {sphereToWorld  :: Transform,
                      worldToSphere  :: Transform,

                      radius         :: Radius,
                      zBounds        :: Range,
                      thetaBounds    :: Range,
                      phiMax         :: Radians,

                      objectBBox     :: BBox,
                      worldBBox      :: BBox,

                      revNormals     :: Bool,
                      revOrientation :: Bool,
                      swapHands      :: Bool}

instance Shape Sphere where
  objectToWorld      = sphereToWorld
  worldToObject      = worldToSphere
  objectBound        = objectBBox
  worldBound         = worldBBox
  reverseNormals     = revNormals
  reverseOrientation = revOrientation
  swapsHandedness    = swapHands
  area               = sphereArea

instance Intersectable Sphere where
  intersect  = sphereIntersect
  intersectP = sphereIntersectP

sphere :: Transform             -- | Object to world
          -> Transform          -- | World to object
          -> Bool               -- | Reverse orientation?
          -> Radius             -- | Radius
          -> DFloat             -- | Z min
          -> DFloat             -- | Z max
          -> Degrees            -- | Phi max
          -> Sphere
sphere o2w w2o ro rad zMin zMax phiMax = Sphere {sphereToWorld  = o2w,
                                                 worldToSphere  = w2o,
                                                 radius         = rad,
                                                 zBounds        = (zMin', zMax'),
                                                 thetaBounds    = (tMin', tMax'),
                                                 phiMax         = pMax',
                                                 objectBBox     = oBox,
                                                 worldBBox      = wBox,
                                                 revNormals     = False,
                                                 revOrientation = ro,
                                                 swapHands      = undefined}
  where zMin' = clamp (-rad, rad) (min zMin zMax)
        zMax' = clamp (-rad, rad) (max zMin zMax)
        tMin' = acos (clamp (-1,1) (zMin/rad))
        tMax' = acos (clamp (-1,1) (zMax/rad))
        pMax' = toRadians (clamp (0,360) phiMax)
        oBox = BBox (Point (-rad) (-rad) zMin') (Point rad rad zMax')
        wBox = o2w `appTrans` oBox

sphereIntersect :: Sphere -> Ray -> Maybe Intersection
sphereIntersect sphere worldRay = sols >>= processSols sphere ray
  where ray     = (worldToSphere sphere) `appTrans` worldRay
        rad     = radius sphere
        a       = normsq (direction ray)
        origin' = vecOfPoint $ origin ray
        b       = 2 * (direction ray `dot` origin')
        c       = normsq origin' - rad*rad
        sols    = solveQuadratic (a,b,c)

-- TODO: Make this more efficient
sphereIntersectP :: Sphere -> Ray -> Bool
sphereIntersectP sphere = isJust . sphereIntersect sphere

processSols :: Sphere -> Ray -> QuadraticSol -> Maybe Intersection
processSols sphere ray (t0, t1) |  t0 > rMinT || t1 < rMinT
                                || t0 < rMinT && t1 > rMaxT = Nothing
                                | clipped   = processClipped sphere ray tHit (t0,t1)
                                | otherwise = succeed
  where (rMinT, rMaxT) = bounds ray
        tHit = if t0 < rMinT then t1 else t0
        p@(Point px' py' pz') = stepRay ray tHit
        rad = radius sphere
        biasX = 1e-5 * rad
        pHit@(Point px py pz) = if px' == 0 && py' == 0
                                then Point biasX py' pz'
                                else p
        phi' = atan2 py px
        phi  = if phi' < 0 then phi' + 2*pi else phi

        (zMin, zMax) = zBounds sphere
        pMax = phiMax sphere
        clipped = zMin > -rad && pz < zMin || zMax < rad && pz > zMax || phi > pMax

        tHit' = if clipped then t1 else tHit
        pHit' = if clipped then stepRay ray tHit' else pHit

        succeed = processDiffGeom sphere ray pHit' phi tHit'

processClipped :: Sphere -> Ray -> DFloat -> QuadraticSol -> Maybe Intersection
processClipped sphere ray tHit (t0, t1) | tHit == t1  = Nothing
                                        | t1 > rMaxT  = Nothing
                                        | outOfBounds = Nothing
                                        | otherwise   = succeed
  where (rMinT, rMaxT) = bounds ray
        (zMin, zMax)   = zBounds sphere
        rad = radius sphere
        checkIntPoint p@(Point x y z) | x == 0 && y == 0 = Point (1e-5 * rad) 0 z
                                      | otherwise        = p

        pHit@(Point x y z) = checkIntPoint $ stepRay ray t1
        phi = let phi' = atan2 y z
              in if phi' < 0 then phi' + 2*pi else phi'

        outOfBounds = (zMin > -rad && z < zMin) ||
                      (zMax < rad && z > zMax) ||
                      (phi > phiMax sphere)

        succeed = processDiffGeom sphere ray pHit phi tHit

processDiffGeom :: Sphere                   -- | Sphere intersected
                   -> Ray                   -- | Ray cast
                   -> Point                 -- | Intersection point (phit)
                   -> DFloat                -- | Intersection angle (phi)
                   -> DFloat                -- | Parameter of intersection (thit)
                   -> Maybe Intersection
processDiffGeom sphere ray phit phi thit = Just res
  where (Point pHitX pHitY pHitZ) = phit
        (thetaMin, thetaMax) = thetaBounds sphere
        thetaRange = thetaMax - thetaMin
        rad = radius sphere
        pMax = phiMax sphere

        u     = phi / pMax
        theta = acos(clamp (-1, 1) (pHitZ / rad))
        v     = (theta - thetaMin) / thetaRange

        zRadius = sqrt (pHitX*pHitX + pHitY*pHitY)
        invZRadius = 1/zRadius
        cosPhi = pHitX * invZRadius
        sinPhi = pHitY * invZRadius
        dpdu = Vector (-pMax * pHitY) (pMax * pHitX) 0
        dpdv = thetaRange *^ Vector (pHitZ * cosPhi) (pHitZ * sinPhi) (-rad * sin theta)

        d2Pduu = (-pMax * pMax) *^ Vector pHitX pHitY 0
        d2Pduv = (thetaRange * pHitZ * pMax) *^ Vector (-sinPhi) cosPhi 0
        d2Pdvv = (-thetaRange * thetaRange) *^ Vector pHitX pHitY pHitZ

        ffE = dpdu `dot` dpdu
        ffF = dpdu `dot` dpdv
        ffG = dpdv `dot` dpdv
        ffN = hat $ dpdu `cross` dpdv
        ffe = ffN `dot` d2Pduu
        fff = ffN `dot` d2Pduv
        ffg = ffN `dot` d2Pdvv

        invEGF2 = 1 / (ffE*ffG - ffF*ffF)
        dndu = normalOfVec $ (((fff*ffF - ffe*ffG) * invEGF2) *^ dpdu) ^+^
                             (((ffe*ffF - fff*ffE) * invEGF2) *^ dpdv)
        dndv = normalOfVec $ (((ffg*ffF - fff*ffG) * invEGF2) *^ dpdu) ^+^
                             (((fff*ffF - ffg*ffE) * invEGF2) *^ dpdv)

        dg = defaultDiffGeom {p = phit,
                              u = u,
                              v = v,
                              dpdu = dpdu,
                              dpdv = dpdv,
                              dndu = dndu,
                              dndv = dndv,
                              shape = Intersectable sphere}

        res = Intersection {tHit         = thit,
                            diffGeom     = computeNormalFor dg,
                            worldToInter = worldToObject sphere,
                            interToWorld = objectToWorld sphere,
                            rayEpsilon   = thit * 5e-4}

sphereArea :: Sphere -> DFloat
sphereArea s = (phiMax s) * (radius s) * (zMax-zMin)
  where (zMin, zMax) = (zBounds s)

sphereSample :: Sphere -> Point -> Range -> Normal -> Point
sphereSample sphere p (u1,u2) ns | inside    = uniformSample (u1,u2) ns
                                 | otherwise = sampleCone sphere p pCenter (u1,u2) ns
  where pCenter   = objectToWorld sphere `appTrans` (Point 0 0 0)
        wc        = hat (pCenter .-. p)
        (wcX,wcY) = coordinateSystem wc
        rad       = radius sphere

        inside = distsq p pCenter - rad*rad < 1e-4

sampleCone = undefined
-- sampleCone :: Sphere -> Point -> Point -> Range -> Normal -> Point
-- sampleCone sphere p (u1,u2) ns = undefined
--   where rad = radius sphere
--         pCenter = objectToWorld sphere `appTrans` (Point 0 0 0)
--         sinThetaMax2 = rad*rad / (distsq p pCenter)
--         cosThetaMax  = sqrt (max 0 (1 - sinThetaMax2))
        -- TODO: Finish this

uniformSample = undefined -- TODO: Remove this

spherePDF :: Sphere -> Point -> Vector -> DFloat
spherePDF = undefined
