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

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

import Control.Arrow
import Control.Monad
import Data.List as List
import Data.Maybe
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Vector as V
import qualified Data.Vector.Storable as StorableV
import qualified Data.ByteString.Lazy.Char8 as L
import Data.Text (unpack,Text,pack)
import Data.List (groupBy,nub)
import Numeric.LinearAlgebra hiding ((<>),Element)
import System.FilePath

import Codec.Picture as Juicy
import Wavefront
import Wavefront.Types
import Mask

data MaterialMesh m = MaterialMesh
    { materialMesh  :: m
    , materialName  :: Maybe (Int,Text)
    , materialMasks :: Map Text Mask
    }

data CurveData = CurveData
    { curves   :: [Curve]
    , curvePt  :: Int -> Location
    , curveMax :: Int
    }


data MeshData = MeshData
    { matMeshes :: [MaterialMesh Mesh]  -- List of uniform-material meshes (and the name of the material).
    , matLib    :: ( V.Vector MtlLib    -- Material definitions.
                   , FilePath )         -- Path to wavefront obj file.
    , matCurves :: CurveData
    }

relativeFrom :: FilePath -> FilePath -> FilePath
relativeFrom path file | isAbsolute file = file
relativeFrom path file = takeDirectory path </> file

loadOBJ :: String -> IO (Either String MeshData)
loadOBJ fname = L.readFile fname >>= \bs -> do
    let  obj@OBJ{..} = Wavefront.parse bs
    -- load materials
    mtlLib <- if V.null objMtlLibs
                then return $ V.singleton (Map.singleton "" $ newMaterial "")
                else mapM (readMtlWithFallback . relativeFrom fname . unpack) objMtlLibs
    return $ Right MeshData
        { matMeshes = objToMesh obj
        , matLib    = (mtlLib,fname)
        , matCurves = objToCurveData obj
        }



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 :: Matrix Float -> V3 Float -> V3 Float
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 :: Matrix Float -> V4 Float -> V4 Float
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 :: Matrix Float -> MeshAttribute -> MeshAttribute
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)
    }

transformLocation :: Matrix Float -> Location -> Location
transformLocation t (Location x y z w) = Location xx yy zz ww
 where
    [xx,yy,zz,ww] = toList $ t #> fromList [x,y,z,w]

locationBoundingBox :: Location -> BoundingBox
locationBoundingBox (Location x y z w) = BoundingBox x x y y z z

uploadOBJToGPU :: Maybe BoundingBox -> MeshData -> IO (([MaterialMesh GPUMesh],CurveData),Matrix Float)
uploadOBJToGPU scalebb (MeshData subModels (mtlLib,objpath) curveData) = do
    let meshbb = foldMap (attribBoundingBox . mAttributes . materialMesh) subModels
                 <> foldMap (foldMap (locationBoundingBox . curvePt curveData) . curvePoints)
                            (curves curveData)
        m = maybe (ident 4) (scaleWithin meshbb) scalebb
        curveData' = case scalebb of
            Just _  -> curveData { curvePt = transformLocation m . curvePt curveData }
            Nothing -> curveData
    putStrLn $ show meshbb
    gpuSubModels <- forM subModels $ \matmesh -> do
        a <- LambdaCubeGL.uploadMeshToGPU (transformMesh m (materialMesh matmesh))
        return matmesh { materialMesh = a }
    return ((gpuSubModels,curveData'),m)

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

vecLocation :: Location -> StorableV.Vector Float
vecLocation (Location x y z w) = StorableV.fromList [x,y,z,w]

objToCurveData :: WavefrontOBJ -> CurveData
objToCurveData OBJ{..} = CurveData
    { curves   = map elValue $ V.toList $ objCurves
    , curvePt  = (objLocations V.!)
    , curveMax = V.length objLocations
    }

faceNormal :: [Location] -> Normal
faceNormal (Location ax ay az _:Location bx by bz _:Location cx cy cz _:_)
    = Normal nx ny nz
 where
    [nx,ny,nz] = toList $ nrml $ cross x y
    a = fromList [ax,ay,az]
    b = fromList [bx,by,bz]
    c = fromList [cx,cy,cz]
    x = b - a
    y = c - b
faceNormal _ = Normal 0 1 0

nrml :: (Linear t c, Fractional t, Normed (c t)) => c t -> c t
nrml v = scale (1 / realToFrac (norm_2 v)) v


objToMesh :: WavefrontOBJ -> [MaterialMesh Mesh]
objToMesh OBJ{..} = [ toMesh faceGroup | faceGroup <- faces ]
 where
  faces = groupBy (\a b -> elMtl a == elMtl b) (V.toList objFaces)
  toMesh l = MaterialMesh mesh mtl gs
     where
        mtl = elMtl $ head l
        computeNormal fs = case 1 {- mtl_illum (lookupMat mtl) -} of
            -- 0 -> Normal 0 1 0
            _ -> faceNormal $ mapMaybe ((objLocations !-) . faceLocIndex) fs
        mesh = Mesh { mAttributes   = Map.fromList
                        [ ("position",  A_V4F position)
                        , ("normal",    A_V3F normal)
                        , ("uvw",       A_V3F texcoord)
                        ]
                    , mPrimitive    = P_Triangles
                    }
        defaultPosition = Location 0 0 0 0
        defaultTexCoord = TexCoord 0 0 0
        v !- i = v V.!? i
        toVertex defaultNormal 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
            )
        (gs,fss) = elementIndices l
        (positions,normals,texcoords) = unzip3 $ concatMap (\fs -> map (toVertex $ computeNormal fs) fs) fss
        position = V.fromList positions
        normal = V.fromList normals
        texcoord = V.fromList texcoords

triangulate :: Face -> [FaceIndex]
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

elementIndices :: [Element Face] -> (Map Text Mask, [[FaceIndex]])
elementIndices els = (spans,map snd ts)
 where
    ts = map ((elGroups &&& elValue) . fmap triangulate) els
    rs = List.scanl' go ((0,0),[]) ts -- scanl :: (b -> a -> b) -> b -> [a] -> [b]
    go ((start,len),_) (gs,vs) = ((start+len,length vs),gs)
    spans = fmap (foldr (maskPlus . Mask . (:[])) (Mask [])
                  .  map (fromIntegral***fromIntegral))
                $ foldr (Map.unionWith (++)) Map.empty $ map singletonSpan rs

    singletonSpan (span,gnames) = Map.fromList $ map (, [span]) gnames


data MaskableObject = MaskableObject
    { maskableObject :: LC.Object
    , groupMasks :: Map Text Mask
    }

objSpan :: LC.Object -> Mask
objSpan obj = case Map.elems (objAttributes obj) of
        Stream { streamLength = x }:_ -> Mask [(0,fromIntegral x)]
        _                             -> Mask [(0,1)]


searchMaterial
    :: V.Vector (Map Text (ObjMaterial, TextureData)) -- ^ Some tail end of this vector will be searched.
    -> (Int, Text)                                    -- ^ Size of tail and material name to search for.
    -> Maybe (ObjMaterial, TextureData)
searchMaterial mtlLib (count,name) = foldr go id (V.drop (V.length mtlLib - count) mtlLib) Nothing
 where
    go m f r = case Map.lookup name m of
        Nothing -> f r
        x       -> x :: Maybe (ObjMaterial,TextureData)

addOBJToObjectArray :: GLStorage -> String
                                 -> [MaterialMesh GPUMesh]
                                 -> V.Vector (Map Text (ObjMaterial,TextureData))
                                 -> IO [MaskableObject]
addOBJToObjectArray storage slotName objMesh mtlLib = forM objMesh $ \matmesh -> do
  obj <- case materialName matmesh >>= searchMaterial mtlLib of
    Nothing -> do
        let slotnm = slotName ++ "1"
        obj <- LambdaCubeGL.addMeshToObjectArray storage
                                                 slotnm
                                                 [ "diffuseTexture"
                                                 , "diffuseColor"
                                                 , "specularReflectivity"]
                                                 (materialMesh matmesh)
        let (white,whiteTex) = (mtlLib V.! 0) Map.! ""
        LC.updateObjectUniforms obj $ do
          "diffuseTexture" @= return whiteTex -- set model's diffuse texture
          "diffuseColor" @= let (r,g,b) = mtl_Kd white in return $ V4 r g b (mtl_Tr white)
          "specularReflectivity" @= let (r,g,b) = mtl_Ks white in return $ V4 r g b (mtl_Ns white)
        return obj
    Just (ObjMaterial{..},t) -> do
        let slotnm = slotName ++ show (if 0 <= mtl_illum && mtl_illum <= 2 then mtl_illum else 2)
        obj <- LambdaCubeGL.addMeshToObjectArray storage
                                                 slotnm
                                                 [ "diffuseTexture"
                                                 , "diffuseColor"
                                                 , "specularReflectivity"]
                                                 (materialMesh matmesh)
        LC.updateObjectUniforms obj $ do
          "diffuseTexture" @= return t -- set model's diffuse texture
          "diffuseColor" @= let (r,g,b) = mtl_Kd in return (V4 r g b mtl_Tr)
          "specularReflectivity" @= let (r,g,b) = mtl_Ks in return (V4 r g b mtl_Ns)
        return obj
  let matmask = maybe Map.empty (`Map.singleton` objSpan obj)
                      (fmap (\(c,n) -> "m:" <> pack (show c) <> ":" <> n) $ materialName matmesh)
  return $ MaskableObject obj (matmask `Map.union` materialMasks matmesh)