content: {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  import Data.Array
  import Data.Foldable (length, concat, sum)
  import Data.List (genericLength)
  import Data.Maybe (mapMaybe)
  import Prelude hiding (length, concat, sum)
  import qualified Data.Map.Strict as Map
  
  {-
  A SimpleMesh consists of only vertices and faces that refer to them.
  A Mesh extends the SimpleMesh to contain edges as well as references to
  adjoining mesh components for each other component, such as a vertex
  also contains what faces it belongs to.
  An isolated edge can be represented as a degenerate face with 2 vertices.
  Faces with 0 or 1 vertices can be thrown out, as they do not contribute to
  the result (they can also propagate NaNs).
  -}
  
  newtype VertexId = VertexId { getVertexId :: Int } deriving (Ix, Ord, Eq, Show)
  newtype EdgeId = EdgeId { getEdgeId :: Int } deriving (Ix, Ord, Eq, Show)
  newtype FaceId = FaceId { getFaceId :: Int } deriving (Ix, Ord, Eq, Show)
  
  data Vertex a = Vertex
    { vertexPoint :: a
    , vertexEdges :: [EdgeId]
    , vertexFaces :: [FaceId]
    } deriving Show
  
  data Edge = Edge
    { edgeVertexA :: VertexId
    , edgeVertexB :: VertexId
    , edgeFaces :: [FaceId]
    } deriving Show
  
  data Face = Face
    { faceVertices :: [VertexId]
    , faceEdges :: [EdgeId]
    } deriving Show
  
  type VertexArray a = Array VertexId (Vertex a)
  type EdgeArray = Array EdgeId Edge
  type FaceArray = Array FaceId Face
  
  data Mesh a = Mesh
    { meshVertices :: VertexArray a
    , meshEdges :: EdgeArray
    , meshFaces :: FaceArray
    } deriving Show
  
  data SimpleVertex a = SimpleVertex { sVertexPoint :: a } deriving Show
  data SimpleFace = SimpleFace { sFaceVertices :: [VertexId] } deriving Show
  
  type SimpleVertexArray a = Array VertexId (SimpleVertex a)
  type SimpleFaceArray = Array FaceId SimpleFace
  
  data SimpleMesh a = SimpleMesh
    { sMeshVertices :: SimpleVertexArray a
    , sMeshFaces :: SimpleFaceArray
    } deriving Show
  
  -- Generic helpers.
  fmap1 :: Functor f => (t -> a -> b) -> (t -> f a) -> t -> f b
  fmap1 g h x = fmap (g x) (h x)
  
  aZipWith :: Ix i1 => (a -> b -> e) -> Array i1 a -> Array i b -> Array i1 e
  aZipWith f a b = listArray (bounds a) $ zipWith f (elems a) (elems b)
  
  average :: (Foldable f, Fractional a) => f a -> a
  average xs = (sum xs) / (fromIntegral $ length xs)
  
  -- Intermediary point types for ultimately converting into a point `a`.
  newtype FacePoint a = FacePoint { getFacePoint :: a } deriving Show
  newtype EdgeCenterPoint a = EdgeCenterPoint { getEdgeCenterPoint :: a } deriving Show
  newtype EdgePoint a = EdgePoint { getEdgePoint :: a } deriving Show
  newtype VertexPoint a = VertexPoint { getVertexPoint :: a } deriving Show
  
  type FacePointArray a = Array FaceId (FacePoint a)
  type EdgePointArray a = Array EdgeId (EdgePoint a)
  type EdgeCenterPointArray a = Array EdgeId (EdgeCenterPoint a)
  type IsEdgeHoleArray = Array EdgeId Bool
  type VertexPointArray a = Array VertexId (VertexPoint a)
  
  -- Subdivision helpers.
  facePoint :: Fractional a => Mesh a -> Face -> FacePoint a
  facePoint mesh = FacePoint . average . (fmap $ vertexPointById mesh) . faceVertices
  
  allFacePoints :: Fractional a => Mesh a -> FacePointArray a
  allFacePoints = fmap1 facePoint meshFaces
  
  vertexPointById :: Mesh a -> VertexId -> a
  vertexPointById mesh = vertexPoint . (meshVertices mesh !)
  
  edgeCenterPoint :: Fractional a => Mesh a -> Edge -> EdgeCenterPoint a
  edgeCenterPoint mesh (Edge ea eb _)
    = EdgeCenterPoint . average $ fmap (vertexPointById mesh) [ea, eb]
  
  allEdgeCenterPoints :: Fractional a => Mesh a -> EdgeCenterPointArray a
  allEdgeCenterPoints = fmap1 edgeCenterPoint meshEdges
  
  allIsEdgeHoles :: Mesh a -> IsEdgeHoleArray
  allIsEdgeHoles = fmap ((< 2) . length . edgeFaces) . meshEdges
  
  edgePoint :: Fractional a => Edge -> FacePointArray a -> EdgeCenterPoint a -> EdgePoint a
  edgePoint (Edge _ _ [_]) _ (EdgeCenterPoint ecp) = EdgePoint ecp
  edgePoint (Edge _ _ faceIds) facePoints (EdgeCenterPoint ecp)
    = EdgePoint $ average [ecp, average $ fmap (getFacePoint . (facePoints !)) faceIds]
  
  allEdgePoints :: Fractional a => Mesh a -> FacePointArray a -> EdgeCenterPointArray a -> EdgePointArray a
  allEdgePoints mesh fps ecps = aZipWith (\e ecp -> edgePoint e fps ecp) (meshEdges mesh) ecps
  
  vertexPoint' :: Fractional a => Vertex a -> FacePointArray a -> EdgeCenterPointArray a -> IsEdgeHoleArray -> VertexPoint a
  vertexPoint' vertex facePoints ecps iehs
    | length faceIds == length edgeIds = VertexPoint newCoords
    | otherwise = VertexPoint avgHoleEcps
    where
      newCoords = (oldCoords * m1) + (avgFacePoints * m2) + (avgMidEdges * m3)
      oldCoords = vertexPoint vertex
      avgFacePoints = average $ fmap (getFacePoint . (facePoints !)) faceIds
      avgMidEdges = average $ fmap (getEdgeCenterPoint . (ecps !)) edgeIds
      m1 = (n - 3) / n
      m2 = 1 / n
      m3 = 2 / n
      n = genericLength faceIds
      faceIds = vertexFaces vertex
      edgeIds = vertexEdges vertex
      avgHoleEcps = average . (oldCoords:) . fmap (getEdgeCenterPoint . (ecps !)) $ filter (iehs !) edgeIds
  
  allVertexPoints :: Fractional a => Mesh a -> FacePointArray a -> EdgeCenterPointArray a -> IsEdgeHoleArray -> VertexPointArray a
  allVertexPoints mesh fps ecps iehs = fmap (\v -> vertexPoint' v fps ecps iehs) (meshVertices mesh)
  
  -- For each vertex in a face, generate a set of new faces from it with its vertex point,
  -- neighbor edge points, and face point. The new faces will refer to vertices in the
  -- combined vertex array.
  newFaces :: Face -> FaceId -> Int -> Int -> [SimpleFace]
  newFaces (Face vertexIds edgeIds) faceId epOffset vpOffset
    = take (genericLength vertexIds)
    $ zipWith3 newFace (cycle vertexIds) (cycle edgeIds) (drop 1 (cycle edgeIds))
    where
      f = VertexId . (+ epOffset) . getEdgeId
      newFace vid epA epB = SimpleFace
        [ VertexId . (+ vpOffset) $ getVertexId vid
        , f epA
        , VertexId $ getFaceId faceId
        , f epB]
  
  subdivide :: Fractional a => SimpleMesh a -> SimpleMesh a
  subdivide simpleMesh
    = SimpleMesh combinedVertices (listArray (FaceId 0, FaceId (genericLength faces - 1)) faces)
    where
      mesh = makeComplexMesh simpleMesh
      fps = allFacePoints mesh
      ecps = allEdgeCenterPoints mesh
      eps = allEdgePoints mesh fps ecps
      iehs = allIsEdgeHoles mesh
      vps = allVertexPoints mesh fps ecps iehs
      edgePointOffset = length fps
      vertexPointOffset = edgePointOffset + length eps
      combinedVertices
        = listArray (VertexId 0, VertexId (vertexPointOffset + length vps - 1))
        . fmap SimpleVertex
        $ concat [ fmap getFacePoint $ elems fps
                 , fmap getEdgePoint $ elems eps
                 , fmap getVertexPoint $ elems vps]
      faces
        = concat $ zipWith (\face fid -> newFaces face fid edgePointOffset vertexPointOffset)
        (elems $ meshFaces mesh) (fmap FaceId [0..])
  
  -- Transform to a Mesh by filling in the missing references and generating edges.
  -- Faces can be updated with their edges, but must be ordered.
  -- Edge and face order does not matter for vertices.
  -- TODO: Discard degenerate faces (ones with 0 to 2 vertices/edges),
  -- or we could transform these into single edges or vertices.
  makeComplexMesh :: forall a. SimpleMesh a -> Mesh a
  makeComplexMesh (SimpleMesh sVertices sFaces) = Mesh vertices edges faces
    where
      makeEdgesFromFace :: SimpleFace -> FaceId -> [Edge]
      makeEdgesFromFace (SimpleFace vertexIds) fid
        = take (genericLength vertexIds)
        $ zipWith (\a b -> Edge a b [fid]) verts (drop 1 verts)
        where
          verts = cycle vertexIds
  
      edgeKey :: VertexId -> VertexId -> (VertexId, VertexId)
      edgeKey a b = (min a b, max a b)
  
      sFacesList :: [SimpleFace]
      sFacesList = elems sFaces
  
      fids :: [FaceId]
      fids = fmap FaceId [0..]
  
      eids :: [EdgeId]
      eids = fmap EdgeId [0..]
  
      faceEdges :: [[Edge]]
      faceEdges = zipWith makeEdgesFromFace sFacesList fids
  
      edgeMap :: Map.Map (VertexId, VertexId) Edge
      edgeMap
        = Map.fromListWith (\(Edge a b fidsA) (Edge _ _ fidsB) -> Edge a b (fidsA ++ fidsB))
        . fmap (\edge@(Edge a b _) -> (edgeKey a b, edge))
        $ concat faceEdges
  
      edges :: EdgeArray
      edges = listArray (EdgeId 0, EdgeId $ (Map.size edgeMap) - 1) $ Map.elems edgeMap
  
      edgeIdMap :: Map.Map (VertexId, VertexId) EdgeId
      edgeIdMap = Map.fromList $ zipWith (\(Edge a b _) eid -> ((edgeKey a b), eid)) (elems edges) eids
  
      faceEdgeIds :: [[EdgeId]]
      faceEdgeIds = fmap (mapMaybe (\(Edge a b _) -> Map.lookup (edgeKey a b) edgeIdMap)) faceEdges
  
      faces :: FaceArray
      faces
        = listArray (FaceId 0, FaceId $ (length sFaces) - 1)
        $ zipWith (\(SimpleFace verts) edgeIds -> Face verts edgeIds) sFacesList faceEdgeIds
  
      vidsToFids :: Map.Map VertexId [FaceId]
      vidsToFids
        = Map.fromListWith (++)
        . concat
        $ zipWith (\(SimpleFace vertexIds) fid -> fmap (\vid -> (vid, [fid])) vertexIds) sFacesList fids
  
      vidsToEids :: Map.Map VertexId [EdgeId]
      vidsToEids
        = Map.fromListWith (++)
        . concat
        $ zipWith (\(Edge a b _) eid -> [(a, [eid]), (b, [eid])]) (elems edges) eids
  
      simpleToComplexVert :: SimpleVertex a -> VertexId -> Vertex a
      simpleToComplexVert (SimpleVertex point) vid
        = Vertex point
        (Map.findWithDefault [] vid vidsToEids)
        (Map.findWithDefault [] vid vidsToFids)
  
      vertices :: VertexArray a
      vertices
        = listArray (bounds sVertices)
        $ zipWith simpleToComplexVert (elems sVertices) (fmap VertexId [0..])
  
  pShowSimpleMesh :: Show a => SimpleMesh a -> String
  pShowSimpleMesh (SimpleMesh vertices faces)
    = "Vertices:\n" ++ (arrShow vertices sVertexPoint)
    ++ "Faces:\n" ++ (arrShow faces (fmap getVertexId . sFaceVertices))
    where
      arrShow a f = concatMap ((++ "\n") . show . (\(i, e) -> (i, f e))) . zip [0 :: Int ..] $ elems a
  
  -- Testing types.
  data Point a = Point a a a deriving (Show)
  
  instance Functor Point where
    fmap f (Point x y z) = Point (f x) (f y) (f z)
  
  zipPoint :: (a -> b -> c) -> Point a -> Point b -> Point c
  zipPoint f (Point x y z) (Point x' y' z') = Point (f x x') (f y y') (f z z')
  
  instance Num a => Num (Point a) where
    (+) = zipPoint (+)
    (-) = zipPoint (-)
    (*) = zipPoint (*)
    negate = fmap negate
    abs = fmap abs
    signum = fmap signum
    fromInteger i = let i' = fromInteger i in Point i' i' i'
  
  instance Fractional a => Fractional (Point a) where
    recip = fmap recip
    fromRational r = let r' = fromRational r in Point r' r' r'
  
  testCube :: SimpleMesh (Point Double)
  testCube = SimpleMesh vertices faces
    where
      vertices = listArray (VertexId 0, VertexId 7)
        $ fmap SimpleVertex
        [ Point (-1) (-1) (-1)
        , Point (-1) (-1) 1
        , Point (-1) 1 (-1)
        , Point (-1) 1 1
        , Point 1 (-1) (-1)
        , Point 1 (-1) 1
        , Point 1 1 (-1)
        , Point 1 1 1]
      faces = listArray (FaceId 0, FaceId 5)
        $ fmap (SimpleFace . (fmap VertexId))
        [ [0, 4, 5, 1]
        , [4, 6, 7, 5]
        , [6, 2, 3, 7]
        , [2, 0, 1, 3]
        , [1, 5, 7, 3]
        , [0, 2, 6, 4]]
  
  testCubeWithHole :: SimpleMesh (Point Double)
  testCubeWithHole
    = SimpleMesh (sMeshVertices testCube) (ixmap (FaceId 0, FaceId 4) id (sMeshFaces testCube))
  
  testTriangle :: SimpleMesh (Point Double)
  testTriangle = SimpleMesh vertices faces
    where
      vertices = listArray (VertexId 0, VertexId 2)
        $ fmap SimpleVertex
        [ Point 0 0 0
        , Point 0 0 1
        , Point 0 1 0]
      faces = listArray (FaceId 0, FaceId 0)
        $ fmap (SimpleFace . (fmap VertexId))
        [ [0, 1, 2]]
  
  main :: IO ()
  main = putStr . pShowSimpleMesh $ subdivide testCube
  