# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
import re, difflib, k1lib
from typing import Dict, Union, List, Optional
__all__ = ["Eqn", "Eqns", "System"]
settings = k1lib.Settings()
settings.spaceBetweenValueSymbol = True
settings.eqnPrintExtras = True
k1lib.settings.add("eqn", settings, "from k1lib.eqn module");
[docs]class Eqn:                                                                       # Eqn
[docs]    def __init__(self, system:"System"):                                         # Eqn
        """Creates a blank equation. Not expected to be instantiated by
the end user."""                                                                 # Eqn
        self.system = system                                                     # Eqn
        self.terms:Dict[str, float] = {}                                         # Eqn 
    def removeZeros(self):                                                       # Eqn
        self.terms = {k: v for k, v in self.terms.items() if abs(v) > 1e-6}; return self # Eqn
    def parse(self, line:str):                                                   # Eqn
        line = line.strip().replace(" +", "+").replace("+ ", "+").replace(" ->", "->").replace("-> ", "->") # Eqn
        reactants, products = line.split("->")                                   # Eqn
        for i, side in enumerate(line.split("->")):                              # Eqn
            sign = i * 2 - 1                                                     # Eqn
            for e in side.split("+"): # side is reactants or products            # Eqn
                e = e.strip()                                                    # Eqn
                number = re.findall("^[0-9.\/]*", e)[0]                          # Eqn
                term, number = (e, 1) if number == "" else (e[e.find(number) + len(number):], eval(str(number))) # Eqn
                term = term.strip(); self.system.terms.add(term)                 # Eqn
                if term not in self: self[term] = 0                              # Eqn
                self[term] += sign * number                                      # Eqn
        return self.removeZeros()                                                # Eqn
[docs]    def save(self):                                                              # Eqn
        """Saves this (potentially new) equation to the system, so that it
can be used directly later on"""                                                 # Eqn
        self.system.parse(str(self)); return self                                # Eqn 
[docs]    def __contains__(self, x:str):                                               # Eqn
        """Whether a term is in this equation"""                                 # Eqn
        return x in self.terms                                                   # Eqn 
[docs]    def __getattr__(self, term:str):                                             # Eqn
        """Gets the value of the term in this equation. Negative if on
consumer side, positive if on producer side"""                                   # Eqn
        if term in self.terms: return self.terms[term]                           # Eqn
        else: return 0                                                           # Eqn 
    def __setitem__(self, idx:str, value:float): self.terms[idx] = value; return self # Eqn
[docs]    def __getitem__(self, idx:str):                                              # Eqn
        """Same as :meth:`__getattr__`"""                                        # Eqn
        return getattr(self, idx)                                                # Eqn 
[docs]    def __iter__(self):                                                          # Eqn
        """Yields key:value pairs"""                                             # Eqn
        for k, v in self.terms.items(): yield k, v                               # Eqn 
[docs]    def __len__(self):                                                           # Eqn
        """Returns number of terms in this equation"""                           # Eqn
        return len(self.terms)                                                   # Eqn 
[docs]    def __hash__(self): return hash(tuple(self.terms.keys()))                    # Eqn 
    def __str__(self):                                                           # Eqn
        a = " + ".join((f"{-v}{k}" for k, v in self.terms.items() if v < 0))     # Eqn
        b = " + ".join((f"{v}{k}" for k, v in self.terms.items() if v > 0))      # Eqn
        return f"{a} -> {b}"                                                     # Eqn
[docs]    def copy(self):                                                              # Eqn
        answer = Eqn(self.system)                                                # Eqn
        answer.terms = dict(self.terms); return answer                           # Eqn 
    def __repr__(self, printExtras=None):                                        # Eqn
        space = " " if settings.spaceBetweenValueSymbol else ""                  # Eqn
        def formatValue(value:float):                                            # Eqn
            if abs(value - 1) < 1e-9: return ""                                  # Eqn
            if abs(value - round(value)) < 1e-9:                                 # Eqn
                return f"{round(value)}{space}"                                  # Eqn
            return f"{round(value, 3)}{space}"                                   # Eqn
        a = " + ".join((f"{formatValue(-v)}{k}" for k, v in self.terms.items() if v < 0)) # Eqn
        b = " + ".join((f"{formatValue(v)}{k}" for k, v in self.terms.items() if v > 0)) # Eqn
        answer = f"{a} \033[1m->\033[0m {b}"                                     # Eqn
        printExtras = printExtras if printExtras is not None else settings.eqnPrintExtras # Eqn
        return answer if not printExtras else f"""{answer}. Can...
- "MJ" in eqn: to check whether this equation has a specific term
- eqn["MJ"], or eqn.MJ: to get the actual value of the term
- eqn.terms: to get dict of all term -> values
- eqn["MJ"] = 5: to modify a term's value
- eqn * 2: to use normal math operations on the entire equation
- eqn1 @ eqn2: to try to zero out some common terms, useful for unit conversions
- eqn1 == eqn2: see if 2 equations are the same, scale invariant
- for term, value in eqn: to loop over every term and its value
- len(eqn): to get number of terms in the equation
- eqn.copy()"""                                                                  # Eqn
    def __mul__(self, number:float):                                             # Eqn
        answer = self.copy()                                                     # Eqn
        answer.terms = {k: v*number for k, v in self.terms.items()}              # Eqn
        return answer.removeZeros()                                              # Eqn
    def __rmul__(self, number:float): return self.__mul__(number)                # Eqn
    def __neg__(self): return -1 * self                                          # Eqn
    def __truediv__(self, number:float):                                         # Eqn
        answer = self.copy()                                                     # Eqn
        answer.terms = {k: v/number for k, v in self.terms.items()}              # Eqn
        return answer.removeZeros()                                              # Eqn
    def __rtruediv__(self, number:float): raise Exception("Can't be divided by a number. It doesn't mean anything") # Eqn
    def __add__(self, eqn):                                                      # Eqn
        answer = self.copy(); answer.terms = {}                                  # Eqn
        for term, value in self: answer[term] = value + eqn[term]                # Eqn
        for term, value in eqn:                                                  # Eqn
            if term not in answer: answer[term] = value + self[term]             # Eqn
        return answer.removeZeros()                                              # Eqn
    def __sub__(self, eqn): return self + -1*eqn                                 # Eqn
    def __eq__(self, eqn):                                                       # Eqn
        if len(self) != len(eqn): return False                                   # Eqn
        if set(self.terms.keys()) != set(eqn.terms.keys()): return False         # Eqn
        term = list(self.terms.keys())[0]                                        # Eqn
        eqn = eqn * self[term] / eqn[term]                                       # Eqn
        for term, value in self:                                                 # Eqn
            if abs(self[term] - eqn[term]) > 1e-9: return False                  # Eqn
        return True                                                              # Eqn
[docs]    def sharedTerms(self, eqn:"Eqn") -> List[str]:                               # Eqn
        """Gets a list of shared terms between this equation and the
specified one."""                                                                # Eqn
        ts = set(self.terms.keys())                                              # Eqn
        return [t for t in eqn.terms.keys() if t in ts]                          # Eqn 
[docs]    def join(self, eqn:"Eqn", term:str) -> "Eqn":                                # Eqn
        """Tries to cancel out this equation with another equation at the
specified term. Example::
    s = eqn.System(\"\"\"a + b -> c + d
    c + 2e -> f\"\"\")
    s.a.c.join(s.c.f, "c") # returns the equation "a + b + 2e -> d + f"
For simpler cases, where the shared term to be joined is obvious, use
:meth:`__matmul__` instead"""                                                    # Eqn
        return self + eqn * (-self[term]/eqn[term])                              # Eqn 
[docs]    def __matmul__(self, eqn:"Eqn") -> "Eqn":                                    # Eqn
        """Convenience method that does the same thing as :meth:`join`.
Example::
    s = eqn.System(\"\"\"a + b -> c + d
    c + 2e -> f\"\"\")
    s.a.c @ s.c.f # returns the equation "a + b + 2e -> d + f"
Preference order of which term to join:
1) If term is on producer side of ``self``, and consumer side of ``eqn``
2) If term is on consumer side of ``self``, and producer side of ``eqn``
3) Other cases"""                                                                # Eqn
        sharedTerms = self.sharedTerms(eqn)                                      # Eqn
        def sortF(term):                                                         # Eqn
            if self[term] > 0 and eqn[term] < 0: return 0                        # Eqn
            if self[term] < 0 and eqn[term] > 0: return 1                        # Eqn
            return 2                                                             # Eqn
        sharedTerms = sorted(sharedTerms, key=sortF)                             # Eqn
        if len(sharedTerms) == 0: return None                                    # Eqn
        return self.join(eqn, sharedTerms[0])                                    # Eqn 
[docs]    def round(self, term:str, amount:float=10) -> "Eqn":                         # Eqn
        """Rounds the equation off, so that the term's value
is the specified amount. For aesthetic purposes mainly. Example::
    s = eqn.System("a + b -> 2c")
    s.a.c.round("c", 5) # returns the equation "2.5a + 2.5b -> 5c"'"""           # Eqn
        if term not in self: raise AttributeError(term)                          # Eqn
        return self * amount / self[term]                                        # Eqn 
[docs]    def __round__(self, term:str=None) -> "Eqn":                                 # Eqn
        """Like :meth:`round`, but more Pythonic?
:param term: Can be any of these:
    - None
    - str
    - Union[int, float]
    - Tuple[str, float]"""                                                       # Eqn
        defaultTerm = list(self.terms.keys())[-1]                                # Eqn
        if term is None: return self.round(defaultTerm, 1)                       # Eqn
        elif isinstance(term, (tuple, list)): return self.round(*term)           # Eqn
        elif isinstance(term, str): return self.round(term, 1)                   # Eqn
        elif k1lib.isNumeric(term): return self.round(defaultTerm, term)         # Eqn
        else: raise AttributeError(f"Don't understand {term}")                   # Eqn  
[docs]class Eqns:                                                                      # Eqns
[docs]    def __init__(self, system:"System", eqns:List[Eqn], focusTerm:str=None):     # Eqns
        """Creates a new list of equations. Not expected to be instantiated
by the end user.
:param system: injected :class:`System`
:param eqns: list of equations
:param focusTerm: if the list of equations are from the result of focusing
    in a single term, then use this parameter to prioritize certain search
    parameters.
"""                                                                              # Eqns
        self.system = system; self.eqns = eqns; self.terms = set()               # Eqns
        for eqn in eqns: self.terms.update(eqn.terms.keys())                     # Eqns
        self.focusTerm = focusTerm                                               # Eqns 
[docs]    def __getitem__(self, idx:Union[int, str]) -> Optional[Eqn]:                 # Eqns
        """If int, return the equation with that index. Not really helpful
for exploring the system of equations, but good for automated scripts
If string, then effectively the same as :meth:`__getattr__`
"""                                                                              # Eqns
        return self.eqns[idx] if isinstance(idx, int) else getattr(self, idx)    # Eqns 
[docs]    def __getattr__(self, term:str) -> Optional[Eqn]:                            # Eqns
        """Picks out a specific :class:`Eqn` that has the specified term.
Prefer shorter equations, and the returned :class:`Eqn` always have the
term on the products side. Meaning::
    eqns = eqn.System("a + 2b -> c").b # gets an Eqns object with that single equation
    eqns.a # gets the equation "c -> a + 2b" instead
This is a convenience way to search for equations. If you need more
granularity, use :meth:`pick` instead"""                                         # Eqns
        chosenEqns = []                                                          # Eqns
        for eqn in self.eqns:                                                    # Eqns
            if term in eqn:                                                      # Eqns
                chosenEqns.append(eqn if eqn[term] > 0 else -eqn)                # Eqns
        chosenEqns = sorted(chosenEqns, key=lambda eqn: len(eqn))                # Eqns
        return None if len(chosenEqns) == 0 else chosenEqns[0]                   # Eqns 
[docs]    def pick(self, *terms:List[str]) -> Optional[Eqn]:                           # Eqns
        """Like the quick method (:meth:`__getattr__`), but here, picks
equations more carefully, with selection for multiple terms. Example::
    s = eqn.System(\"\"\"a + 2b -> c
    b + c -> d
    a -> 3d
    a + b + c -> 2d\"\"\")
    s.a.pick("b", "d") # returns last equation
As you can see, it's impossible to pick out the last equation using
:meth:`__getattr__` alone, as they will all prefer the shorter equations,
so this is where :meth:`pick` can be useful."""                                  # Eqns
        chosenEqns = []; t = self.focusTerm or terms[0]                          # Eqns
        for eqn in self.eqns:                                                    # Eqns
            if all((term in eqn for term in terms)):                             # Eqns
                chosenEqns.append(eqn if eqn[t] > 0 else -eqn)                   # Eqns
        chosenEqns = sorted(chosenEqns, key=lambda eqn: len(eqn))                # Eqns
        return None if len(chosenEqns) == 0 else chosenEqns[0]                   # Eqns 
[docs]    def __dir__(self):                                                           # Eqns
        """Returns the list of terms in every equation here. Useful for
tab completion."""                                                               # Eqns
        return list(self.terms)                                                  # Eqns 
    def __repr__(self):                                                          # Eqns
        end = """Can...
- eqns[i]: to get the 'i'th equation
- eqns.C: to pick out the first equation that has term 'C'"""                    # Eqns
        if self.focusTerm == None:                                               # Eqns
            eqns = "\n".join([f"{i}. {eqn.__repr__(printExtras=False)}" for i, eqn in enumerate(self.eqns)]) # Eqns
            return f"""Equations:\n{eqns}\n\n{end}"""                            # Eqns
        else:                                                                    # Eqns
            consumingEqns = []; producingEqns = []                               # Eqns
            for eqn in self.eqns:                                                # Eqns
                if eqn[self.focusTerm] < 0:                                      # Eqns
                    consumingEqns.append(f"{eqn.__repr__(printExtras=False)}")   # Eqns
                else: producingEqns.append(f"{eqn.__repr__(printExtras=False)}") # Eqns
            consumingEqns = "\n".join([f"{i}. {eqn}" for i, eqn in enumerate(consumingEqns)]) # Eqns
            producingEqns = "\n".join([f"{i}. {eqn}" for i, eqn in enumerate(producingEqns)]) # Eqns
            return f"""Consumers:\n{consumingEqns}\n\nProducers:\n{producingEqns}\n\n{end}""" # Eqns 
[docs]class System:                                                                    # System
[docs]    def __init__(self, strToParse:str=None):                                     # System
        """Creates a new system of equations.
:param strToParse: if specified, then it gets feed into :meth:`parse`"""         # System
        self.terms = set()                                                       # System
        self.eqns = []                                                           # System
        if strToParse is not None: self.parse(strToParse)                        # System 
    def parse(self, lines:str) -> "System":                                      # System
        """Parses extra equations and saves them to this :class:`System`"""      # System
        lines = (line for line in lines.split("\n") if line != "" and not line.startswith("#")) # System
        self.eqns += [Eqn(self).parse(line) for line in lines if not line.startswith("#")] # System
        self.eqns = list(set(self.eqns))                                         # System
        return self                                                              # System
[docs]    def spellCheck(self):                                                        # System
        """Runs a spell check to find out terms that are pretty similar to
each other"""                                                                    # System
        print("Similar terms:"); terms = list(self.terms)                        # System
        for i, iTerm in enumerate(terms):                                        # System
            for j, jTerm in enumerate(terms[i+1:]):                              # System
                if iTerm[:-1] == jTerm[:-1]: continue                            # System
                if abs(len(iTerm) - len(jTerm)) > 2: continue                    # System
                r = difflib.SequenceMatcher(None, iTerm, jTerm).ratio()          # System
                if r < 0.9: continue                                             # System
                print(f"- {round(r*100)}% similar: {iTerm}, {jTerm}")            # System 
[docs]    def __len__(self): return len(self.eqns)                                     # System 
[docs]    def __getitem__(self, idx:int) -> Eqn:                                       # System
        """Picks out the i'th equation from the list of equations. Useful
for automated scripts"""                                                         # System
        return self.eqns[idx]                                                    # System 
[docs]    def __getattr__(self, term:str) -> Eqns:                                     # System
        """Picks out equations that has the term"""                              # System
        return Eqns(self, [eqn for i, eqn in enumerate(self.eqns) if term in eqn], focusTerm=term) # System 
[docs]    def __dir__(self):                                                           # System
        """Returns the list of terms in every equation here. Useful for
tab completion."""                                                               # System
        return list(self.terms)                                                  # System 
    def __repr__(self):                                                          # System
        return f"""System of {len(self)} equations:\n{Eqns(self, self.eqns)}\n
Can...
- s[i]: to get a specific equation
- s.C: to get all equations that involve a specific substance "C"
- s.spellCheck(): to check if there are terms that are close to each other
"""                                                                              # System