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 qualified Data.Char.Small as Unicode
import qualified Data.Char as Char
import Data.Map (Map)
import Data.Set (Set)
import Data.Maybe (mapMaybe, isNothing)
import Data.Tuple.HT (mapPair, mapSnd, swap)
import Data.Bool.HT (if')
import qualified Text.Parsec.Pos as SourcePos
import qualified Text.Parsec as Parsec
import Text.Parsec.String (Parser)
import Text.Parsec ((<|>))
import Text.Printf (printf)
import qualified Options.Applicative as OP
import qualified Shell.Utility.Verbosity as Verbosity
import qualified Shell.Utility.Log as Log
import Shell.Utility.Verbosity (Verbosity)
import Control.Monad (when, void)
import Control.Applicative (liftA2, liftA3, (<$>), (<*))
type Reactant = Map String Integer
consuming :: Parser a -> Parser (String, a)
consuming p = do
input <- Parsec.getInput
posBefore <- Parsec.getPosition
result <- p
posAfter <- Parsec.getPosition
return
(take (SourcePos.sourceColumn posAfter
- SourcePos.sourceColumn posBefore) input,
result)
parseElement :: Parser Reactant
parseElement =
flip Map.singleton 1 <$> liftA2 (:) Parsec.upper (Parsec.many Parsec.lower)
parseMultiplicity :: String -> Parser Integer
parseMultiplicity partStr = do
kStr <- Parsec.many Parsec.digit
case (kStr, read kStr) of
("", _) -> return 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 "%s has multiplicity zero" partStr
(_, k) -> return k
chargeStr :: String
chargeStr = "charge"
parseCharge :: Parser Reactant
parseCharge = do
void $ Parsec.char '^'
k <- parseMultiplicity chargeStr
c <- k <$ Parsec.char '+' <|> (-k) <$ Parsec.char '-'
return $ Map.singleton chargeStr c
parsePart :: Parser Reactant
parsePart =
fmap (Map.unionsWith (+)) $
flip Parsec.sepBy1 (Parsec.optional (Parsec.char '-')) $ do
(partStr,part) <-
consuming $
Parsec.between (Parsec.char '(') (Parsec.char ')') parsePart
<|>
parseCharge
<|>
parseElement
k <- parseMultiplicity partStr
return ((k*) <$> part)
parseReactant :: Parser Reactant
parseReactant = parsePart <* Parsec.eof
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)
type Format = [(Reactant,String)] -> Map Reactant Integer -> String
toUnicode :: String -> String
toUnicode =
let go [] = []
go (c:cs) =
if' (c=='^')
(case mapSnd (splitAt 1) $ span Char.isDigit cs of
(ds, (sign,remnd)) ->
map Unicode.superscript (ds++sign) ++ go remnd) $
if' (Char.isDigit c) (Unicode.subscript c : go cs) $
c : go cs
in go
formatEquation :: Bool -> String -> Format
formatEquation useUnicode arrow rs reactantMap =
let eyeCandy = if useUnicode then toUnicode else id
tagged =
mapMaybe
(\(r,name) ->
let n = reactantMap Map.! r in
case compare n 0 of
EQ -> Nothing
GT -> Just $ (True, ( n, eyeCandy name))
LT -> Just $ (False, (-n, eyeCandy 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 ++ " " ++ arrow ++ " " ++ formatSum rhs
displaySolutions ::
Verbosity -> Format ->
[(Reactant, String)] -> [Map Reactant Integer] -> IO ()
displaySolutions verbosity format reactants sols = do
Fold.for_ sols $ \sol -> do
putStrLn . format 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
info :: OP.Parser a -> OP.ParserInfo a
info p =
OP.info
(OP.helper <*> p)
(OP.fullDesc <>
OP.progDesc "Balance chemical equations")
parser :: OP.Parser (Verbosity, Format, [String])
parser =
liftA3 (,,)
(OP.option (OP.eitherReader Verbosity.parse) $
OP.value Verbosity.normal <>
OP.short 'v' <>
OP.long "verbose" <>
OP.metavar "0..3" <>
OP.help "verbosity")
parseFormatting
(OP.some $
OP.strArgument
(OP.metavar "STRING" <>
OP.help "reactant"))
parseFormatting :: OP.Parser Format
parseFormatting =
pure formatEquation
<*>
OP.switch
(OP.long "unicode" <>
OP.help "format output using Unicode symbols")
<*>
{-
Equilibrium equation: rightwards harpoon over leftwards harpoon
-}
OP.strOption
(OP.long "arrow" <>
OP.metavar "STRING" <>
OP.value "<=>" <>
OP.help "alternative arrow symbol")