from __future__ import annotations

import ast
import re
from itertools import count
from io import StringIO
import numpy as np


class Parameters(dict):
    parameter_counter = count()
    # is one global namespace a good idea?
    namespace: dict = {}

    def __init__(self):
        super().__init__()
        self._mapping: dict = {}

    def __str__(self) -> str:
        return 'Parameters:\n' + '\n'.join([f'{k}: {v}' for k, v in self.items()])

    def __getitem__(self, item) -> Parameter:
        if item in self._mapping:
            return super().__getitem__(self._mapping[item])
        else:
            return super().__getitem__(item)

    def __setitem__(self, key, value):
        self.add_parameter(key, value)

    def __contains__(self, item):
        for v in self.values():
            if item == v.name:
                return True

        return False

    def add(self,
            name: str,
            value: str | float | int = None,
            *,
            var: bool = True,
            lb: str | float = -np.inf,
            ub: str | float = np.inf) -> Parameter:

        par = Parameter(name=name, value=value, var=var, lb=lb, ub=ub)
        key = f'_p{next(Parameters.parameter_counter)}'

        self.add_parameter(key, par)

        return par

    def add_parameter(self, key: str, parameter: Parameter):
        self._mapping[parameter.name] = key
        super().__setitem__(key, parameter)

        parameter.eval_allowed = False
        self.namespace[key] = parameter.value
        parameter.namespace = self.namespace
        parameter.eval_allowed = True

        self.update_namespace()

    def replace_parameter(self, key_out: str, key_in: str, parameter: Parameter):
        self.add_parameter(key_in, parameter)
        for k, v in self._mapping.items():
            if v == key_out:
                self._mapping[k] = key_in
                break

        if key_out in self.namespace:
            del self.namespace[key_out]

    def fix(self):
        for v in self.keys():
            v._value = v.value
            v.namespace = {}

    @staticmethod
    def reset():
        Parameters.namespace = {}

    def get_key(self, name: str) -> str | None:
        for k, v in self.items():
            if name == v.name:
                return k

        return

    def get_state(self):
        return {k: v.get_state() for k, v in self.items()}

    def update_namespace(self):
        for p in self.values():
            try:
                p.value
            except NameError:
                expression = p._expr_disp
                for n, k in self._mapping.items():
                    expression, num_replaced = re.subn(re.escape(n), k, expression)

                p._expr = expression

    def prepare_bounds(self):
        original_values = list(self.values())
        for param in original_values:
            already_with_expression = False
            for mode, value in (('lower', param.lb), ('upper', param.ub)):
                if already_with_expression:
                    raise ValueError('Only one boundary can be an expression')
                already_with_expression = self.parse(param, value, bnd=mode)

    def parse(self, parameter, expression: str, bnd: str = 'lower') -> bool:
        try:
            float(expression)
            return False

        except ValueError:
            pp = ast.parse(expression)
            expr_syms = pp.body[0].value

            if isinstance(expr_syms, ast.BinOp):
                left, op, right = expr_syms.left, expr_syms.op, expr_syms.right

                # check for sign in numerator
                sign = 1
                if isinstance(left, ast.UnaryOp):
                    if isinstance(left.op, (ast.Not, ast.Invert)):
                        raise ValueError('Only `+` and `-` are supported for signs')
                    if isinstance(left.op, ast.USub):
                        sign = -1
                    left = left.operand

                # is expression number / parameter?
                if not (isinstance(left, ast.Constant) and isinstance(op, ast.Div) and isinstance(right, ast.Name)):
                    raise ValueError('Only simple division `const/parameter` are possible')

                result = self.make_proxy_div(parameter, left.value*sign, right.id, bnd=bnd)

                return result

            else:
                raise ValueError('I cannot work under these conditions')

    def make_proxy_div(self, param, num: str | float, denom: str | float, bnd: str = 'upper') -> bool:
        other_param = self[denom]
        other_lb, other_ub = other_param.lb, other_param.ub

        proxy = {'name': f'{param.name}*{other_param.name}', 'value': other_param.value*param.value, 'lb': None, 'ub': None}

        # switching signs is bad for inequalities
        if other_lb < 0 < other_ub:
            raise ValueError('Parameter in expression changes sign')

        if bnd == 'upper':
            # this whole schlamassel is only working for some values as lower bound
            if param.lb not in [-np.inf, 0]:
                raise ValueError('Invalid lower bounds')

            if other_ub < 0 or num < 0:
                # upper bound is always negative, switch boundaries
                proxy['lb'] = num
                proxy['ub'] = param.lb if param.lb == 0 else np.inf
            else:
                # upper bound is always positive
                proxy['lb'] = param.lb
                proxy['ub'] = num

        elif bnd == 'lower':
            if param.ub not in [np.inf, 0]:
                raise ValueError('Invalid upper bound')

            if other_ub <= 0 or num < 0:
                proxy['lb'] = param.lb if param.lb == 0 else np.inf
                proxy['ub'] = num
            else:
                proxy['lb'] = num
                proxy['ub'] = param.ub

        else:
            raise ValueError(f'unknown bound {bnd!r}, use `upper`or `lower`')

        param.set_expression(f'{num}/{denom}')
        self.add(**proxy)
        self.update_namespace()

        return True


class Parameter:
    """
    Container for one parameter
    """

    # TODO Parameter should know its own key
    def __init__(self,
                 name: str,
                 value: float | str,
                 var: bool = True,
                 lb: str | float = -np.inf,
                 ub: str | float = np.inf,
                 ):
        self._value: float | None = None
        self.var: bool = bool(var) if var is not None else True
        self.error:  None | float = None if self.var is False else 0.0
        self.name: str = name
        self.function: str = ""

        self.lb: str | None | float = lb if lb is not None else -np.inf
        self.ub: str | float | None = ub if ub is not None else np.inf
        self.namespace: dict = {}
        self.eval_allowed: bool = True
        self._expr: None | str = None
        self._expr_disp: None | str = None
        self.is_global = False

        if isinstance(value, str):
            self._expr_disp = value
            self._expr = value
            self.var = False
        else:
            if isinstance(self.lb, (int, float)) and isinstance(self.ub, (int, float)):
                if (self.lb <= value <= self.ub) or (not self.var):
                    self._value = value
                else:
                    raise ValueError('Value of parameter is outside bounds')
            else:
                self._value = value

            self.init_val = value

            with np.errstate(divide='ignore'):
                # throws RuntimeWarning for zeros
                self.scale = 10**(np.floor(np.log10(np.abs(self.value))))

            if self.scale == 0:
                self.scale = 1.

    def __str__(self) -> str:
        start = StringIO()
        if self.name:
            if self.function:
                start.write(f"{self.name} ({self.function})")
            else:
                start.write(self.name)

        if self.is_global:
            start.write("*")

        start.write(":  ")

        if self.var:
            start.write(f"{self.value:.4g} +/- {self.error:.4g}, init={self.init_val}")
        else:
            start.write(f"{self.value:.4g}")
            if self._expr is None:
                start.write(" (fixed)")
            else:
                start.write(f" (calc: {self._expr_disp})")

        return start.getvalue()

    def __add__(self, other: Parameter | float | int) -> float:
        if isinstance(other, (float, int)):
            return self.value + other
        elif isinstance(other, Parameter):
            return self.value + other.value

    def __radd__(self, other: Parameter | float) -> float:
        return self.__add__(other)

    @property
    def scaled_value(self) -> float:
        return self.value / self.scale

    @scaled_value.setter
    def scaled_value(self, value: float) -> None:
        self._value = value * self.scale

    @property
    def value(self) -> float | None:
        if self._value is not None:
            return self._value

        if self._expr is not None and self.eval_allowed:
            return eval(self._expr, {}, self.namespace)

        return

    def set_expression(self, expr: str):
        self._value = None
        self._expr_disp = expr
        self._expr = expr
        self.var = False

    @property
    def scaled_error(self) -> None | float:
        if self.error is not None:
            return self.error / self.scale

        return

    @scaled_error.setter
    def scaled_error(self, value) -> None:
        self.error = value * self.scale

    def get_state(self) -> dict:
        state_dict = {
            'name': self.name,
            'value': self._value,
            'error': self.error,
            'init_val': self.init_val,
            'var': self.var,
            'lb': self.lb,
            'ub': self.ub,
            'scale': self.scale,
            'function': self.function,
            '_expr': self._expr,
            '_expr_disp': self._expr_disp,
            'is_global': self.is_global,
         }
        return state_dict

    @staticmethod
    def set_state(state: dict) -> Parameter:
        par = Parameter(state.get('name'), state.pop('value'))
        for k, v in state.items():
            setattr(par, k, v)

        return par

    @property
    def full_name(self) -> str:
        name = self.name
        if self.function:
            name += f" ({self.function})"

        return name

    def copy(self) -> Parameter:
        if self._expr:
            val = self._expr_disp
        else:
            val = self._value
        para_copy = Parameter(name=self.name, value=val, var=self.var, lb=self.lb, ub=self.ub)
        para_copy._expr = self._expr
        para_copy.namespace = self.namespace
        para_copy.is_global = self.is_global
        para_copy.error = self.error
        para_copy.function = self.function

        return para_copy

    def fix(self):
        self._value = self.value
        self.namespace = {}

