# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
import re, logging, warnings
from typing import Union, Tuple, List
from .atom import Atom, mo, NoFreeElectrons, OctetFull
from .substance import alkane, alcohol
import k1lib.cli as cli
__all__ = ["parse"]
mainNumsDict = {"meth": 1, "eth": 2, "prop": 3, "but": 4, "pent": 5, "hex": 6, "hept": 7, "oct": 8, "non": 9, "dec": 10, "undec": 11, "dodec": 12, "tridec": 13, "tetradec": 14, "pentadec": 15, "hexadec": 16, "heptadec": 17, "octadec": 18, "nonadec": 19}
bondsDict = {"ane": 1, "ene": 2, "yne": 3}
cis_trans = "(?P<cis_trans>[, \-]?((cis)|(trans))$)"
def join(a): return "(" + "|".join(f"({p})" for p in a) + ")"                    # join
pres1 = ["n-", "o-", "p-"]                                                       # join
pres2 = ["mono", "di", "tri", "tetra", "penta", "hexa"]                          # join
pres3 = ["iso", "cyclo", "sec-", "tert-"]                                        # join
pres = f"{join(pres1)}?[\-]?{join(pres2)}?[\-]?{join(pres3)}?"                   # join
main_pre = f"(?P<main_pre>(((cis-)|(trans-))?[0-9,\-])*{pres})"                  # join
#main_pre = "(?P<main_pre>(n-)?(cyclo))"                                         # join
mids = list(mainNumsDict.keys()) + ["benz", "naphthal", "pyridine", "xylenol", "ether", "amine"] # join
main_mid = "(?P<main_mid>" + "|".join(f"({m})" for m in mids) + ")"              # join
sufs = list(bondsDict.keys()) + ["anol", "ate", "adiene", "aldehyde", "anone"]   # join
main_suf = f"(?P<main_suf>{join(sufs)})"                                         # join
main = f"(?P<main>{main_pre}?{main_mid}{main_suf}?{cis_trans}?)$"; mainC = re.compile(f"^{main}") # join
branch_pre = f"(?P<branch_pre>(((cis-)|(trans-))?[0-9,\-])*{pres})"              # join
bMids = mids + ["fluoro", "chloro", "bromo", "iodo", "nitro", "phenyl", "perfluoro"] # join
branch_mid = "(?P<branch_mid>" + join(bMids) + ")"                               # join
branch_suf = "(?P<branch_suf>yl)"                                                # join
branch = f"(?P<branch>{branch_pre}?{branch_mid}{branch_suf}?[ ]?)"; branchC = re.compile(f"^{branch}$") # join
fullStrict = f"^{branch}*{main}"                                                 # join
fullStrictC = re.compile(fullStrict)                                             # join
full = f"{branch}*{main}"                                                        # join
fullC = re.compile(full)                                                         # join
def recognizeInit(ogS:str) -> Union[Tuple[List[str], List[str]], str]:           # recognizeInit
    """Tries to recognize the input sequence. Returns the sequence if can't understand, else
returns tuple (mainParts, branches)."""                                          # recognizeInit
    s = ogS; m = re.fullmatch(fullStrictC, s)                                    # recognizeInit
    if m is None: return ogS                                                     # recognizeInit
    s = s[:m.start("main")]                                                      # recognizeInit
    sMain = m.group("main")                                                      # recognizeInit
    sMains = [m.group("main_pre") or "", m.group("main_mid") or "", m.group("main_suf") or ""] # recognizeInit
    sMains = ["".join(sMains), *sMains]                                          # recognizeInit
                                                                                 # recognizeInit
    groups = []                                                                  # recognizeInit
    for _ in range(int(10)):                                                     # recognizeInit
        m = re.match(branch, s)                                                  # recognizeInit
        if m is None or m.end() == 0: break                                      # recognizeInit
        groups.append(s[:m.end()].strip("-,"))                                   # recognizeInit
        s = s[m.end():]                                                          # recognizeInit
    if len(s) != 0: return None                                                  # recognizeInit
    return [sMains, groups]                                                      # recognizeInit
def preLocs(sPre, *defaults) -> List[int]:                                       # preLocs
    """If ``sPre`` looks like "3,1,4-" or sth like that, then extract all the location infos.""" # preLocs
    answer = []                                                                  # preLocs
    try: answer = sPre.strip("-,").split(",") | cli.replace("-", "") | cli.toInt() | cli.deref() # preLocs
    except: pass                                                                 # preLocs
    return defaults if len(answer) == 0 else answer                              # preLocs
def detectModifiers(sPre:str, has:set) -> str:                                   # detectModifiers
    if "cyclo" in sPre: has.add("cyclo"); sPre = sPre.replace("cyclo", "")       # detectModifiers
    if "iso" in sPre: has.add("iso"); sPre = sPre.replace("iso", "")             # detectModifiers
    if "sec" in sPre: has.add("sec"); sPre = sPre.replace("sec", "")             # detectModifiers
    if "tert" in sPre: has.add("tert"); sPre = sPre.replace("tert", "")          # detectModifiers
    return sPre                                                                  # detectModifiers
def fixTridec(sPre:str, sMid:str) -> Tuple[str, str]:                            # fixTridec
    """Handles special case for tridec, tetradec, ..., as "tri" leaks into sPre and not sMid for detection""" # fixTridec
    if sMid == "dec": # special case for tridec, tetradec and whatnot            # fixTridec
        for prefix in ["tri", "tetra", "penta", "hexa"]:                         # fixTridec
            if sPre.endswith(prefix):                                            # fixTridec
                sPre = sPre[:-len(prefix)]                                       # fixTridec
                sMid = prefix + sMid; break                                      # fixTridec
    return sPre, sMid                                                            # fixTridec
def recognizeMain(a:Union[List[str], str]) -> Union[Atom, str]:                  # recognizeMain
    """Tries to recognize main part of the molecule (eg. n-cyclohexane). Returns
:class:`~k1lib.mo.Atom` if recognized, else the input string"""                  # recognizeMain
    s = ""                                                                       # recognizeMain
    try:                                                                         # recognizeMain
        if isinstance(a, str):                                                   # recognizeMain
            m = re.fullmatch(mainC, a)                                           # recognizeMain
            a = [a, m.group("main_pre"), m.group("main_mid"), m.group("main_suf")] # recognizeMain
        s, sPre, sMid, sSuf = a                                                  # recognizeMain
        sPre, sMid, sSuf = (sPre or "").strip().lstrip("-"), (sMid or "").strip(), (sSuf or "").strip() # recognizeMain
        has = set(); sPre = detectModifiers(sPre, has)                           # recognizeMain
        if sMid in mainNumsDict:                                                 # recognizeMain
            sPre, sMid = fixTridec(sPre, sMid)                                   # recognizeMain
            carbons = mainNumsDict[sMid]                                         # recognizeMain
            if sSuf in bondsDict:                                                # recognizeMain
                a = alkane(carbons); bonds = bondsDict[sSuf] # main atom         # recognizeMain
                if bonds > 1: b = a.next(times=preLocs(sPre, 1)[0]-1); b.bond(b.next(), bonds-1) # recognizeMain
            elif sSuf == "anol":                                                 # recognizeMain
                a = alcohol(carbons, preLocs(sPre, 1)[0]) # main atom            # recognizeMain
            elif sSuf == "adiene":                                               # recognizeMain
                a = alkane(carbons)                                              # recognizeMain
                for loc in preLocs(sPre, 1, 2):                                  # recognizeMain
                    b = a.next(times=loc-1); b(b.next())                         # recognizeMain
            else: return s                                                       # recognizeMain
        elif sMid == "xylenol" and sSuf.strip() == "":                           # recognizeMain
            a = mo.benzene; a(mo.H2O)                                            # recognizeMain
            for loc in preLocs(sPre): a.next(times=loc-1)(mo.CH4)                # recognizeMain
        else: return s                                                           # recognizeMain
        if "cyclo" in has: a.endChain.bond(a)                                    # recognizeMain
        if "iso" in has: a.moveLastCTo2ndC()                                     # recognizeMain
        if "tert" in has: a.moveLastCTo2ndC().moveLastCTo2ndC()                  # recognizeMain
        return a                                                                 # recognizeMain
    except Exception as e:                                                       # recognizeMain
        print(s); raise e                                                        # recognizeMain
diTriTetra = {term: i+1 for i, term in enumerate(pres2)}                         # recognizeMain
def detectDiTriTetra(sPre:str) -> Tuple[str, str]:                               # detectDiTriTetra
    """If exists, returns (sPre, num), else return (sPre, None)"""               # detectDiTriTetra
    for k, v in diTriTetra.items():                                              # detectDiTriTetra
        if sPre.endswith(k):                                                     # detectDiTriTetra
            return sPre.replace(k, ""), v                                        # detectDiTriTetra
    return sPre, None                                                            # detectDiTriTetra
def detectCisTrans(sPre:str, has) -> str:                                        # detectCisTrans
    for term in ["cis", "trans"]:                                                # detectCisTrans
        if term in sPre:                                                         # detectCisTrans
            has.add(term); sPre = sPre.replace(term, "")                         # detectCisTrans
    return sPre                                                                  # detectCisTrans
def recognizeGroup(s:str, mainA:Atom) -> Union[Atom, str]:                       # recognizeGroup
    """Tries to recognize a group (eg. "1-cyclopentyl"), and adds it to the main atom ``mainA``.

:return: if can't recognize, then returns the group string, else returns the main atom""" # recognizeGroup
    if mainA is None: mainA = alkane(10)                                         # recognizeGroup
    m = re.fullmatch(branchC, s); ops = [] # operations to generate the branch   # recognizeGroup
    sPre, sMid, sSuf = [m.group("branch_pre"), m.group("branch_mid"), m.group("branch_suf")] # recognizeGroup
    has = set(); sPre = detectCisTrans(detectModifiers(sPre, has), has)          # recognizeGroup
    if sMid in mainNumsDict:                                                     # recognizeGroup
        if sSuf != "yl": return s                                                # recognizeGroup
        carbons = mainNumsDict[sMid]                                             # recognizeGroup
        ops.append(lambda _: alkane(carbons))                                    # recognizeGroup
        if "cyclo" in has: ops.append(lambda a: a.endChain.bond(a))              # recognizeGroup
        if "iso" in has: ops.append(lambda a: a.moveLastCTo2ndC())               # recognizeGroup
        if "tert" in has: ops.append(lambda a: a.moveLastCTo2ndC().moveLastCTo2ndC()) # recognizeGroup
    elif sMid == "chloro": ops.append(lambda _: mo.Cl(mo.H))                     # recognizeGroup
    elif sMid == "fluoro": ops.append(lambda _: mo.F (mo.H))                     # recognizeGroup
    elif sMid == "bromo":  ops.append(lambda _: mo.Br(mo.H))                     # recognizeGroup
    elif sMid == "iodo":   ops.append(lambda _: mo.I (mo.H))                     # recognizeGroup
    elif sMid == "phenyl": ops.append(lambda _: mo.benzene)                      # recognizeGroup
    else: return s#, sPre, sMid, sSuf                                            # recognizeGroup
    def gen(a=None):                                                             # recognizeGroup
        for op in ops: a = op(a)                                                 # recognizeGroup
        return a                                                                 # recognizeGroup
    # now attaching group to main chain...                                       # recognizeGroup
    sPre, locWord = detectDiTriTetra(sPre); locs = preLocs(sPre)                 # recognizeGroup
    logging.debug(f"locs: {locs}, locWord: {locWord}, sPre: {sPre}")             # recognizeGroup
    if len(locs) > 0 and locWord is not None: # both available, check consistency # recognizeGroup
        logging.debug(f"locs and locWord both available")                        # recognizeGroup
        assert locWord == len(locs), f"s: {s}, locWord: {locWord}, locs: {locs}" # recognizeGroup
        for loc in locs: mainA.next(times=loc-1)(gen())                          # recognizeGroup
        return mainA                                                             # recognizeGroup
    elif len(locs) > 0 and locWord is None:                                      # recognizeGroup
        if len(locs) > 1: warnings.warn(f"Group {s} has attachment positions, but not di,tri,tetra words!") # recognizeGroup
        for loc in locs: mainA.next(times=loc-1)(gen())                          # recognizeGroup
        return mainA                                                             # recognizeGroup
    else:                                                                        # recognizeGroup
        logging.debug(f"Can't detect any position, will be adding stuff randomly to the molecule") # recognizeGroup
        a = mainA                                                                # recognizeGroup
        for _ in range(max(len(locs), locWord or 0, 1)):                         # recognizeGroup
            try:                                                                 # recognizeGroup
                while a.availableBonds == 0: a = a.next()                        # recognizeGroup
            except RecursionError: raise RuntimeError(f"Molecule has no more free carbons to attach group {s} to") # recognizeGroup
            a(gen())                                                             # recognizeGroup
            pass                                                                 # recognizeGroup
        return mainA                                                             # recognizeGroup
def parse(s:str, quiet:bool=False) -> Union[Atom, str]:                          # parse
    """Tries to recognize molecule. Returns molecule :class:`~k1lib.mo.Atom`
if recognized, else the molecule's string."""                                    # parse
    def p(s):                                                                    # parse
        if not quiet: print(s)                                                   # parse
    try:                                                                         # parse
        # initial recognition                                                    # parse
        out = recognizeInit(s); has = set()                                      # parse
        if isinstance(out, str): return s                                        # parse
        mainParts, groups = out                                                  # parse
        # form main chain                                                        # parse
        mainA = recognizeMain(mainParts)                                         # parse
        if isinstance(mainA, str):                                               # parse
            p(f"Recognized {s}, but couldn't form the main chain {mainParts[0]}"); return s # parse
        # check for special perfluoro branch case                                # parse
        gs = []                                                                  # parse
        for g in groups:                                                         # parse
            if "perfluoro" in g:                                                 # parse
                if "perfluoro" != g: p(f"For some reason, group {g} is not exactly 'perfluoro'."); return s # parse
                has.add("perfluoro")                                             # parse
            else: gs.append(g)                                                   # parse
        for g in gs:                                                             # parse
            mainA = recognizeGroup(g, mainA)                                     # parse
            if isinstance(mainA, str):                                           # parse
                p(f"Recognized {s}, the main chain too, but don't recognize group {g}") # parse
        if "perfluoro" in has: mainA.perfluoro_ize()                             # parse
        return mainA                                                             # parse
    except OctetFull:                                                            # parse
        p("Some of the atoms can't form new bonds, as their octet is full. Please revise your structure!"); return s # parse
    except NoFreeElectrons:                                                      # parse
        p("Some of the atoms can't form new bonds, as they have no free electrons left to contribute to the bond. Please revise your structure!"); return s # parse
