{-# LANGUAGE OverloadedStrings #-}

import qualified Data.Map as Map
import qualified Data.Text as StrictText
import qualified Data.Text.Lazy as Text
import qualified Data.Text.Lazy.IO as Text
import qualified Data.List as List
import qualified System.IO as IO
import qualified Data.ByteString.Lazy as BS

import System.Environment 
import Control.Monad 
import Data.Maybe
import System.Console.GetOpt
import Control.Applicative
import Data.Binary
import qualified Vowpal as Vowpal
import Debug.Trace

data Opts = Opts { nolearning   :: Bool 
                 , bitsize      :: Int
                 , maxlabels    :: Int 
                 , path         :: Maybe FilePath 
                 , help         :: Bool }
            deriving (Show)
                     
defaults :: Opts 
defaults =  Opts { nolearning = False
                 , bitsize = 26
                 , maxlabels = maxBound
                 , path = Nothing 
                 , help = False }

options :: [ OptDescr (Opts -> Opts) ]
options = 
  [ Option "" ["no-learning"] (NoArg (\ opt -> opt { nolearning = True } ))
    "only prediction, no learning"
  , Option "" ["size"]        (ReqArg (\arg opt -> opt { bitsize = read arg } )  "INT")
    "model size in bits"
  , Option "" ["max-labels"]  (ReqArg (\arg opt -> opt { maxlabels = read arg }) "INT")
    "maximum number of the label set"
  , Option "h" ["help"]       (NoArg (\opt -> opt { help = True } ))
    "print help"
  ]

parseOptions :: [String] -> Either [String] Opts
parseOptions args = 
  case getOpt Permute options args of
    (os, as, [])  -> Right $ List.foldl' (flip ($)) defaults { path = listToMaybe as }  os
    (_, _,  errs)  -> Left errs
                                     
main :: IO ()
main = do
  args <- getArgs  
  case parseOptions args of
    Left errs                  -> error $ concat errs
    Right opts | help opts     -> do
      when (help opts) $ do 
        putStrLn (usageInfo "Usage: progressive [OPTION...] model-path" options)
    Right opts | isNothing (path opts) -> error "Missing required argument model-path"
    Right opts -> progressive opts  

progressive :: Opts -> IO ()
progressive opts =
  case nolearning opts of
    False -> do
       let size = show . bitsize $ opts
           -- | size of model. Use 29 bits if you have enough memory
           out = maybe (error "Missing path") id . path $ opts
           -- | Path where model will be saved
           vw_opts = ["--adaptive", "--bit_precision", size]
       xys <- map parseLine . Text.lines <$> Text.getContents
       m0 <- Vowpal.emptyIO (maxlabels opts) vw_opts out
       IO.hSetBuffering IO.stdout IO.LineBuffering
       let step m xy = do
             (p:ps,m') <- Vowpal.predictUpdateIO m xy
             Text.putStrLn . Text.fromStrict . paste 
               . map fst
               $ p:ps
             return m'
       m <- foldM step m0 xys
       Vowpal.save m 
    True -> do
       m <- Vowpal.load . maybe (error "Missing path") id . path $ opts
       IO.hSetBuffering IO.stdout IO.LineBuffering
       xys <- map  parseLine . Text.lines <$> Text.getContents
       let step xy = do       
             p:ps <- fst <$> Vowpal.predictIO m (fst xy)
             Text.putStrLn . Text.fromStrict . paste 
               . map fst
               $ p:ps
       flip mapM_ xys step

paste :: [StrictText.Text] -> StrictText.Text
paste = StrictText.concat . List.intersperse ","

parseLine :: Text.Text -> ([StrictText.Text], [StrictText.Text])
parseLine ln = 
    let (x:xs) = StrictText.words . Text.toStrict $ ln
        ys = StrictText.splitOn "," x
    in (xs,ys)