{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
module LoadMesh where

import LambdaCube.GL as LambdaCubeGL -- renderer
import LambdaCube.GL.Mesh as LambdaCubeGL
import MtlParser

import Control.Monad
import Data.Maybe
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Vector as V
import qualified Data.ByteString as SB
import qualified Data.ByteString.Lazy.Char8 as L
import Data.Text (unpack,Text)
import Data.List (groupBy,nub)
import Numeric.LinearAlgebra hiding ((<>))

import Codec.Picture as Juicy
import Wavefront
import Wavefront.Types
import Data.Aeson

type MeshData = ( [(Mesh,Maybe Text)]  -- List of uniform-material meshes (and the name of the material).
                , MtlLib               -- Material definitions.
                )

loadOBJ :: String -> IO (Either String MeshData)
loadOBJ fname = L.readFile fname >>= \bs -> do
    let  obj@OBJ{..} = parse bs
    -- load materials
    mtlLib <- mconcat . V.toList <$> mapM (readMtl . unpack) objMtlLibs
    return $ Right (objToMesh obj,mtlLib)


data BoundingBox = BoundingBox
    { minX :: Float
    , maxX :: Float
    , minY :: Float
    , maxY :: Float
    , minZ :: Float
    , maxZ :: Float
    }
 deriving (Eq,Ord,Show)

instance Semigroup BoundingBox where
    a <> b = BoundingBox
        { minX = if minX b < minX a then minX b else minX a
        , maxX = if maxX b > maxX a then maxX b else maxX a
        , minY = if minY b < minY a then minY b else minY a
        , maxY = if maxY b > maxY a then maxY b else maxY a
        , minZ = if minZ b < minZ a then minZ b else minZ a
        , maxZ = if maxZ b > maxZ a then maxZ b else maxZ a
        }
instance Monoid BoundingBox where mempty = BoundingBox 0 0 0 0 0 0

attribBoundingBox :: Map String MeshAttribute -> BoundingBox
attribBoundingBox attrib = case Map.lookup "position" attrib of
        Just (A_V3F vs) -> V.foldr (\(V3 x y z  ) bb -> bb <> BoundingBox x x y y z z) mempty vs
        Just (A_V4F vs) -> V.foldr (\(V4 x y z _) bb -> bb <> BoundingBox x x y y z z) mempty vs
        _               -> mempty

bbnorm :: BoundingBox -> Float
bbnorm (BoundingBox x0 x y0 y z0 z) = sqrt $ (x-x0)^2 + (y-y0)^2 + (z-z0)^2

scaleWithin :: BoundingBox -> BoundingBox -> Matrix Float
scaleWithin meshbb scalebb =
    if meshbb <> scalebb /= scalebb
        || (let {m=bbnorm meshbb; s=bbnorm scalebb} in m < 0.1*s)
        then let tr0 = (4><4) [ 1,0,0, negate $ (minX meshbb + maxX meshbb)/2
                              , 0,1,0, negate $ (minY meshbb + maxY meshbb)/2
                              , 0,0,1, negate $ (minZ meshbb + maxZ meshbb)/2
                              , 0,0,0, 1 ]
                 sc = (4><4) [s,0,0,0
                             ,0,s,0,0
                             ,0,0,s,0
                             ,0,0,0,1]
                 s = minimum [sx,sy,sz]
                 sx = (maxX scalebb - minX scalebb) / (maxX meshbb - minX meshbb)
                 sy = (maxY scalebb - minY scalebb) / (maxY meshbb - minY meshbb)
                 sz = (maxZ scalebb - minZ scalebb) / (maxZ meshbb - minZ meshbb)
                 tr1 = (4><4) [ 1,0,0, (minX scalebb + maxX scalebb)/2
                              , 0,1,0, (minY scalebb + maxY scalebb)/2
                              , 0,0,1, (minZ scalebb + maxZ scalebb)/2
                              , 0,0,0, 1 ]
             in tr1 <> sc <> tr0
        else ident 4

transV3 t (V3 x y z) = let v = t #> fromList [x,y,z,1] in V3 (v!0/v!3) (v!1/v!3) (v!2/v!3)
transV4 t (V4 x y z w) = let v = t #> fromList [x,y,z,w] in V4 (v!0) (v!1) (v!2) (v!3)

tranformAttribute t (A_V3F v) = A_V3F $ transV3 t <$> v
tranformAttribute t (A_V4F v) = A_V4F $ transV4 t <$> v

transformMesh :: Matrix Float -> Mesh -> Mesh
transformMesh t m = m
    { mAttributes = Map.adjust (tranformAttribute t) "position" (mAttributes m)
    }

uploadOBJToGPU :: Maybe BoundingBox -> MeshData -> IO (([(GPUMesh, Maybe Text)], MtlLib),Matrix Float)
uploadOBJToGPU scalebb (subModels,mtlLib) = do
    let meshbb = foldMap (attribBoundingBox . mAttributes . fst) subModels :: BoundingBox
        m = maybe (ident 4) (scaleWithin meshbb) scalebb
    gpuSubModels <- forM subModels $ \(mesh,mat) -> LambdaCubeGL.uploadMeshToGPU (transformMesh m mesh) >>= \a -> return (a,mat)
    return ((gpuSubModels,mtlLib),m)

uploadMtlLib :: MtlLib -> IO (Map Text (ObjMaterial,TextureData))
uploadMtlLib mtlLib = do
  -- collect used textures
  let usedTextures = nub . concatMap (maybeToList . mtl_map_Kd) $ Map.elems mtlLib
      whiteImage = Juicy.ImageRGB8 $ Juicy.generateImage (\_ _ -> Juicy.PixelRGB8 255 255 255) 1 1
      checkerImage = Juicy.ImageRGB8 $ Juicy.generateImage (\x y -> if mod (x + y) 2 == 0 then Juicy.PixelRGB8 0 0 0 else Juicy.PixelRGB8 255 255 0) 2 2
  checkerTex <- LambdaCubeGL.uploadTexture2DToGPU checkerImage
  -- load images and upload to gpu
  textureLib <- forM (Map.fromList $ zip usedTextures usedTextures) $ \fname -> Juicy.readImage fname >>= \case
    Left err  -> putStrLn (fname ++": "++err) >> return checkerTex
    Right img -> LambdaCubeGL.uploadTexture2DToGPU img
  whiteTex <- LambdaCubeGL.uploadTexture2DToGPU whiteImage
  -- pair textures and materials
  return $ (\a -> (a, maybe whiteTex (fromMaybe checkerTex . flip Map.lookup textureLib) . mtl_map_Kd $ a)) <$> mtlLib

objToMesh :: WavefrontOBJ -> [(Mesh,Maybe Text)]
objToMesh OBJ{..} = [(toMesh faceGroup, elMtl . head $ faceGroup) | faceGroup <- faces] where
  faces = groupBy (\a b -> elMtl a == elMtl b) (V.toList objFaces)
  toMesh l = Mesh
    { mAttributes   = Map.fromList
        [ ("position",  A_V4F position)
        , ("normal",    A_V3F normal)
        , ("uvw",       A_V3F texcoord)
        ]
    , mPrimitive    = P_Triangles
    } where
        triangulate (Triangle a b c) = [a,b,c]
        triangulate (Quad a b c d) = [a,b,c, c,d,a]
        triangulate (Face a b c l) = a : b : c : concatMap (\(x,y) -> [a,x,y]) (zip (c:l) l) -- should work for convex polygons without holes
        defaultPosition = Location 0 0 0 0
        defaultNormal = Normal 0 0 0
        defaultTexCoord = TexCoord 0 0 0
        v !- i = v V.!? i
        toVertex FaceIndex{..} = ( let Location x y z w = fromMaybe defaultPosition (objLocations !- faceLocIndex)            in V4 x y z w
                                 , let Normal x y z     = fromMaybe defaultNormal ((objNormals !-) =<< faceNorIndex)          in V3 x y z
                                 , let TexCoord x y z   = fromMaybe defaultTexCoord ((objTexCoords !-) =<< faceTexCoordIndex) in V3 x y z
                                 )
        (position,normal,texcoord) = V.unzip3 . V.concat . map (V.fromList . map toVertex . triangulate . elValue) $ l