from re import findall


def molar_mass(element):
    "get molar mass (g/mol) of chemistry elements"
    d = {
        "H": 1.00794,
        "He": 4.002602,
        "Li": 6.941,
        "Be": 9.012182,
        "B": 10.811,
        "C": 12.0107,
        "N": 14.0067,
        "O": 15.9994,
        "F": 18.9984032,
        "Ne": 20.1797,
        "Na": 22.98976928,
        "Mg": 24.305,
        "Al": 26.9815386,
        "Si": 28.0855,
        "P": 30.973762,
        "S": 32.065,
        "Cl": 35.453,
        "Ar": 39.948,
        "K": 39.0983,
        "Ca": 40.078,
        "Sc": 44.955912,
        "Ti": 47.867,
        "V": 50.9415,
        "Cr": 51.9961,
        "Mn": 54.938045,
        "Fe": 55.845,
        "Co": 58.933195,
        "Ni": 58.6934,
        "Cu": 63.546,
        "Zn": 65.409,
        "Ga": 69.723,
        "Ge": 72.64,
        "As": 74.9216,
        "Se": 78.96,
        "Br": 79.904,
        "Kr": 83.798,
        "Rb": 85.4678,
        "Sr": 87.62,
        "Y": 88.90585,
        "Zr": 91.224,
        "Nb": 92.90638,
        "Mo": 95.94,
        "Tc": 98.9063,
        "Ru": 101.07,
        "Rh": 102.9055,
        "Pd": 106.42,
        "Ag": 107.8682,
        "Cd": 112.411,
        "In": 114.818,
        "Sn": 118.71,
        "Sb": 121.760,
        "Te": 127.6,
        "I": 126.90447,
        "Xe": 131.293,
        "Cs": 132.9054519,
        "Ba": 137.327,
        "La": 138.90547,
        "Ce": 140.116,
        "Pr": 140.90465,
        "Nd": 144.242,
        "Pm": 146.9151,
        "Sm": 150.36,
        "Eu": 151.964,
        "Gd": 157.25,
        "Tb": 158.92535,
        "Dy": 162.5,
        "Ho": 164.93032,
        "Er": 167.259,
        "Tm": 168.93421,
        "Yb": 173.04,
        "Lu": 174.967,
        "Hf": 178.49,
        "Ta": 180.9479,
        "W": 183.84,
        "Re": 186.207,
        "Os": 190.23,
        "Ir": 192.217,
        "Pt": 195.084,
        "Au": 196.966569,
        "Hg": 200.59,
        "Tl": 204.3833,
        "Pb": 207.2,
        "Bi": 208.9804,
        "Po": 208.9824,
        "At": 209.9871,
        "Rn": 222.0176,
        "Fr": 223.0197,
        "Ra": 226.0254,
        "Ac": 227.0278,
        "Th": 232.03806,
        "Pa": 231.03588,
        "U": 238.02891,
        "Np": 237.0482,
        "Pu": 244.0642,
        "Am": 243.0614,
        "Cm": 247.0703,
        "Bk": 247.0703,
        "Cf": 251.0796,
        "Es": 252.0829,
        "Fm": 257.0951,
        "Md": 258.0951,
        "No": 259.1009,
        "Lr": 262,
        "Rf": 267,
        "Db": 268,
        "Sg": 271,
        "Bh": 270,
        "Hs": 269,
        "Mt": 278,
        "Ds": 281,
        "Rg": 281,
        "Cn": 285,
        "Nh": 284,
        "Fl": 289,
        "Mc": 289,
        "Lv": 292,
        "Ts": 294,
        "Og": 294,
        "": 0,
    }

    return d[element]


def formula_parser(formula: str):
    """
    a chemical formular parser

    :param formula: formula string
    :returns: dictionary with element (key) and number (value)
    """
    s = findall("([A-Z][a-z]?)([0-9]*)", formula)
    d = dict(s)
    for k, v in d.items():
        if v == "":
            d[k] = 1
        else:
            d[k] = int(v)
    return d


def molecular_weight(formula: str):
    """
    Calculate molecular weight using a chemical formula

    :param formula: chemical formula
    :returns: molecular weight

    ------
    Example
    ------
    >>> molecular_weight("C6H12O6")
    >>> 180.15588
    """
    "calculate molecular weight"
    d = formula_parser(formula)
    mw = sum([molar_mass(k) * v for k, v in d.items()])
    return mw


def pure_unit(unit: str):
    """
    Get a pure unit string by removing elements in string (e.g., mmol C d-1)

    :param unit: unit string with elements
    :returns: pure unit string

    ------
    Example
    ------
    >>> pure_unit("mmol C d-1")
    >>> "mmol d-1"
    """
    # an element pool
    l = ["C", "N", "P", "O", "Fe", "Si", "Ca", "S"]
    # add leading whitespace
    l = [" " + i for i in l]

    # remove element in unit string
    for element in l:
        if element in unit:
            unit = unit.replace(element, "")

    return unit


def format_base_unit(input_str, latex_format=True):
    """
    Convert the base unit to the desired format
    
    base unit means it is not a compound unit (e.g., m2, m-2, m/s, m s-1),
    which should be handled by the function `format_unit`
            
    :param input_str: a base unit string 
    :return: a formatted string
    
    ----------
    Example:
    
    >>> convert_a_unit('m2')
    '$m^2$'
    >>> convert_a_unit('m-2')
    '$m^{-2}$'
    """
    if input_str == 'n/a':
        return None

    ## if the input string contains a '^' character, then remove it
    if '^' in input_str:
        input_str = input_str.replace('^', '')

    if '**' in input_str:
        input_str = input_str.replace('**', '^')
        
    if 'o/oo' in input_str:
        input_str = input_str.replace('o/oo', '‰')
        
    if 'degree' in input_str:
        input_str = input_str.replace('degree', '°')

    ## when exponent is negative  (e.g., m-2, m-3)
    if '-' in input_str:
        parts = input_str.split('-')
        ## two parts: base and exponent
        base = parts[0]
        exponent = parts[1]
        
        if len(parts) == 2 and exponent.isdigit():
            exponent = int(exponent)
            if exponent != 0:
                # Convert to the desired format using string formatting
                # 3 layers of curly braces are needed to print a single curly brace
                output_str = f"{parts[0]}$^{{-{exponent}}}$"
                ## if not use LaTeX format, then remove dollar sign
                if not latex_format: output_str = output_str.replace('$', '')                
                return output_str
            
    ## when exponent is positive (e.g., m2, m3)
    ## also slice into two parts: base and exponent
    if len(input_str) > 1:
        base = input_str[:-1]
        exponent = input_str[-1]
        
        if exponent.isdigit():
            exponent = int(exponent)
            if exponent != 0:
                # Convert to the desired format using string formatting
                # 3 layers of curly braces are needed to print a single curly brace
                output_str = f"{base}$^{{{exponent}}}$"
                ## if not use LaTeX format, then remove dollar sign
                if not latex_format: output_str = output_str.replace('$', '')
                return output_str
    
    return input_str

def format_unit(input, *args, **kwargs):
    """
    convert a unit in a string to the desired format, which can be used in plot labels and 
    facilitate the unit calculation
    
    :param input: a unit string 
    :return: a  formatted string
    
    ----------
    Example
    
    >>> format_unit('mol kg-1')
    'mol kg$^{-1}$'
    >>> format_unit('mmol m-2 d-1')
    'mmol m$^{-2}$ d$^{-1}$'
    """
    if input == 'n/a':
        return None
    unit_list = input.split(" ")
    unit_list = [format_base_unit(i, *args, **kwargs) for i in unit_list]

    clean_string = ' '.join(unit_list)
    return clean_string
