#!/usr/bin/env nix-script
#!>haskell
#! haskell | attoparsec hmatrix

{- Balance chemical reactions using linear algebra. -}

{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative        (many, (<|>))
import Data.Function              (on)
import Data.Maybe                 (fromMaybe)
import Data.List                  (groupBy, sort, nub)
import Data.Text                  (pack)
import Data.Attoparsec.Text
import Numeric.LinearAlgebra      (Field, ident, inv, rank)
import Numeric.LinearAlgebra.Data

type Atom     = String
type Element  = (Atom, Z)
type Compound = [Element]
type Reaction = ([Compound], [Compound])


-- Parsers --

element :: Parser Element
element = do
  atom <- ((++) <$> upper <*> lower) <|> upper
  quantity <- option 1 decimal
  return (atom, quantity)
  where upper  = charRange "A-Z"
        lower  = charRange "a-z"

compound :: Parser Compound
compound = reduce <$> many element

hand :: Parser [Compound]
hand = compound `sepBy` spaced "+"

reaction :: Parser Reaction
reaction = (,) <$> (hand <* spaced "->") <*> hand


-- parser helpers --

-- | Parse while ignoring enclosing whitespace
spaced :: Parser a -> Parser () 
spaced p = skipSpace >> p >> skipSpace

-- | Create a parser from a characters range
charRange :: String -> Parser String
charRange = fmap pure . satisfy . inClass


-- stochiometrycal matrix --

-- | Reduce a 'Compound' to his empirical formula
reduce :: Compound -> Compound
reduce = map count . groupBy ((==) `on` fst) . sort
  where count ((s,n):xs) = (,) s (n + sum (map snd xs))

-- | Find a vector basis for the reaction
basisFor :: Reaction -> [Atom]
basisFor = nub . map fst . concat . fst

-- | Calculate the coordinats vector
coords :: [Atom] -> Compound -> Vector Z 
coords basis v = fromList (map (flip count v) basis)
  where count e = fromMaybe 0 . lookup e 

-- | Calculate stochoimetrical matrix
stoch :: Reaction -> Matrix Z 
stoch r@(x, y) = fromColumns (a ++ b)
  where col = coords (basisFor r)
        a   = map col x
        b   = map (negate . col) y


-- linear algebra --

-- | Find a base of Ker A
kernel :: Matrix Z -> Vector Z
kernel a = cmap (round . min) t 
  where a' = extend (fromZ a) :: Matrix R
        t  = toColumns (inv a') !! (cols a'-1)
        min = (/ minimum (toList t))

-- | Extend a Matrix to a square
extend :: Field k => Matrix k -> Matrix k
extend a
  | m == n    = a
  | otherwise = a === fromRows padd
  where (m,n) = size a
        id    = ident n
        padd  = map (toColumns id !!) [rank a-1..m-1]


main :: IO ()
main = do
  input <- pack <$> prompt
  case parseOnly reaction input of
    Right react -> print (kernel $ stoch react)
    Left err    -> putStrLn ("Failed to parse: " ++ err)
  main
  where prompt = putStr "eq> " >> getLine