<commit_before>import itertools
from dataclasses import dataclass

import scipy
from pymatgen.core import Composition
from scipy.special import comb
from typing import List, Tuple, Union

import numpy as np
from pymatgen.entries.computed_entries import ComputedStructureEntry
from tqdm.autonotebook import tqdm

from piro.data import GAS_RELEASE, ST, H, GASES
from piro.utils import get_v, get_ST


@dataclass
class BalancedReaction:
    reduced_formulas: Tuple[str]
    target_entry: ComputedStructureEntry
    coeffs: np.array
    removed_coeff_indexes: List[int]
    main_formulas: np.array


class SkipReaction(Exception):
    pass


def get_balanced_reaction(
        target_entry: ComputedStructureEntry,
        reduced_formulas: Tuple[str],
        precursors: List[ComputedStructureEntry],
        v_elements_key: str,
        allow_gas_release: bool = False
) -> BalancedReaction:
    if len(set(reduced_formulas)) != len(reduced_formulas):
        raise SkipReaction('Reaction has duplicate formula')

    target_c = target_entry.data[v_elements_key]
    c = np.vstack([p.data[v_elements_key] for p in precursors])

    if np.any(np.sum(c, axis=0) == 0.0):
        raise SkipReaction('Reaction as sum 0 c')

    try:
        coeffs = np.linalg.solve(c.T, target_c)
    except:
        # need better handling here.
        raise SkipReaction('Reaction could not solve for coeffs')

    if np.any(np.abs(coeffs) > 100):
        raise SkipReaction('Reaction has a coeff over 100')

    if np.any(coeffs < 0.0):
        if not allow_gas_release:
            raise SkipReaction('Reaction has a gas release')
        else:
            if not set(np.array(reduced_formulas)[coeffs < 0.0]).issubset(GAS_RELEASE):
                raise SkipReaction('Reaction gas releases are not expected gases')

    removed_coeff_indexes = []
    for i in reversed(range(len(coeffs))):
        if np.abs(coeffs[i]) < 0.00001:
            c = np.delete(c, i, 0)
            coeffs = np.delete(coeffs, i)
            removed_coeff_indexes.append(i)

    effective_rank = scipy.linalg.lstsq(c.T, target_c)[2]
    if effective_rank < len(coeffs):
        # Removes under-determined reactions.
        # print(effective_rank, precursor_formulas, \
        # [prec_.composition.reduced_formula for prec_ in precursors],coeffs)
        raise SkipReaction('Reaction is under-determined')

    main_formulas = np.array([
        p for i, p in enumerate(reduced_formulas) if i not in removed_coeff_indexes
    ])

    return BalancedReaction(
        reduced_formulas,
        target_entry,
        coeffs,
        removed_coeff_indexes,
        main_formulas,
    )


class Reaction:

    def __init__(<commit_msg>more cleanup<commit_after>import itertools
from dataclasses import dataclass

import scipy
from pymatgen.core import Composition
from scipy.special import comb
from typing import List, Tuple, Union

import numpy as np
from pymatgen.entries.computed_entries import ComputedStructureEntry
from tqdm.autonotebook import tqdm

from piro.data import GAS_RELEASE, ST, H, GASES
from piro.utils import get_v, get_ST


@dataclass
class BalancedReaction:
    reduced_formulas: Tuple[str]
    target_entry: ComputedStructureEntry
    coeffs: np.array
    removed_coeff_indexes: List[int]
    main_formulas: np.array


class SkipReaction(Exception):
    pass


def get_balanced_reaction(
        target_entry: ComputedStructureEntry,
        reduced_formulas: Tuple[str],
        precursors: List[ComputedStructureEntry],
        v_elements_key: str,
        allow_gas_release: bool = False
) -> BalancedReaction:
    if len(set(reduced_formulas)) != len(reduced_formulas):
        raise SkipReaction('Reaction has duplicate formula')

    target_c = target_entry.data[v_elements_key]
    c = np.vstack([p.data[v_elements_key] for p in precursors])

    if np.any(np.sum(c, axis=0) == 0.0):
        raise SkipReaction('Reaction as sum 0 c')

    try:
        coeffs = np.linalg.solve(c.T, target_c)
    except:
        # need better handling here.
        raise SkipReaction('Reaction could not solve for coeffs')

    if np.any(np.abs(coeffs) > 100):
        raise SkipReaction('Reaction has a coeff over 100')

    if np.any(coeffs < 0.0):
        if not allow_gas_release:
            raise SkipReaction('Reaction has a gas release')
        else:
            if not set(np.array(reduced_formulas)[coeffs < 0.0]).issubset(GAS_RELEASE):
                raise SkipReaction('Reaction gas releases are not expected gases')

    removed_coeff_indexes = []
    for i in reversed(range(len(coeffs))):
        if np.abs(coeffs[i]) < 0.00001:
            c = np.delete(c, i, 0)
            coeffs = np.delete(coeffs, i)
            removed_coeff_indexes.append(i)

    effective_rank = scipy.linalg.lstsq(c.T, target_c)[2]
    if effective_rank < len(coeffs):
        # Removes under-determined reactions.
        raise SkipReaction('Reaction is under-determined')

    main_formulas = np.array([
        p for i, p in enumerate(reduced_formulas) if i not in removed_coeff_indexes
    ])

    return BalancedReaction(
        reduced_formulas,
        target_entry,
        coeffs,
        removed_coeff_indexes,
        main_formulas,
    )


class Reaction:

    def __init__(
{'fraction_non_alphanumeric': 0.0860460689130021, 'fraction_numerical': 0.008376166000380734, 'mean_word_length': 3.057142857142857, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2924087', 'n_tokens_mistral': 1782, 'n_tokens_neox': 1697, 'n_words': 419}