sloane-0.1: sloane.hs

{-# LANGUAGE DeriveDataTypeable #-}

import Prelude hiding (all)
import System.Console.CmdArgs
import Data.Maybe (fromJust, isJust)
import Data.Function (on)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Control.Monad (when)
import System.IO.Unsafe (unsafePerformIO)
import Network.HTTP (simpleHTTP, getRequest, getResponseBody)
import Network.URL (importURL, exportURL, add_param)
import Data.List (groupBy, sortBy, intersperse)
import Data.Ord (comparing)

newtype OEISEntry = OEISEntry [(Char, String)]

oeisURL = fromJust $ importURL "http://oeis.org/search?fmt=text"
oeisKeys = "ISTUVWXNDHFYAOEeptoKC"

data Sloane = Sloane { keys  :: String
                     , all   :: Bool
                     , limit :: Int
                     , terms :: String
                     }
              deriving (Data, Typeable)

sloane = cmdArgsMode $ Sloane 
         { keys = "SN" &= typ "KEYS" 
           &= help "Keys of fields to print, http://oeis.org/eishelp1.html (default: SN)"
         , all = False 
           &= help "Print all fields"
         , limit = 5 &= name "n"
           &= help "Limit the number of entries retrieved (default: 5)"
         , terms = def &= argPos 0 &= typ "SEARCH-TERMS"
         } 
         &= versionArg [summary "sloane 0.1"]
         &= summary "Search Sloane's On-Line Encyclopedia of Integer Sequences"

-- The maximum number of results returned when searching the OEIS.
maxResults :: IORef Int
maxResults = unsafePerformIO $ newIORef 5

-- Get the maximum number of results returned when searching OEIS
getMaxResults :: IO Int
getMaxResults = readIORef maxResults

-- Set the maximum number of results returned when searching OEIS
setMaxResults :: Int -> IO ()
setMaxResults = writeIORef maxResults

nonempty = not . null

parseOEISEntry :: String -> OEISEntry
parseOEISEntry = OEISEntry
                 . map merge
                 . groupBy ((==) `on` fst)
                 . sortBy (comparing fst)
                 . map (split . tail) . filter nonempty . lines
    where
      split xs@(x:_) = (x, xs)
      merge kvs = let (k:_, vs) = unzip kvs in (k, init $ unlines vs)

parseOEISEntries :: String -> [OEISEntry]
parseOEISEntries = map parseOEISEntry
                   . filter nonempty
                   . map unlines
                   . groupBy (\_ x -> nonempty x)
                   . reverse . drop 2 . reverse . drop 5 . lines

-- Returns the given fields of an 'OEISEntry'. For instance, @fields "NFA" e@
-- would return the name, formula, and auhtor fields of @e@.
fields :: String -> OEISEntry -> [String]
fields ks (OEISEntry e) = [ fromJust d | k<-ks, let d = lookup k e, isJust d ]

-- Returns a list of matching entries in OEIS
lookupOEIS :: String -> IO [OEISEntry]
lookupOEIS s = do
  n <- getMaxResults
  parseOEISEntries `fmap` (simpleHTTP (getRequest (url n)) >>= getResponseBody)
    where
      url n = exportURL $ oeisURL `add_param` ("n", show n) `add_param` ("q", s)

main = do
  args <- cmdArgsRun sloane
  setMaxResults $ limit args
  es <- lookupOEIS . filter (`notElem` "[]") $ terms args
  when (nonempty es) $ putStrLn ""
  let ks = if all args then oeisKeys else keys args
  mapM_ putStrLn [ unlines $ fields ks e | e <- es ]