content: {-# LANGUAGE ConstraintKinds     #-}
  {-# LANGUAGE FlexibleContexts    #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeOperators       #-}
  -- |
  -- Module      : Data.Array.Accelerate.Math.Kmeans
  -- Copyright   : [2014..2020] Trevor L. McDonell
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  -- An implementation of K-means clustering in Accelerate using Lloyd's
  -- algorithm. This heuristic technique finds a solution by iteratively improving
  -- an initial guess. The algorithm takes as input the number of clusters to find
  -- and makes an initial guess at the center of each cluster. Then proceeds as
  -- follows:
  --
  --   1. Assign each point to the cluster to which it is closest. This forms the
  --   new set of clusters.
  --
  --   2. Compute the centroid of the cluster (the average of all points in the
  --   cluster.
  --
  --   3. Repeat until the centroid locations stabilise (or some maximum iteration
  --   limit is reached).
  --
  -- Adapted from K-Means sample from "Parallel and Concurrent Programming in
  -- Haskell", (c) Simon Marlow, 2013.
  --
  
  module Kmeans ( Point, Cluster, Id, kmeans )
    where
  
  import Prelude                                          as P
  import Data.Array.Accelerate                            as A
  
  
  -- This implementation works on 2D points. In future, generalise this to some
  -- arbitrary "feature vector".
  --
  type Point a = (a, a)
  
  -- Clusters consist of the centroid location as well as its identifier
  --
  type Id = Word32
  type Cluster a = (Id, (a, a))
  
  idOfCluster :: Elt a => Exp (Cluster a) -> Exp Id
  idOfCluster = A.fst
  
  centroidOfCluster :: Elt a => Exp (Cluster a) -> Exp (Point a)
  centroidOfCluster = A.snd
  
  -- We'll use this as an intermediate structure; it contains the number of points
  -- in the set as well as the sum of the x and y coordinates respectively.
  --
  type PointSum a = (Word32, (a, a))
  
  
  -- Get the distance (squared) between two points. Since we only compare this
  -- value, we elide the square root.
  --
  distance :: A.Num a => Exp (Point a) -> Exp (Point a) -> Exp a
  distance u v =
    let (x1,y1) = unlift u
        (x2,y2) = unlift v
    in
    (x1-x2) P.^ (2::Int) + (y1-y2) P.^ (2::Int)
  
  
  -- For each of the given points, return the cluster Id that that that point is
  -- closest to.
  --
  findClosestCluster
      :: forall a. (A.RealFloat a, P.RealFloat a)
      => Acc (Vector (Cluster a))
      -> Acc (Vector (Point a))
      -> Acc (Vector Id)
  findClosestCluster clusters points =
    A.map (\p -> A.fst $ A.sfoldl (nearest p) z (constant Z) clusters) points
    where
      z = constant (-1, inf)
  
      nearest :: Exp (Point a) -> Exp (Id, a) -> Exp (Cluster a) -> Exp (Id, a)
      nearest p st c =
        let d  = A.snd st
            d' = distance p (centroidOfCluster c)
        in
        d' A.< d ? ( lift (idOfCluster c, d') , st )
  
  
  -- Given a vector of points and a vector of clusters we, we first locate the
  -- closest cluster to each point, assign that point to their closest cluster,
  -- and compute the centroid of the cluster. This yields the new centroid
  -- locations.
  --
  makeNewClusters
      :: forall a. (A.RealFloat a, P.RealFloat a, A.FromIntegral Word32 a)
      => Acc (Vector (Point a))
      -> Acc (Vector (Cluster a))
      -> Acc (Vector (Cluster a))
  makeNewClusters points clusters
    = pointSumToCluster
    . makePointSum
    . findClosestCluster clusters
    $ points
  
  -- TLM: This setup might be quicker, because it forces the result of
  --      findClosestCluster to be evaluated, which overall reduces memory
  --      traffic. However, this is hitting a sharing recovery bug in Accelerate
  --      so can't be used right now ):
  --
  --      As per tip in issue #148, we can get around the bug by making the first
  --      argument to pipe closed. It turns out that the first version is quicker!
  --
  --  = A.uncurry findClosestCluster >-> pointSumToCluster . makePointSum $ A.lift (clusters, points)
    where
      npts        = size points
      nclusters   = size clusters
  
      -- Turn the PointSum intermediate structure into the clusters, by averaging
      -- the cumulative (x,y) positions.
      --
      pointSumToCluster :: Acc (Vector (PointSum a)) -> Acc (Vector (Cluster a))
      pointSumToCluster ps =
        A.generate (A.shape ps)
                   (\ix -> lift (A.fromIntegral (unindex1 ix), average (ps ! ix)))
  
      average :: Exp (PointSum a) -> Exp (Point a)
      average ps =
        let (n, xy) = unlift ps   :: (Exp Word32, Exp (Point a))
            (x, y)  = unlift xy
        in
        lift (x / A.fromIntegral n, y / A.fromIntegral n) -- TLM: what if there are no points in the cluster??
  
      -- Reduce along the rows of 'pointSum' to get the cumulative (x,y) position
      -- and number of points assigned to each centroid.
      --
      makePointSum :: Acc (Vector Id) -> Acc (Vector (PointSum a))
      makePointSum = A.fold1 addPointSum . compute . pointSum
  
      -- The point sum is an intermediate 2D array (it gets fused away, so does
      -- not exist in memory). The points are laid out along the innermost
      -- dimension (rows), and down the column is each of the clusters.
      --
      -- For each point, we put its (x,y) coordinates into the row corresponding
      -- to whichever cluster it is closest to, and zeros in each of the other
      -- rows.
      --
      pointSum :: Acc (Vector Id) -> Acc (Array DIM2 (PointSum a))
      pointSum nearest =
        A.generate (lift (Z:.nclusters:.npts))
                   (\ix -> let Z:.i:.j = unlift ix    :: Z :. Exp Int :. Exp Int
                               near    = nearest ! index1 j
  
                               yes     = lift (constant 1, points ! index1 j)
                               no      = constant (0, (0, 0))
                           in
                           near A.== A.fromIntegral i ? ( yes, no ))
  
      addPointSum :: Exp (PointSum a) -> Exp (PointSum a) -> Exp (PointSum a)
      addPointSum x y =
        let (c1, u) = unlift x    :: (Exp Word32, Exp (Point a))
            (c2, v) = unlift y    :: (Exp Word32, Exp (Point a))
            (x1,y1) = unlift u    :: (Exp a, Exp a)
            (x2,y2) = unlift v    :: (Exp a, Exp a)
        in
        lift (c1+c2, lift (x1+x2, y1+y2) :: Exp (Point a))
  
  {--
      -- Alternative to computing the PointSum structure.
      --
      -- This method uses a forward permutation with atomic instructions to create
      -- the array directly (this method is closer to what one might write
      -- sequentially). This avoids a parallel reduction, but has very high
      -- contention. Overall performance much lower, as:
      --
      --   number of clusters << number of points
      --
      makePointSum :: Acc (Vector (PointSum a))
      makePointSum = A.permute addPointSum zeros near input
        where
          zeros   = A.fill (constant (Z:.nclusters)) (constant (0,(0,0)))
          input   = A.zip (A.fill (A.shape points) (constant 1)) points
          near ix = index1 (A.fromIntegral (nearest ! ix))
  --}
  
  
  -- To complete the k-means algorithm, we loop repeatedly generating new clusters
  -- positions, until the positions converge (or some maximum iteration limit is
  -- reached?)
  --
  kmeans :: forall a. (A.RealFloat a, P.RealFloat a, A.FromIntegral Word32 a)
         => Acc (Vector (Point a))        -- the points to cluster
         -> Acc (Vector (Cluster a))      -- initial cluster positions (guess)
         -> Acc (Vector (Cluster a))
  kmeans points clusters
    = A.asnd
    $ A.awhile (A.uncurry keepGoing)
               (\cs -> let (_, old) = unlift cs   :: (Acc (Vector (Cluster a)), Acc (Vector (Cluster a)))
                           new      = makeNewClusters points old
                       in
                       lift (old,new))
               (lift (clusters, makeNewClusters points clusters))
    where
      keepGoing :: Acc (Vector (Cluster a)) -> Acc (Vector (Cluster a)) -> Acc (Scalar Bool)
      keepGoing xs ys
        = A.or
        $ A.zipWith (\c1 c2 -> let (x1,y1) = unlift (centroidOfCluster c1)
                                   (x2,y2) = unlift (centroidOfCluster c2)
                               in
                               abs (x1-x2) A.> 0.01 A.|| abs (y1-y2) A.> 0.01) xs ys
  
  
  -- The largest non-infinite floating point number
  --
  inf :: forall a. P.RealFloat a => a
  inf = P.encodeFloat m n
    where
      a           = undefined :: a
      b           = P.floatRadix a
      e           = P.floatDigits a
      (_, e')     = P.floatRange a
      m           = b P.^ e - 1
      n           = e' - e
  
  