"""GTJA 公式预处理器。

将 GTJA 原始语法清洗为框架可识别的 DSL 语法。
修复了原版公式中的拼写错误、歧义重载、嵌套三元运算符等问题。
"""

import re

def clean_gtja_formula(formula: str) -> str:
    """将 GTJA 原始语法清洗为框架可识别的 DSL 语法。"""

    formula = formula.strip()

    # 0. 清洗中文标点符号和空格
    formula = formula.replace("（", "(").replace("）", ")")
    formula = formula.replace("，", ",").replace("–", "-")
    formula = formula.replace("【", "[").replace("】", "]")

    # 1. 替换基础算术运算符和逻辑运算符
    formula = formula.replace("./", "/").replace(".*", "*").replace("^", "**")
    formula = formula.replace("||", "|").replace("&&", "&")

    # 2. 宏替换 (基础衍生宏)
    replacements = {
        r"\bRET\b": "(CLOSE / DELAY(CLOSE, 1) - 1)",
        r"\bVWAP\b": "(AMOUNT / VOLUME)",
    }
    for old, new in replacements.items():
        formula = re.sub(old, new, formula, flags=re.IGNORECASE)

    # 3. 修复原版 GTJA 公式库中的各处天坑笔误 (Typo)
    typo_mapping = {
        r"\bHGIH\b": "HIGH",                                         # Alpha 159 拼写错误
        r"\bDELAT\b": "DELTA",                                       # Alpha 131 拼写错误
        r"\?STD\(CLOSE\s*:\s*20\)\s*,\s*0": "? STD(CLOSE, 20) : 0",  # Alpha 23 冒号与逗号打反
        r"CLOSE\s*:\s*20": "CLOSE, 20",                              # 其他可能存在的冒号误触
        r"(?<=-)L\b": "LOW",                                         # Alpha 52 极简缩写: -L)
        r"\)\(": ")*(",                                              # Alpha 166 缺乘号: (20-2)(SUM...
        r"\(CLOSE/DELAY\(CLOSE,1\),20\)": "(CLOSE/DELAY(CLOSE,1)-1)",# Alpha 166 多余参数与格式错乱
        r"\*SIGN\(DELTA\(CLOSE,\s*7\)\)\s*:\s*\(-1\s*\*VOLUME\)\)\)": "*SIGN(DELTA(CLOSE, 7))) : (-1 * VOLUME))", # Alpha 180 括号位置打错
        r"\bOR\b": "|",                                              # Alpha 182 异常逻辑符
        r"\bAND\b": "&",                                             # 异常逻辑符
    }
    for bad, good in typo_mapping.items():
        formula = re.sub(bad, good, formula, flags=re.IGNORECASE)

    # 4. 修复条件表达式中的赋值符为比较符
    # 把 = 变成 ==，但避开 <=, >=, !=, ==
    formula = re.sub(r"(?<![=<>!])=(?![=])", "==", formula)

    # 5. 智能解析多态重载函数 (RANK, MEAN) -> (cs_/ts_)
    def resolve_overloaded_funcs(f: str) -> str:
        for target in ["RANK", "MEAN"]:
            while True:
                match = re.search(rf"(?<!cs_)(?<!ts_)\b{target}\s*\(", f, flags=re.IGNORECASE)
                if not match:
                    break

                start_idx = match.start()
                paren_start = match.end() - 1

                depth = 1
                comma_count = 0
                paren_end = -1
                for i in range(paren_start + 1, len(f)):
                    if f[i] == '(': depth += 1
                    elif f[i] == ')':
                        depth -= 1
                        if depth == 0:
                            paren_end = i
                            break
                    elif f[i] == ',' and depth == 1:
                        comma_count += 1

                if paren_end == -1:
                    break

                args_str = f[paren_start+1:paren_end].strip()
                arg_count = comma_count + 1 if args_str else 0

                prefix = "ts_" if arg_count > 1 else "cs_"
                f = f[:start_idx] + prefix + target.lower() + f[paren_start:]
        return f

    formula = resolve_overloaded_funcs(formula)

    # 6. 三元运算符安全转换 (Condition) ? True : False -> if_(Condition, True, False)
    # 自右向左匹配，完美解决复杂嵌套
    def ternary_to_if(f: str) -> str:
        max_iterations = 100
        iteration = 0

        while "?" in f and iteration < max_iterations:
            q_idx = f.rfind("?")
            if q_idx == -1: break

            depth = 0
            c_idx = -1
            for i in range(q_idx + 1, len(f)):
                if f[i] == '(': depth += 1
                elif f[i] == ')': depth -= 1
                elif f[i] == ':' and depth == 0:
                    c_idx = i
                    break

            if c_idx == -1:
                f = f[:q_idx] + "_" + f[q_idx+1:]
                continue

            depth = 0
            a_start = 0
            for i in range(q_idx - 1, -1, -1):
                if f[i] == ')': depth += 1
                elif f[i] == '(':
                    depth -= 1
                    if depth < 0:
                        a_start = i + 1
                        break
                elif f[i] == ',' and depth == 0:
                    a_start = i + 1
                    break

            depth = 0
            c_end = len(f)
            for i in range(c_idx + 1, len(f)):
                if f[i] == '(': depth += 1
                elif f[i] == ')':
                    depth -= 1
                    if depth < 0:
                        c_end = i
                        break
                elif f[i] == ',' and depth == 0:
                    c_end = i
                    break

            A_str = f[a_start:q_idx].strip()
            B_str = f[q_idx+1:c_idx].strip()
            C_str = f[c_idx+1:c_end].strip()

            replacement = f"if_({A_str}, {B_str}, {C_str})"
            f = f[:a_start] + replacement + f[c_end:]
            iteration += 1

        return f

    formula = ternary_to_if(formula)

    # 7. 函数名直接映射 (GTJA -> DSL)
    function_mapping = {
        r"\bDELAY\s*\(": "ts_delay(",
        r"\bDELTA\s*\(": "ts_delta(",
        r"\bSTD\s*\(": "ts_std(",
        r"\bMAX\s*\(": "max_(",
        r"\bMIN\s*\(": "min_(",
        r"\bSUM\s*\(": "ts_sum(",
        r"\bVAR\s*\(": "ts_var(",
        r"\bCOV\s*\(": "ts_cov(",
        r"\bCOVIANCE\s*\(": "ts_cov(",
        r"\bCORR\s*\(": "ts_corr(",
        r"\bSMA\s*\(": "ts_sma(",
        r"\bSMEAN\s*\(": "ts_sma(",
        r"\bMA\s*\(": "ts_mean(",
        r"\bWMA\s*\(": "ts_wma(",
        r"\bDECAYLINEAR\s*\(": "ts_decay_linear(",
        r"\bHIGHDAY\s*\(": "ts_argmax(",
        r"\bLOWDAY\s*\(": "ts_argmin(",
        r"\bCOUNT\s*\(": "ts_count(",
        r"\bPROD\s*\(": "ts_prod(",
        r"\bSUMAC\s*\(": "ts_sumac(",
        r"\bTSRANK\s*\(": "ts_rank(",
        r"\bTSMAX\s*\(": "ts_max(",
        r"\bTSMIN\s*\(": "ts_min(",
        r"\bLOG\s*\(": "log(",
        r"\bEXP\s*\(": "exp(",
        r"\bSQRT\s*\(": "sqrt(",
        r"\bSIGN\s*\(": "sign(",
        r"\bABS\s*\(": "abs(",
        r"\bATAN\s*\(": "atan(",
    }
    for gtja_func, dsl_func in function_mapping.items():
        formula = re.sub(gtja_func, dsl_func, formula, flags=re.IGNORECASE)

    # 8. 字段名映射
    field_mapping = {
        r"\bCLOSE\b": "close",
        r"\bOPEN\b": "open",
        r"\bHIGH\b": "high",
        r"\bLOW\b": "low",
        r"\bVOLUME\b": "vol",
        r"\bVOL\b": "vol",
        r"\bAMOUNT\b": "amount",
        r"\bPRE_CLOSE\b": "pre_close",
        r"\bCHANGE\b": "change",
        r"\bPCT_CHG\b": "pct_chg",
    }
    for gtja_field, dsl_field in field_mapping.items():
        formula = re.sub(gtja_field, dsl_field, formula, flags=re.IGNORECASE)

    # 9. 智能补全默认缺省参数
    def add_default_args(f: str, func_name: str, default_val: str, required_args: int) -> str:
        pattern = f"{func_name}("
        result =[]
        i = 0
        while i < len(f):
            if f[i:i+len(pattern)] == pattern:
                paren_start = i + len(pattern) - 1

                depth = 1
                paren_end = -1
                for j in range(paren_start + 1, len(f)):
                    if f[j] == '(': depth += 1
                    elif f[j] == ')':
                        depth -= 1
                        if depth == 0:
                            paren_end = j
                            break

                if paren_end == -1:
                    result.append(f[i])
                    i += 1
                    continue

                args_content = f[paren_start+1:paren_end]

                # 正确统计顶层逗号数量 (修复嵌套逗号被计入的 Bug)
                depth_comma = 0
                comma_count = 0
                for ch in args_content:
                    if ch == '(': depth_comma += 1
                    elif ch == ')': depth_comma -= 1
                    elif ch == ',' and depth_comma == 0:
                        comma_count += 1

                arg_count = comma_count + 1 if args_content.strip() else 0

                if arg_count < required_args:
                    result.append(f"{func_name}({args_content}, {default_val})")
                else:
                    result.append(f[i:paren_end+1])
                i = paren_end + 1
            else:
                result.append(f[i])
                i += 1
        return "".join(result)

    formula = add_default_args(formula, "ts_delay", "1", 2)
    formula = add_default_args(formula, "ts_delta", "1", 2)
    formula = add_default_args(formula, "ts_std", "20", 2)
    formula = add_default_args(formula, "ts_corr", "5", 3)
    formula = add_default_args(formula, "ts_sma", "1", 3)

    return formula


def filter_unsupported_formulas(formula: str) -> bool:
    """检查公式是否包含不支持的函数/算子。"""
    unsupported_patterns =[
        r"\bREGBETA\b",       # OLS Beta
        r"\bREGRESI\b",       # OLS 残差
        r"\bSEQUENCE\b",      # 生成时间序列(作自变量)
        r"\bSELF\b",          # 循环递归引用
        r"\bBANCHMARK\w*\b",  # 基准指数(修正匹配 BANCHMARKINDEXCLOSE 等连写)
        r"\bINDEX\b",         # 宏观变量引入
        r"\bMKT\b", r"\bSMB\b", r"\bHML\b", # Fama-French 因子
        r"\bDTM\b", r"\bDBM\b", r"\bTR\b", r"\bHD\b", r"\bLD\b", # 复杂的外部黑盒宏
        r"\bFILTER\b",        # 条件屏蔽函数
        r"\bSUMIF\b",         # 条件求和函数
    ]
    for pattern in unsupported_patterns:
        if re.search(pattern, formula, re.IGNORECASE):
            return False
    return True

