# %%
# 公式规范化
import re
from math import log, log10, e

MARK = '`'


class Strip:
    '''
    除去空格
    '''

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        s = s.replace(' ', '')
        return s, coefficient, input_symbol_set


class StripZh:
    '''
    去除中文
    '''

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        s = ''.join(list(filter(lambda c: not('\u4e00' <= c <= '\u9fa5'), s)))
        return s, coefficient, input_symbol_set


class CvtTwoAlphaToOne:
    '''
    多字母符号替换为1个
    '''

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        exclude_set = ['log', 'ln']
        if (tmp := re.findall(r'[a-zA-Z]{2,}', s)):  # 去除2个字母
            for x in tmp:
                if x not in exclude_set:
                    s = s.replace(x, x[0])  # 替换为1个字母
        return s, coefficient, input_symbol_set


class ReplaceAndGetCoefficient:
    '''
    变量字母替换系数值（整数除外）
    '''

    def __init__(self, mark=MARK):
        self.mark = mark

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        # pattern = r'-?\d+\.\d*|-?\d*\.\d+'
        pattern = r'-?(?:(?:\d+\.?\d*)|(?:\.?\d+))(?:[eE][-+]?\d+)?'
        if (tmp := re.findall(pattern, s)):  # 获取系数(仅小数)
            coe_n = ord('a')
            for x in tmp:
                if self._skip_coe(x):
                    continue
                # assert len(re.findall(x, s)) == 1, '系数多于1个'   # 校验系数数量
                coefficient[chr(coe_n)] = float(x)
                if x[0] == '-':
                    s = s.replace(
                        x, '+' + f'{self.mark}{chr(coe_n)}{self.mark}', 1)    # 系数替换为字母
                else:
                    s = s.replace(x, f'{self.mark}{chr(coe_n)}{self.mark}', 1)
                coe_n += 1
            s = s.replace('=+', '=')
        return s, coefficient, input_symbol_set

    def _skip_coe(self, x):
        return '.' not in x and 'e' not in x and 'E' not in x  # 整数不视为系数


class GetInputAlpha:
    '''
    获取系数
    '''

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        s_tmp = s.split('=')[-1]
        exclude_set = ['log', 'e']
        if (tmp := re.findall(r'[A-Z]{1}', s_tmp)):  # 获取输入参数
            input_symbol_set = set(filter(lambda x: x not in exclude_set, tmp))
        return s, coefficient, input_symbol_set


class GetAndReplaceInputAlpha:
    '''
    获取输入
    '''

    def __init__(self, mark=MARK):
        self.mark = mark

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        for x in input_symbol_set:
            s = s.replace(x, f'{self.mark}{x}{self.mark}')
        return s, coefficient, input_symbol_set


class ConvWithLog:
    '''
    公式左边带log，需要转换
    '''
    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        op = {'log': '10', 'ln': 'e'}
        for k, v in op.items():
            if s.startswith(k):
                s = f'{s.split("=")[0]}={v}^({s.split("=")[-1]})'
                break
        return s, coefficient, input_symbol_set


class ReplacePyOperator:
    '''
    数学计算符替换为python计算符
    '''

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        py_op_pre = {'log10': 'log'}    # 统一操作
        for k, v in py_op_pre.items():
            s = s.replace(k, v)
        py_op = {'^': '**', 'log': 'log10', 'ln': 'log'}    # 地理资源院中模型，默认log底数为10
        for k, v in py_op.items():
            s = s.replace(k, v)
        return s, coefficient, input_symbol_set


class GetValExpress:
    '''
    获取表达式
    '''

    def __call__(self, s: str, coefficient: dict, input_symbol_set: set):
        s = s.split('=')[-1]
        return s, coefficient, input_symbol_set


class Formula:
    '''
    公式处理
    '''
    class FeedDict:
        '''
        用值填充变量字母
        '''

        def __init__(self, mark=MARK) -> None:
            self.mark = mark

        def __call__(self, s: str, input_val: dict) -> str:
            for k, v in input_val.items():
                s = s.replace(f'{self.mark}{k}{self.mark}', str(v))
            for k, v in {'+-': '-'}.items():    # fix
                s = s.replace(k, v)
            return s

    def __init__(self, s_formula: str):
        self.s, self.coefficient, self.input_symbol_set = self.standard_formula(
            s_formula)
        self.express = self.feed_dict(self.s, self.coefficient)

    def standard_formula(self, s: str):
        coefficient, input_symbol_set = {}, set()
        for handler in [
            Strip(),
            StripZh(),
            CvtTwoAlphaToOne(),
            ReplaceAndGetCoefficient(),
            GetInputAlpha(),
            GetAndReplaceInputAlpha(),
            ConvWithLog(),
            ReplacePyOperator(),
            GetValExpress(),
        ]:
            s, coefficient, input_symbol_set = handler(
                s, coefficient, input_symbol_set)
        return s, coefficient, input_symbol_set

    # def feed_dict(self, s, input_val: dict):
    #     for k, v in input_val.items():
    #         s = s.replace(k, str(v))
    #     for k, v in {'+-': '-'}.items():    # fix
    #         s = s.replace(k, v)
    #     return s

    def feed_dict(self, s: str, input_val: dict):
        return self.FeedDict()(s, input_val)

    def forward(self, x: dict):
        express = self.feed_dict(self.express, x)
        return eval(express)

    def __call__(self, x: dict):
        return self.forward(x)

    def eval(self, x: dict, ndigits=5):
        return round(self.forward(x), ndigits)

    def __str__(self) -> str:
        return self._print_color(f'out = {self.s} = {self.express}')

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

    def _print_color(self, s: str, mark=MARK) -> str:
        pattern = r'`[A-Za-z]+`'
        for x in re.findall(pattern, s):
            s = s.replace(x, f'\033[0;31;40m{x.replace(mark, "")}\033[0m')
        return s


if __name__ == '__main__':
    for s in [
        'Wr = -3E-05*((D^2)*H) - 0.1704',
        'log(TWs)=-1.293+0.827*log((D^2)*H)',
        'log(TWs)=-1.113+0.882*log10((D^2)*H)',
        'y = 7.9753*e^(0.00021*(D^2)*H)',
    ]:
        formula = Formula(s)
        # print(formula({'D':1.1, 'H': 2.2}))
        print(formula, '=', formula.eval({'D': 1.1, 'H': 2.2}, ndigits=3))
