IO.hs

module Waterfall.IO
( WaterfallIOException (..)
, WaterfallIOExceptionCause (..)
  -- * Solid Writers
, writeSolid
, writeSTL
, writeAsciiSTL
, writeSTEP
, writeGLTF
, writeGLB
, writeOBJ
  -- * Solid Readers
  -- 
  -- | Load a `Waterfall.Solid` from a file
  --
  -- At present, the "read*" functions do slightly less validation on the loaded solid 
  -- than they arguably should  
  -- and may succeed when reading solids that may generate invalid geometry when processed
, readSolid
, readSTL
, readSTEP
, readGLTF
, readGLB
, readOBJ
) where 

import Waterfall.Internal.Solid (Solid(..), PaintMap(..), emptyPaintMap, acquireSolid, paintMapIsEmpty)
import qualified Waterfall.Internal.Remesh as Remesh
import qualified OpenCascade.BRepMesh.IncrementalMesh as BRepMesh.IncrementalMesh
import qualified OpenCascade.StlAPI.Writer as StlWriter
import qualified OpenCascade.StlAPI.Reader as StlReader
import qualified OpenCascade.STEPControl.StepModelType as StepModelType
import qualified OpenCascade.STEPControl.Reader as STEPReader
import qualified OpenCascade.XSControl.Reader as XSControl.Reader
import qualified OpenCascade.IFSelect.ReturnStatus as IFSelect.ReturnStatus
import qualified OpenCascade.TDocStd.Document as TDocStd.Document
import qualified OpenCascade.Message.ProgressRange as Message.ProgressRange
import qualified OpenCascade.RWGltf.CafWriter as RWGltf.CafWriter
import qualified OpenCascade.NCollection.IndexedDataMap as NCollection.IndexedDataMap
import qualified OpenCascade.NCollection.Types as NCollection
import qualified OpenCascade.RWGltf.CafReader as RWGltf.CafReader
import qualified OpenCascade.RWObj.CafWriter as RWObj.CafWriter
import qualified OpenCascade.RWObj.CafReader as RWObj.CafReader
import qualified OpenCascade.RWMesh.Types as RWMesh
import qualified OpenCascade.RWMesh.CafReader as RWMesh.CafReader
import qualified OpenCascade.TDocStd.Types as TDocStd
import qualified OpenCascade.XCAFDoc.DocumentTool as XCafDoc.DocumentTool
import qualified OpenCascade.XCAFDoc.ShapeTool as XCafDoc.ShapeTool
import qualified OpenCascade.XCAFDoc.ColorType as XCAFDoc.ColorType
import qualified OpenCascade.TopoDS.Types as TopoDS
import qualified OpenCascade.TopoDS.Shape as TopoDS.Shape
import OpenCascade.Handle (Handle)
import OpenCascade.Inheritance (upcast)
import qualified OpenCascade.Quantity as Quantity
import qualified OpenCascade.Quantity.Color as Quantity.Color
import qualified OpenCascade.Quantity.ColorRGBA as Quantity.ColorRGBA
import qualified OpenCascade.Quantity.TypeOfColor as Quantity.TypeOfColor
import Control.Monad.IO.Class (liftIO)
import Control.Monad (unless, when, forM, forM_, (<=<))
import Waterfall.Internal.Finalizers (acquireForeignPtr, fromAcquirePtr)
import Data.Acquire ( Acquire, withAcquire )
import Foreign.Ptr (Ptr)
import Data.Char (toLower)
import System.FilePath (takeExtension)
import Control.Exception (Exception, throwIO)
import OpenCascade.TDF.Label (Label)
import Waterfall.Paint (Colour(..), paintColour, defaultPaint, Paint, DoubleSidedness(..), paintName, paintOpacity, paintEmissiveColour, paintMetallic, paintDoubleSidedness, paintRoughness)
import qualified OpenCascade.XCAFDoc.ColorTool as XCafDoc.ColourTool
import Control.Lens ((^.), Getting)
import System.IO (hPutStrLn, stderr)
import qualified OpenCascade.STEPCAFControl.Writer as StepCAFControl.Writer
import qualified OpenCascade.XCAFDoc.Types as XCafDoc
import qualified OpenCascade.XCAFDoc.VisMaterial as XCafDoc.VisMaterial
import qualified OpenCascade.XCAFDoc.VisMaterialPBR as XCafDoc.VisMaterialPBR
import qualified OpenCascade.NCollection.Vec3 as NCollection.Vec3
import Data.Maybe (fromMaybe)
import qualified OpenCascade.XCAFDoc.VisMaterialTool as XCafDoc.VisMaterialTool
import Control.Applicative ((<|>), Alternative)
import qualified OpenCascade.Graphic3D.AlphaMode as AlphaMode
import qualified OpenCascade.Graphic3D.TypeOfBackfacingModel as TypeOfBackfacingModel

import Data.Map (Map)
import qualified Data.Map as M

-- | The type of exceptions thrown by IO actions defined in `Waterfall.IO`
data WaterfallIOException = 
    WaterfallIOException 
      { ioExceptionCause :: WaterfallIOExceptionCause
      , ioExceptionFilePath :: FilePath 
      }
    deriving Show

instance Exception WaterfallIOException

-- | Reason for an IO action to have failed
data WaterfallIOExceptionCause = 
    -- | Something went wrong when accessing a file,
    -- eg. a write to a file path that is unreachable,
    -- or a read from a file in the wrong format 
    FileError  
    -- | The contents of a file could not be converted into a `Waterfall.Solid`
    -- e.g the file did not contain a solid object
    | BadGeometryError
    -- | The `readSolid`/`writeSolid` functions could not infer the correct file format from a filepath
    | UnrecognizedFormatError 
    deriving (Show, Eq)


extensionToFormats :: String -> Maybe (Double -> FilePath -> Solid -> IO(), FilePath -> IO Solid)
extensionToFormats s =
    let ext = fmap toLower . takeExtension $ s 
     in case ext of  
        ".stl" -> Just (writeSTL, readSTL)
        ".step" -> Just (const writeSTEP, readSTEP)
        ".gltf" -> Just (writeGLTF, readGLTF)
        ".glb" -> Just (writeGLB, readGLB)
        ".obj" -> Just (writeOBJ, readOBJ)
        _ -> Nothing

-- | Write a `Solid` to a file, work out the format from the file extension
-- 
-- Errors if passed a filename with an unrecognized extension
--
-- Because BRep representations of objects can store arbitrary precision curves,
-- but some of the supported file formats store triangulated surfaces, 
-- this function takes a "deflection" argument used to discretize curves.
--
-- The deflection is the maximum allowable distance between a curve and the generated triangulation.
writeSolid :: Double -> FilePath -> Solid -> IO ()
writeSolid res filepath = 
    case extensionToFormats filepath of
        Just (writer, _) -> writer res filepath
        Nothing -> const $ throwIO (WaterfallIOException UnrecognizedFormatError filepath)

writeSTLAsciiOrBinary :: Bool -> Double -> FilePath -> Solid -> IO ()
writeSTLAsciiOrBinary asciiMode linDeflection filepath (Solid ptr _paintMap) = (`withAcquire` pure) $ do
    s <- acquireForeignPtr ptr
    mesh linDeflection s
    writer <- StlWriter.new
    liftIO $ do
            StlWriter.setAsciiMode writer asciiMode
            res <- StlWriter.write writer s filepath
            unless res (throwIO (WaterfallIOException FileError filepath))
    return ()

-- | Write a `Solid` to a (binary) STL file at a given path
--
-- Because BRep representations of objects can store arbitrary precision curves,
-- but STL files store triangulated surfaces, 
-- this function takes a "deflection" argument used to discretize curves.
--
-- The deflection is the maximum allowable distance between a curve and the generated triangulation.
writeSTL :: Double -> FilePath -> Solid -> IO ()
writeSTL = writeSTLAsciiOrBinary False

-- | Write a `Solid` to an Ascii STL file at a given path
--
-- Because BRep representations of objects can store arbitrary precision curves,
-- but STL files store triangulated surfaces, 
-- this function takes a "deflection" argument used to discretize curves.
--
-- The deflection is the maximum allowable distance between a curve and the generated triangulation.
writeAsciiSTL :: Double -> FilePath -> Solid -> IO ()
writeAsciiSTL = writeSTLAsciiOrBinary True

-- | Write a `Solid` to a STEP file at a given path
--
-- STEP files can be imported by [FreeCAD](https://www.freecad.org/)
writeSTEP :: FilePath -> Solid -> IO ()
writeSTEP filepath (Solid ptr paintMap) = (`withAcquire` pure) $ do
    s <- acquireForeignPtr ptr
    writer <- StepCAFControl.Writer.new
    doc <- cafDocument s paintMap
    liftIO $ StepCAFControl.Writer.setColorMode writer (not . paintMapIsEmpty $ paintMap)
    resTransfer <- liftIO $ StepCAFControl.Writer.transfer writer doc StepModelType.Asls
    unless resTransfer (liftIO . throwIO $ WaterfallIOException BadGeometryError filepath)
    resWrite <- liftIO $ StepCAFControl.Writer.write writer filepath
    unless (resWrite == IFSelect.ReturnStatus.Done) (liftIO . throwIO $ WaterfallIOException FileError filepath)

colourToOCColor :: Colour -> Acquire (Ptr Quantity.Color)
colourToOCColor (Colour r g b) =
    Quantity.Color.new r g b Quantity.TypeOfColor.SRGB
    
colourToOCVec3 :: Colour -> Acquire (Ptr (NCollection.Vec3 Float))
colourToOCVec3 (Colour r g b) =
    NCollection.Vec3.newVec3Float (realToFrac r) (realToFrac g) (realToFrac b)

paintLookup :: Alternative f => Paint -> Getting (f a) Paint (f a) -> f a
paintLookup paint lens = paint ^. lens <|> defaultPaint ^. lens

paintToMaterial :: Ptr (Handle XCafDoc.VisMaterialTool) -> Paint ->  Acquire (Ptr Label)
paintToMaterial tool paint = do
    pbrMaterial <- XCafDoc.VisMaterialPBR.new
    material <- XCafDoc.VisMaterial.new

    let opacity = fromMaybe 1.0 (paintLookup paint paintOpacity)

    forM_ (paintLookup paint paintColour) 
        (liftIO . XCafDoc.VisMaterialPBR.setBaseColor pbrMaterial 
            <=< flip Quantity.ColorRGBA.new opacity
            <=< colourToOCColor
        )
    forM_ (paintLookup paint paintEmissiveColour) 
        (liftIO . XCafDoc.VisMaterialPBR.setEmissiveFactor pbrMaterial 
            <=< colourToOCVec3
        )
    forM_ (paintLookup paint paintMetallic)
        (liftIO . XCafDoc.VisMaterialPBR.setMetallic pbrMaterial) 

    forM_ (paintLookup paint paintRoughness)
        (liftIO . XCafDoc.VisMaterialPBR.setRoughness pbrMaterial) 
        
    forM_ (paintLookup paint paintName)
        (liftIO . XCafDoc.VisMaterial.setRawName material)

    let backfacing DoubleSided = TypeOfBackfacingModel.DoubleSided
        backfacing SingleSided = TypeOfBackfacingModel.BackCulled

    forM_ (paintLookup paint paintDoubleSidedness)
        (liftIO . XCafDoc.VisMaterial.setFaceCulling material . backfacing)
    
    liftIO $ XCafDoc.VisMaterial.setPbrMaterial material pbrMaterial
    let name = fromMaybe "" (paint ^. paintName)

    let alphaMode = if opacity < 1.0 then AlphaMode.BlendAuto else AlphaMode.Opaque
    liftIO $ XCafDoc.VisMaterial.setAlphaMode material alphaMode 0.5

    XCafDoc.VisMaterialTool.addMaterial tool material name

buildMaterialMap :: Ptr (Handle XCafDoc.VisMaterialTool) -> [Paint] ->  Acquire (Map Paint (Ptr Label))
buildMaterialMap tool paints = 
    let dup a = (a, a)
    in forM (M.fromList (dup <$> paints)) (paintToMaterial tool)

addColourToCafWriter :: Ptr TopoDS.Shape -> Ptr Label -> PaintMap -> Acquire ()
addColourToCafWriter s shapeLabel paintMap = do
    colourTool <- XCafDoc.DocumentTool.colorTool shapeLabel
    materialTool <- XCafDoc.DocumentTool.visMaterialTool shapeLabel
    case paintMap of
        UniformPaint paint -> do
            forM_ (paint ^. paintColour) $ \color -> do
                ocColour <- colourToOCColor color
                colourWasSet <- liftIO $ XCafDoc.ColourTool.setShapeColor colourTool s ocColour XCAFDoc.ColorType.ColorSurf
                liftIO . unless colourWasSet $ hPutStrLn stderr "Inconsistency: Base Shape not found in CAF document"

            material <- paintToMaterial materialTool paint
            liftIO $ XCafDoc.VisMaterialTool.setShapeMaterial materialTool shapeLabel material
        FacePaints facePaints -> do
            forM_ (defaultPaint ^. paintColour) $ \color -> do
                ocColour <- colourToOCColor color
                colourWasSet <- liftIO $ XCafDoc.ColourTool.setShapeColor colourTool s ocColour XCAFDoc.ColorType.ColorSurf
                liftIO . unless colourWasSet $ hPutStrLn stderr "Inconsistency: Base Shape not found in CAF document"
            materialMap <- buildMaterialMap materialTool (snd <$> facePaints)
            forM_ facePaints $ \(face, paint) -> do
                face' <- acquireForeignPtr face
                forM_ (paint ^. paintColour) $ \colour -> do
                        ocColour <- colourToOCColor colour
                        colourWasSet <- liftIO $ XCafDoc.ColourTool.setShapeColor colourTool (upcast face') ocColour XCAFDoc.ColorType.ColorSurf
                        liftIO . unless colourWasSet $ hPutStrLn stderr "Inconsistency: Face not found in CAF document"
                forM_ (M.lookup paint materialMap) 
                    (liftIO . XCafDoc.VisMaterialTool.setShapeMaterialFromShape materialTool (upcast face'))

mesh :: Double -> Ptr TopoDS.Shape -> Acquire ()
mesh linDeflection s = do
    mesher <- BRepMesh.IncrementalMesh.fromShapeAndLinDeflection s linDeflection
    liftIO $ BRepMesh.IncrementalMesh.perform mesher

cafDocument
    :: Ptr TopoDS.Shape 
    -> PaintMap 
    -> Acquire (Ptr (Handle TDocStd.Document))
cafDocument s paintMap = do
    doc <- TDocStd.Document.fromStorageFormat ""
    mainLabel <- TDocStd.Document.main doc
    shapeTool <- XCafDoc.DocumentTool.shapeTool mainLabel
    shapeLabel <- XCafDoc.ShapeTool.addShape shapeTool s True False
    addColourToCafWriter s shapeLabel paintMap
    return doc

writeGLTFOrGLB :: Bool -> Double -> FilePath -> Solid -> IO ()
writeGLTFOrGLB binary linDeflection filepath solid = (`withAcquire` pure) $  do
    writer <- RWGltf.CafWriter.new filepath binary
    meta <- NCollection.IndexedDataMap.newAsciiStringMap
    progress <- Message.ProgressRange.new
    s <- acquireSolid solid
    mesh linDeflection s
    doc <- cafDocument s (solidPaintMap solid)
    liftIO $ RWGltf.CafWriter.perform writer doc meta progress

-- | Write a `Solid` to a glTF file at a given path
--
-- glTF, or Graphics Library Transmission Format is a JSON based format 
-- used for three-dimensional scenes and models
--
-- Because BRep representations of objects can store arbitrary precision curves,
-- but glTF files store triangulated surfaces, 
-- this function takes a "deflection" argument used to discretize curves.
--
-- The deflection is the maximum allowable distance between a curve and the generated triangulation.
writeGLTF :: Double -> FilePath -> Solid -> IO ()
writeGLTF = writeGLTFOrGLB False

-- | Write a `Solid` to a glb file at a given path
--
-- glb is the binary variant of the glTF file format
--
-- Because BRep representations of objects can store arbitrary precision curves,
-- but glTF files store triangulated surfaces, 
-- this function takes a "deflection" argument used to discretize curves.
--
-- The deflection is the maximum allowable distance between a curve and the generated triangulation.
writeGLB :: Double -> FilePath -> Solid -> IO ()
writeGLB = writeGLTFOrGLB True

-- | Write a `Solid` to an obj file at a given path
--
-- Wavefront OBJ is a simple ascii file format that stores geometric data.
--
-- Because BRep representations of objects can store arbitrary precision curves,
-- but obj files store triangulated surfaces, 
-- this function takes a "deflection" argument used to discretize curves.
--
-- The deflection is the maximum allowable distance between a curve and the generated triangulation.
writeOBJ :: Double -> FilePath -> Solid -> IO ()
writeOBJ linDeflection filepath solid = (`withAcquire` pure) $ do
    writer <- RWObj.CafWriter.new filepath
    meta <- NCollection.IndexedDataMap.newAsciiStringMap
    progress <- Message.ProgressRange.new
    s <- acquireSolid solid
    mesh linDeflection s
    doc <- cafDocument s (solidPaintMap solid)
    liftIO $ RWObj.CafWriter.perform writer doc meta progress

-- | Read a `Solid` from a file at a given path
-- 
-- Throws an error if loading fails, or if it's unable to work out
-- the intended file format from the path
readSolid :: FilePath -> IO Solid
readSolid filepath = 
    case extensionToFormats filepath of 
        Nothing -> throwIO (WaterfallIOException UnrecognizedFormatError filepath)
        Just (_, reader) -> reader filepath

remeshOrThrow :: FilePath -> Ptr TopoDS.Shape -> Acquire (Ptr TopoDS.Shape)
remeshOrThrow filepath shape = do
    remeshed <- Remesh.remesh shape
    case remeshed of 
        Just solid -> pure solid
        Nothing -> liftIO . throwIO $ WaterfallIOException BadGeometryError filepath

-- | Read a `Solid` from an STL file at a given path
readSTL :: FilePath -> IO Solid
readSTL filepath = fmap (`Solid` emptyPaintMap) . fromAcquirePtr $ do
    shape <- TopoDS.Shape.new
    reader <- StlReader.new
    res <- liftIO $ StlReader.read reader shape filepath
    unless res $ liftIO . throwIO $ WaterfallIOException FileError filepath
    remeshOrThrow filepath shape

-- | Read a `Solid` from a STEP file at a given path
readSTEP :: FilePath -> IO Solid
readSTEP filepath = fmap (`Solid` emptyPaintMap) . fromAcquirePtr $ do
    reader <- STEPReader.new
    status <- liftIO $ XSControl.Reader.readFile (upcast reader) filepath
    _ <- liftIO $ XSControl.Reader.transferRoots (upcast reader)
    shape <- XSControl.Reader.oneShape (upcast reader)
    unless (status == IFSelect.ReturnStatus.Done) (liftIO . throwIO $ WaterfallIOException FileError filepath)
    shapeIsNull <- liftIO $ TopoDS.Shape.isNull shape
    when shapeIsNull (liftIO . throwIO $ WaterfallIOException BadGeometryError filepath)
    return shape

cafReader :: Acquire (Ptr RWMesh.CafReader) -> FilePath -> IO Solid
cafReader mkReader filepath = fmap (`Solid` emptyPaintMap) . fromAcquirePtr $ do
    reader <- mkReader
    doc <- TDocStd.Document.fromStorageFormat ""
    progress <- Message.ProgressRange.new
    _ <- liftIO $ RWMesh.CafReader.setDocument reader doc
    res <- liftIO $ RWMesh.CafReader.perform reader filepath progress
    unless res (liftIO . throwIO $ WaterfallIOException FileError filepath)
    remeshOrThrow filepath =<< RWMesh.CafReader.singleShape reader

-- | Read a `Solid` from a GLTF file at a given path
--
-- This should support reading both the GLTF (json) and GLB (binary) formats
readGLTF :: FilePath -> IO Solid
readGLTF  = cafReader $ do
    reader <- RWGltf.CafReader.new 
    liftIO $ RWGltf.CafReader.setDoublePrecision reader True
    liftIO $ RWMesh.CafReader.setFileLengthUnit (upcast reader) 1
    return (upcast reader)

-- | Alias for `readGLTF`
readGLB :: FilePath -> IO Solid
readGLB = readGLTF

-- | Read a `Solid` from an obj file at a given path
readOBJ :: FilePath -> IO Solid
readOBJ  = cafReader $ do
    reader <- RWObj.CafReader.new 
    liftIO $ RWObj.CafReader.setSinglePrecision reader False
    liftIO $ RWMesh.CafReader.setFileLengthUnit (upcast reader) 1
    return (upcast reader)