"""
Market State Engine v1 — 因果可计算（无未来函数）

仅使用截至当前 8h K 线已收盘信息：
  EMA50/200、ADX、EMA slope、价格相对 MA200 距离

输出 0–100 分数 + 主导状态标签（argmax），供 Decision Gate 使用。
禁止用事后涨跌路径标注 cycle。
"""
from __future__ import annotations

import numpy as np
import pandas as pd
import talib.abstract as ta


def _clip01(x: pd.Series) -> pd.Series:
	return x.clip(lower=0.0, upper=1.0)


def compute_market_state_8h(df: pd.DataFrame) -> pd.DataFrame:
	"""
	在原生 8h OHLCV 上计算状态分数。
	返回列: accumulation_score, markup_score, distribution_score,
	         markdown_score, range_score, market_state, allow_spring, allow_utad
	"""
	out = df.copy()
	out["ema50"] = ta.EMA(out, timeperiod=50)
	out["ema200"] = ta.EMA(out, timeperiod=200)
	out["adx"] = ta.ADX(out, timeperiod=14)

	# slope: 过去 6 根 8h（约 2 天），仅用历史
	out["ema_slope"] = (out["ema50"] - out["ema50"].shift(6)) / out["ema50"].shift(6).replace(0, np.nan)
	out["dist_ema200"] = (out["close"] - out["ema200"]) / out["ema200"].replace(0, np.nan)

	bull = (out["close"] > out["ema200"]) & (out["ema50"] > out["ema200"])
	bear = (out["close"] < out["ema200"]) & (out["ema50"] < out["ema200"])
	range_m = (~bull) & (~bear)

	slope = out["ema_slope"].fillna(0.0)
	dist = out["dist_ema200"].fillna(0.0)
	adx = out["adx"].fillna(0.0)

	# ---- 分数：连续、因果、可解释 ----
	# accumulation: 仍处熊偏结构，但下跌斜率缓和 / 略抬升（吸筹语境）
	accum = (
		0.45 * bear.astype(float)
		+ 0.35 * _clip01((slope + 0.02) / 0.04)  # slope 从 -2%→+2% 映射
		+ 0.20 * _clip01((0.05 + dist) / 0.10)  # 仍在 MA200 下方但不极端深
	) * 100.0

	# markup: 牛偏 + 正斜率 + 价格在 MA200 上方
	markup = (
		0.40 * bull.astype(float)
		+ 0.35 * _clip01(slope / 0.02)
		+ 0.25 * _clip01(dist / 0.08)
	) * 100.0

	# distribution: 牛偏但斜率走平/向下（顶部语境）
	distrib = (
		0.40 * bull.astype(float)
		+ 0.40 * _clip01((-slope) / 0.015)
		+ 0.20 * _clip01((0.12 - dist.abs()) / 0.12)
	) * 100.0

	# markdown: 熊偏 + 明显负斜率
	markdown = (
		0.45 * bear.astype(float)
		+ 0.40 * _clip01((-slope) / 0.02)
		+ 0.15 * _clip01((-dist) / 0.10)
	) * 100.0

	# range: 非明确牛熊，或 ADX 偏低
	range_s = (
		0.50 * range_m.astype(float)
		+ 0.30 * _clip01((22.0 - adx) / 22.0)
		+ 0.20 * (1.0 - bull.astype(float)) * (1.0 - bear.astype(float))
	) * 100.0

	out["accumulation_score"] = accum.clip(0, 100)
	out["markup_score"] = markup.clip(0, 100)
	out["distribution_score"] = distrib.clip(0, 100)
	out["markdown_score"] = markdown.clip(0, 100)
	out["range_score"] = range_s.clip(0, 100)

	# 主导状态：与归因研究同一套因果规则（非事后路径标注）
	# bear+非急跌斜率 → accumulation；bull+正斜率 → markup；…
	state = np.full(len(out), "range", dtype=object)
	state[(bear) & (slope < -0.01)] = "markdown"
	state[(bear) & (slope >= -0.01)] = "accumulation"
	state[(bull) & (slope > 0.005)] = "markup"
	state[(bull) & (slope <= 0.005)] = "distribution"
	out["market_state"] = state

	# 默认 Gate v1.1：状态集合（soft 阈值由 apply_decision_gate 覆盖）
	out = apply_decision_gate(out, mode="state_set")
	return out


def apply_decision_gate(
	df: pd.DataFrame,
	*,
	mode: str = "state_set",
	q_sum: float = 100.0,
	q_bad: float = 55.0,
) -> pd.DataFrame:
	"""
	Decision Gate（因果）。

	mode:
	  - state_set: state ∈ {accumulation, markup} / UTAD 镜像
	  - soft_sum:  state_set 且 (accum+markup) >= q_sum
	  - soft_bad_cap: state_set 且 max(distrib, range, markdown) <= q_bad
	"""
	out = df.copy()
	state = out["market_state"]
	spring_state = state.isin(["accumulation", "markup"])
	utad_state = state.isin(["distribution", "markdown"])

	good_sum = out["accumulation_score"] + out["markup_score"]
	bad_max = out[["distribution_score", "range_score", "markdown_score"]].max(axis=1)
	# UTAD 镜像：good = distrib+markdown；bad = accum/range
	utad_good_sum = out["distribution_score"] + out["markdown_score"]
	utad_bad_max = out[["accumulation_score", "range_score", "markup_score"]].max(axis=1)

	if mode == "state_set":
		out["allow_spring"] = spring_state
		out["allow_utad"] = utad_state
	elif mode == "soft_sum":
		out["allow_spring"] = spring_state & (good_sum >= float(q_sum))
		out["allow_utad"] = utad_state & (utad_good_sum >= float(q_sum))
	elif mode == "soft_bad_cap":
		out["allow_spring"] = spring_state & (bad_max <= float(q_bad))
		out["allow_utad"] = utad_state & (utad_bad_max <= float(q_bad))
	else:
		raise ValueError(f"unknown gate mode: {mode}")

	out["gate_mode"] = mode
	out["gate_q_sum"] = float(q_sum)
	out["gate_q_bad"] = float(q_bad)
	return out


def spring_gate_mask(dataframe: pd.DataFrame, suffix: str = "_8h") -> pd.Series:
	col = f"allow_spring{suffix}"
	if col not in dataframe.columns:
		return pd.Series(True, index=dataframe.index)
	return dataframe[col].fillna(False).astype(bool)


def utad_gate_mask(dataframe: pd.DataFrame, suffix: str = "_8h") -> pd.Series:
	col = f"allow_utad{suffix}"
	if col not in dataframe.columns:
		return pd.Series(True, index=dataframe.index)
	return dataframe[col].fillna(False).astype(bool)