{-# LANGUAGE TemplateHaskell, NoMonomorphismRestriction, OverloadedStrings #-}
module Ipopt.PP where

import Data.List
import Data.Ord
import Ipopt.Raw
import Ipopt.NLP
import Ipopt.AnyRF
import Text.PrettyPrint.ANSI.Leijen hiding ((<>))

import Text.Printf
import qualified Data.IntMap as IM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Generic as VG
import Data.Vector (Vector)
import qualified Data.Vector as V
import Control.Monad.State
import Control.Monad.Writer
import Foreign.C.Types(CDouble(..))
import Control.Lens
import Data.List
import Data.Foldable (toList,for_, foldMap)

-- * lenses for IpOptSolved
makeLensesWith (lensRules & lensField  .~ stripPrefix "_ipOptSolved_")  ''IpOptSolved

-- * pretty printing
ppSoln state0 problem = flip evalStateT state0 $ runWriterT $ do
    s <- lift problem
    st <- get

    tell (dullgreen "status: " <> (s^.status.to colorStatus))
    br
    tell $ "obj_tot" <> double (s^.objective)
    join $ uses (nlpfun . funF) $ \(AnyRF f) -> case sortBy (comparing fst) $
                                toList (f (s^.x&VG.convert)) `zip` [1 .. ] of
        [_] -> return ()
        [] -> return ()
        xs -> for_ xs $ \(x,i) -> tell $ "obj" <> int i <> colon <$> double x

    for_ (st ^. variablesInv . from ixMap . to IM.toList) $ \(k,desc) -> do
        br
        tell $ string desc <> "(" <> int k <> ")" <> "=" <>
            string (printf "%.3g" (s ^?! x . ix k))
    br
    tell $ "g" <$> foldMap (\e -> mempty <$$> double e) (s ^. g & VG.convert :: V.Vector Double)

    return s

-- * internal

statusOk :: ApplicationReturnStatus -> Bool
statusOk x = case x of
  SolveSucceeded -> True
  SolvedToAcceptableLevel -> True
  UserRequestedStop -> True
  FeasiblePointFound -> True
  _ -> False

colorStatus x = (if statusOk x then id else black . onred) (string (show x))

br = tell (mempty <$$> mempty)