diagrams-0.2.1.3: example/wordcloud.hs

module Main where

-- Cobbled together from snippets and ideas by Chris Done
-- http://github.com/chrisdone/wordcloud/

import Control.Monad (liftM)
import Data.Char (isLetter, toLower)
import Data.List (sortBy, foldl')
import Data.Ord (comparing)
import qualified Data.Map as M

import System.Random (getStdGen, randomRs)
import System.Environment (getArgs)

import Graphics.Rendering.Diagrams
import Data.Colour (withOpacity)
import Data.Colour.SRGB.Linear (rgb)

import CloudPacker (diagram)

type Weight = Int
type Word = String
type Histogram = [(Word,Weight)]
type Colour = [Double]

-- Word histograms.
--
histogramByFreq :: [Word] -> String -> Histogram
histogramByFreq badws = list . table where
    table = filterByGood badws . histogram . words . map toLetter
    list = take 150 . sortBy (flip (comparing snd)) . M.toAscList

toLetter c | isLetter c = c
           | otherwise  = ' '

histogram = foldl' (flip $ flip (M.insertWith' $ const (+1)) 1) M.empty

filterByGood badws = M.filterWithKey (\x y -> goodWord x) where
    goodWord [_] = False
    goodWord w   = not $ any (==(map toLower w)) badws -- No articles.

stopwords = words "import qualified hiding data newtype type deriving instance do if then else case of let where"


main = do
    args <- getArgs
    let input = case args of
                    "-":_       -> getContents
                    filename:_  -> readFile filename
                    _           -> readFile "wordcloud.hs"
    weightedwords <- histogramByFreq stopwords `liftM` input
    rands <- groupsOf 3 `liftM` randomRs (0,1) `liftM` getStdGen
    let sizedwords = mkWords weightedwords rands
    renderAs PNG "wordcloud.png" (Width 500) $ diagram sizedwords


-- Sizing and colouring the text according to the
-- weight given in the histogram.
mkWords :: Histogram -> [Colour] -> [Diagram]
mkWords wwds cols = zipWith (mkWord maxweight) wwds cols
    where maxweight = snd $ head wwds

mkColour [r,g,b] a = rgb r g b `withOpacity` max a 0.1

mkWord :: Weight -> (Word,Weight) -> Colour -> Diagram
mkWord mx (s,w) col = fc c $ lw 0 $ tf "URW Bookman L" $ textPath sz s
    where sz = fromIntegral w * 10
          c  = mkColour col (max 0.2 (fromIntegral w/fromIntegral mx))

groupsOf :: Int -> [a] -> [[a]]
groupsOf n [] = []
groupsOf n xs = let (grp, remainder) = splitAt n xs
                in grp : groupsOf n remainder