CodeObject/storage/zeta/rs/completed/2006.rs

use std::{
    collections::HashMap,
    io::{read_to_string, stdin},
};

#[derive(Debug)]
enum Token {
    Atom { symbol: String },
    AddOp,
    EqOp,
    Num { num: u64 },
    Undef,
}

struct Atom {
    symbol: String,
}

struct Molecular {
    composites: Vec<(Atom, u64)>,
}

struct Expr {
    composites: Vec<(Molecular, u64)>,
}

// stmt := Expr = Expr
struct Stmt {
    lhs: Expr,
    rhs: Expr,
}

/*
    Lex
    Tokenizing
*/
fn lex(line: &mut String) -> Vec<Token> {
    let mut tokens = vec![];

    let mut chars = line.chars();

    let mut symb = String::new();

    let mut tk = Token::Undef;

    let mut pull_tk = |tk: &mut Token, tokens: &mut Vec<Token>, symb: &mut String| {
        if symb.len() == 0 {
            return;
        }
        let mut current_tk = std::mem::replace(tk, Token::Undef);

        match current_tk {
            Token::Num { ref mut num } => {
                if let Ok(parsed) = symb.parse::<u64>() {
                    *num = parsed;
                }
                symb.clear();
            }
            Token::Atom { ref mut symbol } => {
                *symbol = std::mem::take(symb);
            }
            _ => {}
        }
        tokens.push(current_tk);
    };

    let mut flg = false;

    while let Some(curr) = chars.next() {
        if curr.is_whitespace() {
            pull_tk(&mut tk, &mut tokens, &mut symb);
            continue;
        }
        if curr.is_uppercase() {
            if flg {
                pull_tk(&mut tk, &mut tokens, &mut symb);
            }

            flg = true;
            symb.push(curr);
            tk = Token::Atom {
                symbol: String::new(),
            };
        } else if curr.is_lowercase() {
            symb.push(curr);
        } else if curr.is_numeric() {
            if flg {
                if !matches!(tk, Token::Num { .. }) {
                    pull_tk(&mut tk, &mut tokens, &mut symb);
                }
            }

            flg = true;
            symb.push(curr);
            tk = Token::Num { num: 0 };
        } else if curr == '+' {
            if flg {
                pull_tk(&mut tk, &mut tokens, &mut symb);
            }
            flg = false;
            tk = Token::Undef;
            tokens.push(Token::AddOp);
        } else if curr == '=' {
            if flg {
                pull_tk(&mut tk, &mut tokens, &mut symb);
            }
            flg = false;
            tk = Token::Undef;
            tokens.push(Token::EqOp);
        }
    }
    pull_tk(&mut tk, &mut tokens, &mut symb);

    tokens
}

/*
    Parse
*/

fn parse_molecular(tokens: &Vec<Token>, ptr: &mut usize) -> Molecular {
    let mut comps = vec![];
    while *ptr < tokens.len() {
        let Token::Atom { symbol } = &tokens[*ptr] else {
            break;
        };

        let atom = Atom {
            symbol: symbol.clone(),
        };

        *ptr += 1;
        if *ptr < tokens.len() {
            if let Token::Num { num } = &tokens[*ptr] {
                *ptr += 1;
                comps.push((atom, *num));
            } else {
                comps.push((atom, 1));
            }
        } else {
            comps.push((atom, 1));
        }
    }
    Molecular { composites: comps }
}

fn parse_expr(tokens: &Vec<Token>, ptr: &mut usize) -> Expr {
    let mut comps = vec![];
    while *ptr < tokens.len() {
        if let Token::EqOp = &tokens[*ptr] {
            break;
        }

        if let Token::AddOp = &tokens[*ptr] {
            *ptr += 1;
        }

        let mut cnt = 1;
        if let Token::Num { num } = &tokens[*ptr] {
            *ptr += 1;
            cnt = *num;
        }

        let molecular = parse_molecular(tokens, ptr);

        comps.push((molecular, cnt));
    }

    Expr { composites: comps }
}

fn parse_stmt(tokens: &Vec<Token>, ptr: &mut usize) -> Stmt {
    let lhs = parse_expr(tokens, ptr);
    *ptr += 1;
    let rhs = parse_expr(tokens, ptr);

    Stmt { lhs, rhs }
}

fn count_amount(expr: &Expr) -> HashMap<String, u64> {
    let mut counter = HashMap::new();

    for (mol, cnt_mol) in expr.composites.iter() {
        for (atom, cnt_atom) in mol.composites.iter() {
            counter
                .entry(atom.symbol.clone())
                .and_modify(|cnt| *cnt += cnt_atom * cnt_mol)
                .or_insert(cnt_atom * cnt_mol);
        }
    }

    counter
}

fn main() {
    let mut line = String::new();
    stdin().read_line(&mut line).unwrap();

    let mut case = 0;
    while !line.starts_with('#') {
        case += 1;
        line = line.trim().to_string();
        let tks = lex(&mut line);

        let stmt = parse_stmt(&tks, &mut 0);

        let lhs_counter = count_amount(&stmt.lhs);
        let rhs_counter = count_amount(&stmt.rhs);

        // calculate destoryed and created atoms
        let mut calcs = HashMap::new();
        rhs_counter.iter().for_each(|(symb, cnt)| {
            calcs
                .entry(symb.clone())
                .and_modify(|c| *c += *cnt as i64)
                .or_insert(*cnt as i64);
        });

        lhs_counter.iter().for_each(|(symb, cnt)| {
            calcs
                .entry(symb.clone())
                .and_modify(|c| *c -= *cnt as i64)
                .or_insert(-(*cnt as i64));
        });

        let flg = calcs.values().any(|cnt| *cnt != 0);

        if !flg {
            println!("Equation {} is balanced.", case);
        } else {
            println!("Equation {} is unbalanced.", case);

            let mut calcs = calcs
                .into_iter()
                .map(|(symb, cnt)| (symb, cnt))
                .collect::<Vec<_>>();

            calcs.sort();

            for (symb, cnt) in calcs.iter() {
                if *cnt == 1 {
                    println!("You have created {} atom of {}.", cnt, symb);
                } else if *cnt > 1 {
                    println!("You have created {} atoms of {}.", cnt, symb);
                } else if *cnt == -1 {
                    println!("You have destroyed {} atom of {}.", -cnt, symb);
                } else if *cnt < -1 {
                    println!("You have destroyed {} atoms of {}.", -cnt, symb);
                }
            }
            println!();
        }

        line.clear();
        stdin().read_line(&mut line).unwrap();
    }
}