varius

from typing import Any, Dict

VARIABLE_STORAGE: Dict[str, Dict[str, float]] = {"default": dict()}
EXPRESSION_STORAGE: Dict[str, Dict[str, float]] = {"default": dict()}


class MagicGlobals:
    latex: bool = True  # use latex
    cv: str = "default"  # current version
    float_digit: int = 5
    repr_indent: int = 2


def reset_version(version: str):
    if version not in VARIABLE_STORAGE:
        raise KeyError(f"The version `{version}` does not exist.")
    VARIABLE_STORAGE[version] = dict()
    EXPRESSION_STORAGE[version] = dict()
    if MagicGlobals.cv == version:
        MagicGlobals.cv == "default"


def reset_all():
    VARIABLE_STORAGE.clear()
    EXPRESSION_STORAGE.clear()
    VARIABLE_STORAGE["default"] = dict()
    EXPRESSION_STORAGE["default"] = dict()
    MagicGlobals.cv == "default"


from .printer import *
from .scope import Scope as note
from .variable import Expression, Variable


def set_latex(use: bool = True):
    MagicGlobals.latex = use
    if use:
        latex_use_cdot()


set_latex(use=is_ipython())


vr = Variable
ex = Expression

__all__ = ["note", "vr", "ex", "show"]

__version__ = "0.1.10"

class Scope:
    "Scope: alias `note`"

    def __init__(self, version: str = "default", copy: Optional[str] = None):
        self.prev = G.cv
        self.version = version
        if version not in VS:
            new_version(version)
        if copy is not None:
            self.load(copy)

    @property
    def variables(self) -> Dict:
        return {latex_to_plain(k.name): v for k, v in VS[self.version].items()}

    @property
    def expressions(self) -> Dict:
        return {latex_to_plain(k): v for k, v in ES[self.version].items()}

    def load(self, other_version):
        duplicate(self.version, other_version)

    def __enter__(self):
        G.cv = self.version
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        G.cv = self.prev

    def __repr__(self) -> str:
        INDENT = G.repr_indent
        lines = []
        lines.append(f"Version: {self.version}")

        lines.append(" " * INDENT + "Variables:")

        for k, v in self.variables.items():
            v = latex_to_plain(str(v))
            lines.append(" " * INDENT * 2 + f"{k} = {v}")
        lines.append(" " * INDENT + "Expressions:")
        for k, v in self.expressions.items():
            v = latex_to_plain(str(v))
            lines.append(" " * INDENT * 2 + f"{k} = {v}")

        return "\n".join(lines)

class Variable(sympy.Symbol):
    """An abstract variable that represents a numerical quantity."""

    _registry: Dict = dict()

    @classmethod
    def ls(cls) -> List:
        return list(cls._registry.values())

    def __new__(
        cls,
        name: str,
        value: Optional[Union[float, int]] = None,
        is_text: bool = True,
        **assumptions: Any,
    ):
        if is_text:
            name = r"\text{" + name + r"}"
        instance = super(Variable, cls).__new__(cls, name, **assumptions)

        if G.cv is not None and value is not None:
            instance._set_value(G.cv, value)

        cls._registry[instance.plain_name] = instance

        return instance

    @property
    def plain_name(self):
        return latex_to_plain(self.name)

    @property
    def value(self) -> Union[float, int]:
        if G.cv is not None:
            return self.get(G.cv)
        else:
            raise RuntimeError("Current version is `None`.")

    @value.setter
    def value(self, value: Union[float, int]):

        if G.cv is not None:
            self._set_value(G.cv, value)
        else:
            raise RuntimeError("Current version is `None`.")

    def __call__(self, value: Union[float, int]):
        self.value = value

    def get(self, version: str) -> Optional[Union[float, int]]:

        assert isinstance(version, str)

        if version in VS:
            if self in VS[version]:
                return sympy.Float(VS[version][self], G.float_digit)
            else:
                return None
        else:
            raise KeyError(f"Version `{version}` does not exist.")

    def _set_value(self, version: str, value: Union[float, int]):
        if not isinstance(value, (float, int)):
            raise TypeError(
                f"Assigned value should be float or int but get {type(value)}"
            )
        if version not in VS:
            VS[version] = dict()
            ES[version] = dict()

        VS[version][self] = sympy.Float(value, G.float_digit)

    def __repr__(self) -> str:
        return self.plain_name

    __str__ = __repr__

    @property
    def latex_repr(self):
        try:
            v = sympy.latex(self.value)
            return sympy.latex(self) + f"= {v}"
        except KeyError:
            return sympy.latex(self)

class Expression(numbers.Number):
    """Expression in symbolic variables."""

    _registry: Dict = dict()

    @classmethod
    def ls(cls) -> List:
        return list(cls._registry.values())

    def __init__(self, name: str, expr, is_text: bool = True):
        if is_text:
            name = r"\text{" + name + r"}"
        self.name = name
        self.expr = expr

        self._registry[self.plain_name] = self

    @property
    def plain_name(self):
        return latex_to_plain(self.name)

    @property
    def plain_expr(self):
        return latex_to_plain(r"{}".format(self.expr))

    def __call__(self, version: Optional[str] = None):
        if version is None:
            version = G.cv
        res = eval_expr(self.expr, version)
        ES[version][self.plain_name] = res

        return res

    def __add__(self, other):
        if isinstance(other, Expression):
            return self.expr + other.expr
        else:
            return self.expr + other

    def __sub__(self, other):
        if isinstance(other, Expression):
            return self.expr - other.expr
        else:
            return self.expr - other

    def __mul__(self, other):
        if isinstance(other, Expression):
            return self.expr * other.expr
        else:
            return self.expr * other

    def __truediv__(self, other):
        if isinstance(other, Expression):
            return self.expr / other.expr
        else:
            return self.expr / other

    def __floordiv__(self, other):
        if isinstance(other, Expression):
            return self.expr // other.expr
        else:
            return self.expr // other

    def __pow__(self, other):
        if isinstance(other, Expression):
            return self.expr ** other.expr
        else:
            return self.expr ** other

    __div__ = __truediv__

    def grad(
        self, *args: Variable, evaluate: bool = True, version: Optional[str] = None
    ) -> Dict:
        if len(args) > 0:
            grads = {x: Gradient(self, x) for x in args}
        else:
            grads = {x: Gradient(self, x) for x in Variable._registry.values()}

        if evaluate:
            grads = {v: g(version) for v, g in grads.items()}

        return grads

    def diff(self, variable: Variable):
        return self.expr.diff(variable)

    @property
    def value(self):
        return self.__call__()

    def __repr__(self):
        return f"{self.plain_name} = {self.plain_expr}"

    __str__ = __repr__

    @property
    def latex_repr(self, evaluate: bool = True, version: Optional[str] = None):
        lhs = self.name
        rhs = sympy.latex(self.expr)
        eq = lhs + " = " + rhs
        if not evaluate:
            return eq

        res = self.__call__(version)

        res = sympy.latex(res)
        if res == rhs:
            return eq
        eq += " = " + res
        return eq

def show(*args):

    for x in args:

        if G.latex:
            if hasattr(x, "latex_repr"):
                display(Math(x.latex_repr))
            else:
                display(x)
        else:
            if hasattr(x, "latex_repr"):
                print(latex_to_plain(x.latex_repr))
            else:
                print(x)