chemical-equation-0.0: src/Common.hs

module Common where

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import Data.Map (Map)
import Data.Set (Set)
import Data.Maybe (mapMaybe, isNothing)
import Data.Tuple.HT (mapPair, mapSnd, swap)

import qualified Text.Parsec as Parsec
import Text.Parsec.String (Parser)

import Text.Printf (printf)

import qualified Shell.Utility.Log as Log
import Shell.Utility.Verbosity (Verbosity)

import Control.Monad (when)
import Control.Applicative (liftA2, (<$>))



type Reactant = Map String Integer


parseReactant :: Parser Reactant
parseReactant =
   fmap (Map.fromListWith (+)) $
   Parsec.many1 $ do
      element <- liftA2 (:) Parsec.upper (Parsec.many Parsec.lower)
      kStr <- Parsec.many Parsec.digit
      case (kStr, read kStr) of
         ("", _) -> return (element, 1)
         {- ToDo:
         should be an unrecoverable parser error
         in order to not generate confusing additional errors,
         e.g. when parsing C0
         -}
         (_, 0) ->
            Parsec.unexpected $
            printf "element %s has multiplicity zero" element
         (_, k) -> return (element, k)

preprocessArguments ::
   Verbosity -> [String] -> IO (Set Reactant, [(Reactant, String)])
preprocessArguments verbosity args = do
   let (brokenArgs,reactants) =
         ListHT.unzipEithers $
         map (\arg -> flip (,) arg <$> Parsec.parse parseReactant arg arg)
            args
   case brokenArgs of
      _:_ -> fail $ concatMap ("\n\n"++) $ map show brokenArgs
      [] -> do
         let m = Map.fromListWith NonEmptyC.append $
                     map (mapSnd NonEmpty.singleton) reactants
         let duplicatesGroups =
               Map.filter
                  (\xs -> case xs of NonEmpty.Cons _ (_:_) -> True; _ -> False)
                  m
         Fold.for_ duplicatesGroups $ \duplicates ->
            Log.warn verbosity $
               printf "duplicate reactants: %s\n"
                  (List.intercalate ", " $ NonEmpty.flatten duplicates)

         return
            (Map.keysSet m,
               map snd $ filter fst $ snd $
               Trav.mapAccumL
                  (\seenSoFar p@(reactant,_name) ->
                     (Set.insert reactant seenSoFar,
                        (Set.notMember reactant seenSoFar, p)))
                  Set.empty reactants)



mapFromListWithoutDuplicates :: (Ord k) => [(k,a)] -> Either [k] (Map k a)
mapFromListWithoutDuplicates xs =
   let m = Map.fromListWith (\_ _ -> Nothing) $ map (mapSnd Just) xs in
   case Trav.sequence m of
      Just table -> Right table
      Nothing -> Left $ Map.keys $ Map.filter isNothing m


cancelIntegerVector :: (Functor f, Foldable f) => f Integer -> f Integer
cancelIntegerVector v =
   let d = Fold.foldl gcd 0 v
   in fmap (flip div d) v


integerLinearCombination :: Map Reactant Integer -> Map String Integer
integerLinearCombination =
   Map.unionsWith (+) . Map.elems .
   Map.mapWithKey (\r k -> fmap (k*) r)

formatEquation :: [(Reactant,String)] -> Map Reactant Integer -> String
formatEquation rs reactantMap =
   let tagged =
         mapMaybe
            (\(r,name) ->
               let n = reactantMap Map.! r in
               case compare n 0 of
                  EQ -> Nothing
                  GT -> Just $ (True,(n,name))
                  LT -> Just $ (False,(-n,name)))
            rs
       (lhs,rhs) =
         (case tagged of ((True,_):_) -> id; _ -> swap) $
         mapPair (map snd, map snd) $
         List.partition fst tagged
       formatSum =
         List.intercalate " + " .
         map (\(n,str) -> if n==1 then str else printf "%d %s" n str)
   in  formatSum lhs ++ " <=> " ++ formatSum rhs

displaySolutions ::
   Verbosity -> [(Reactant, String)] -> [Map Reactant Integer] -> IO ()
displaySolutions verbosity reactants sols = do
   Fold.for_ sols $ \sol -> do
      putStrLn . formatEquation reactants $ sol
      let probe = integerLinearCombination sol
      when (Fold.any (0/=) probe) $
         Log.warn verbosity $ "probe failed\n" ++ show probe ++ "\n"

   let unused =
         if null sols
         then Set.fromList $ map fst reactants
         else Map.keysSet $
               Map.filter id $ Map.unionsWith (&&) $ map (fmap (0==)) sols
   when (not $ Set.null unused) $
      Log.warn verbosity $
      printf "Unused reactants: %s\n" $ List.intercalate ", " $ map snd $
      filter (flip Set.member unused . fst) $ reactants