import re

import numpy as np

from . import constants as const

_SI_units = ["kg", "m", "s", "A", "K", "cd", "mol"]

_units = {
    "V": {"kg": 1, "m": 2, "s": -3, "A": -1, "K": 0, "cd": 0, "mol": 0},
    "C": {"kg": 0, "m": 0, "s": 1, "A": 1, "K": 0, "cd": 0, "mol": 0},
    "N": {"kg": 1, "m": 1, "s": -2, "A": 0, "K": 0, "cd": 0, "mol": 0},
    "J": {"kg": 1, "m": 2, "s": -2, "A": 0, "K": 0, "cd": 0, "mol": 0},
    "W": {"kg": 1, "m": 2, "s": -3, "A": 0, "K": 0, "cd": 0, "mol": 0},
    "Pa": {"kg": 1, "m": -1, "s": -2, "A": 0, "K": 0, "cd": 0, "mol": 0},
    "Ω": {"kg": 1, "m": 2, "s": -3, "A": -2, "K": 0, "cd": 0, "mol": 0},
}

_units_scale_conversion = {
    "eV": (const.qe, "V"),
    "bar": (1e5, "Pa"),
    "atm": (101325, "Pa"),
}

_unit_shift_conversion = {
    "°C": (273.15, "K"),
}

_unit_scale = {
    "y": -24,
    "z": -21,
    "a": -18,
    "f": -15,
    "p": -12,
    "n": -9,
    "u": -6,
    "m": -3,
    "c": -2,
    "d": -1,
    "h": 2,
    "k": 3,
    "M": 6,
    "G": 9,
    "T": 12,
    "P": 15,
    "E": 18,
    "Z": 21,
    "Y": 24,
    "": 0,
}

_unit_revert_scale = {}

_SI_units = ["kg", "m", "s", "A", "K", "cd", "mol"]

_units = {
    "V": {"kg": 1, "m": 2, "s": -3, "A": -1},
    "C": {"s": 1, "A": 1},
    "N": {"kg": 1, "m": 1, "s": -2},
    "J": {"kg": 1, "m": 2, "s": -2},
    "W": {"kg": 1, "m": 2, "s": -3},
    "Pa": {"kg": 1, "m": -1, "s": -2},
}

_unit_scale = {
    "y": -24,
    "z": -21,
    "a": -18,
    "f": -15,
    "p": -12,
    "n": -9,
    "u": -6,
    "m": -3,
    "c": -2,
    "d": -1,
    "h": 2,
    "k": 3,
    "M": 6,
    "G": 9,
    "T": 12,
    "P": 15,
    "E": 18,
    "Z": 21,
    "Y": 24,
    "": 0,
}


def par_parse(s):
    parsed = []
    count = 0
    opening = None
    closing = 0
    for i, x in enumerate(s):
        if x == "(":
            if opening is None:
                opening = i
            count += 1
        elif x == ")":
            count -= 1
            if count == 0 and opening is not None:
                parsed += [s[closing:opening], par_parse(s[opening + 1 : i])]
                closing = i + 1
                opening = None
    if closing < len(s):
        parsed.append(s[closing:])
    return parsed


def op_parse(s):
    r = []
    for x in s:
        if isinstance(x, list):
            r.append(op_parse(x))
        else:
            r += [x for x in re.split(r"(\*|/|\^)", x) if len(x) > 0]
    return r


def parse(unit):
    sp = par_parse(unit)
    sp = op_parse(sp)
    sp = u_parse(sp)
    sp = op_exec(sp)
    return sp


def num_parse(s):
    pass


def u_parse(s):
    if isinstance(s, list):
        sub = [u_parse(y) for y in s]
        return sub
    for x in "*/^":
        if x in s:
            return s
    result = None
    if re.match(r"\-?[0-9]+(\.[0-9]+)?", s):
        result = unit({}, float(s))
    elif s in _SI_units:
        result = unit(s)
    elif s in _units:
        result = unit(_units[s])
    elif s[0] in _unit_scale:
        if s[1:] in _SI_units:
            result = unit(s[1:], 10 ** (_unit_scale[s[0]]))
        elif s[1:] in _units:
            result = unit(_units[s[1:]], 10 ** (_unit_scale[s[0]]))
    elif len(s) == 2 and s[1] == "g" and x[0] in _unit_scale:
        result = unit("kg", 10 ** (_unit_scale[s[0]] - 3))
    elif s == "g":
        result = unit("kg", 1e-3)
    elif s in _units_scale_conversion:
        u = _units_scale_conversion[s]
        result = unit(u[1], u[0])
    return result


def op_exec(s):
    s = [op_exec(x) if isinstance(x, list) else x for x in s]
    while "^" in s:
        i = s.index("^")
        a = s[i - 1]
        b = s[i + 1]
        s = s[: i - 1] + [a**b] + s[i + 2 :]
    while "/" in s:
        i = s.index("/")
        s = s[: i - 1] + [s[i - 1] / s[i + 1]] + s[i + 2 :]
    while "*" in s:
        i = s.index("*")
        s = s[: i - 1] + [s[i - 1] * s[i + 1]] + s[i + 2 :]
    return s[0]


class unit:
    def __init__(self, u, value=1):
        self.value = 1
        if isinstance(u, str):
            if u in _SI_units:
                self.units = {u: 1}
            else:
                p = parse(u)
                self.units = p.units
                self.value = p.value
        elif isinstance(u, dict):
            self.units = {x: u[x] for x in u}
        else:
            raise TypeError(type(u))
        self.value *= value

    def __mul__(self, b):
        value = self.value
        units = self.units
        if isinstance(b, unit):
            for x in b.units:
                units[x] = units.get(x, 0) + b.units[x]
            value *= b.value
            return unit(units, value)
        return unit(self.units, self.value * b)

    __rmul__ = __mul__

    def __div__(self, b):
        value = self.value
        units = self.units
        if isinstance(b, unit):
            for x in b.units:
                units[x] = units.get(x, 0) - b.units[x]
            value /= b.value
            return unit(units, value)
        return unit(self.units, self.value / b)

    def __rdiv__(self, b):
        value = 1 / self.value
        units = {x: -self.units[x] for x in self.units}
        if isinstance(b, unit):
            for x in b.units:
                units[x] = units.get(x, 0) + b.units[x]
            value *= b.value
            return unit(units, value)
        return unit(units, b / self.value)

    def __pow__(self, n):
        if isinstance(n, unit):
            assert n.units == {}
            return unit(
                {x: n.value * self.units[x] for x in self.units}, self.value**n.value
            )
        return unit({x: n * x for x in self.units}, self.value**n)

    __truediv__ = __div__
    __rtruediv__ = __rdiv__

    def __repr__(self):
        if self.units == {}:
            return str(self.value)
        u = "*".join(
            [
                f"{x}^{self.units[x]}" if self.units[x] != 1 else x
                for x in self.units
                if self.units[x] != 0
            ]
        )
        if self.value == 1:
            return u
        return f"{self.value:.3e}*" + u

    __str__ = __repr__


class SIunit(np.ndarray):
    def __new__(cls, input_array, u=None):
        if u is None:
            u = {}
        obj = np.asarray(input_array).view(cls)
        obj.unit = unit(u)
        # Finally, we must return the newly created object:
        return obj

    def __array_finalize__(self, obj):
        # see InfoArray.__array_finalize__ for comments
        if obj is None:
            return
        self.unit = getattr(obj, "unit", {})

    def __repr__(self):
        return np.ndarray.__repr__(self)[:-1] + ", unit=" + str(self.unit) + ")"

    def __mul__(self, b):
        r = np.ndarray.__mul__(self, b)
        if isinstance(b, SIunit):
            r.unit = r.unit * b.unit
        else:
            r.unit = self.unit
        return r

    def __add__(self, b):
        r = np.ndarray.__add__(self, b)
        if isinstance(b, SIunit):
            for x in self.units:
                assert r.units[x] == b.units[x]
            for x in b.units:
                assert r.units[x] == b.units[x]
            return r
