"""
Utility functions for handling the units in `Equations`.
"""

import re

from brian2.core.variables import Variable
from brian2.parsing.expressions import parse_expression_dimensions
from brian2.parsing.statements import parse_statement
from brian2.units.fundamentalunits import (
    fail_for_dimension_mismatch,
    get_dimensions,
    get_unit,
)

__all__ = ["check_dimensions", "check_units_statements"]


def check_dimensions(expression, dimensions, variables):
    """
    Compares the physical dimensions of an expression to expected dimensions in
    a given namespace.

    Parameters
    ----------
    expression : str
        The expression to evaluate.
    dimensions : `Dimension`
        The expected physical dimensions for the `expression`.
    variables : dict
        Dictionary of all variables (including external constants) used in
        the `expression`.

    Raises
    ------
    KeyError
        In case on of the identifiers cannot be resolved.
    DimensionMismatchError
        If an unit mismatch occurs during the evaluation.
    """
    expr_dims = parse_expression_dimensions(expression, variables)
    expected = repr(get_unit(dimensions))
    err_msg = (
        f"Expression '{expression.strip()}' does not have the expected unit {expected}"
    )
    fail_for_dimension_mismatch(expr_dims, dimensions, err_msg)


def check_units_statements(code, variables):
    """
    Check the units for a series of statements. Setting a model variable has to
    use the correct unit. For newly introduced temporary variables, the unit
    is determined and used to check the following statements to ensure
    consistency.

    Parameters
    ----------
    code : str
        The statements as a (multi-line) string
    variables : dict of `Variable` objects
        The information about all variables used in `code` (including
        `Constant` objects for external variables)

    Raises
    ------
    KeyError
        In case on of the identifiers cannot be resolved.
    DimensionMismatchError
        If an unit mismatch occurs during the evaluation.
    """
    variables = dict(variables)
    # Avoid a circular import
    from brian2.codegen.translation import analyse_identifiers

    newly_defined, _, unknown = analyse_identifiers(code, variables)

    if len(unknown):
        raise AssertionError(
            "Encountered unknown identifiers, this should not "
            f"happen at this stage. Unknown identifiers: {unknown}"
        )

    code = re.split(r"[;\n]", code)
    for line in code:
        line = line.strip()
        if not len(line):
            continue  # skip empty lines

        varname, op, expr, comment = parse_statement(line)
        if op in ("+=", "-=", "*=", "/=", "%="):
            # Replace statements such as "w *=2" by "w = w * 2"
            expr = f"{varname} {op[0]} {expr}"
        elif op == "=":
            pass
        else:
            raise AssertionError(f'Unknown operator "{op}"')

        expr_unit = parse_expression_dimensions(expr, variables)

        if varname in variables:
            expected_unit = variables[varname].dim
            fail_for_dimension_mismatch(
                expr_unit,
                expected_unit,
                "The right-hand-side of code "
                'statement "%s" does not have the '
                "expected unit {expected}" % line,
                expected=expected_unit,
            )
        elif varname in newly_defined:
            # note the unit for later
            variables[varname] = Variable(
                name=varname, dimensions=get_dimensions(expr_unit), scalar=False
            )
        else:
            raise AssertionError(
                f"Variable '{varname}' is neither in the variables "
                "dictionary nor in the list of undefined "
                "variables."
            )
