# indie:lang_version = 5
from math import nan, floor, sqrt
from indie import (
    indicator, param, plot, color, source,
    MutSeriesF, SeriesF, Color, Optional, MainContext, sec_context, algorithm
)
from indie.algorithms import Ma, Sma, Adx, Atr, Wma
from indie.math import cross_over, cross_under


# ─── Hull Moving Average ──────────────────────────────────────────────────────
@algorithm
def HullMa(self, src: SeriesF, length: int) -> SeriesF:
    half_len = length // 2
    sq_len = floor(sqrt(length))
    wma_half = Wma.new(src, half_len)
    wma_full = Wma.new(src, length)
    delta = MutSeriesF.new(2.0 * wma_half[0] - wma_full[0])
    return Wma.new(delta, sq_len)


# ─── Generic Trend MA (SMA / EMA / WMA / HMA) ────────────────────────────────
@algorithm
def TrendMa(self, src: SeriesF, length: int, ma_type: str) -> SeriesF:
    result: Optional[SeriesF]
    if ma_type == 'HMA':
        result = HullMa.new(src, length)
    else:
        result = Ma.new(src, length, ma_type)
    return result.value()


# ─── HTF Secondary Context (SMA 200 на старшем ТФ) ───────────────────────────
@sec_context
def HtfCtx(self):
    htf_sma = Sma.new(self.close, 200)
    return self.close[0], htf_sma[0]


# ─── Main ─────────────────────────────────────────────────────────────────────
@indicator('MACD + SMA 200 Enhanced', overlay_main_pane=True)
@param.source('src', default=source.CLOSE, title='Source')
# MACD
@param.int('fast_length', default=12, min=1, title='MACD Fast Length')
@param.int('slow_length', default=26, min=1, title='MACD Slow Length')
@param.int('signal_length', default=9, min=1, title='MACD Signal Length')
@param.str('macd_ma_type', default='SMA', title='MACD MA Type', options=['SMA', 'EMA'])
# Trend MA
@param.int('trend_length', default=200, min=2, title='Trend MA Length')
@param.str('trend_type', default='SMA', title='Trend MA Type', options=['SMA', 'EMA', 'WMA', 'HMA'])
# Visualization
@param.bool('show_bar_color', default=True, title='Enable Bar Color?')
@param.bool('show_mas', default=True, title='Enable Moving Averages?')
@param.bool('show_background', default=True, title='Enable Trend Background?')
# ADX Filter
@param.bool('use_adx', default=False, title='Enable ADX Filter?')
@param.int('adx_length', default=14, min=1, title='ADX Length')
@param.float('adx_threshold', default=25.0, min=1.0, title='ADX Threshold')
# ATR Filter
@param.bool('use_atr', default=False, title='Enable ATR Volatility Filter?')
@param.int('atr_length', default=14, min=1, title='ATR Length')
# MTF Filter
@param.bool('use_mtf', default=False, title='Enable MTF Filter?')
@param.time_frame('htf_tf', default='1D', title='HTF Timeframe')
# Plots
@plot.line('fast_ma', title='Fast MA')
@plot.line('slow_ma', title='Slow MA', line_width=2)
@plot.line('trend_ma', title='Trend MA', line_width=4)
@plot.fill('fast_ma', 'trend_ma', title='Trend Ribbon')
@plot.background(title='Trend Regime')
@plot.bar_color(title='Bar Color')
@plot.marker(style=plot.marker_style.CIRCLE, position=plot.marker_position.BELOW, size=7, title='Bull Signal',
    display_options=plot.MarkerDisplayOptions(pane=True, status_line=False, price_label=False))
@plot.marker(style=plot.marker_style.CIRCLE, position=plot.marker_position.ABOVE, size=7, title='Bear Signal',
    display_options=plot.MarkerDisplayOptions(pane=True, status_line=False, price_label=False))
class Main(MainContext):
    def __init__(self, htf_tf):
        self._htf_close, self._htf_sma = self.calc_on(HtfCtx, time_frame=htf_tf)

    def calc(self, src, fast_length, slow_length, signal_length, macd_ma_type,
             trend_length, trend_type,
             show_bar_color, show_mas, show_background,
             use_adx, adx_length, adx_threshold,
             use_atr, atr_length, use_mtf):

        # ── Moving averages ──────────────────────────────────────────────────
        fast_ma = Ma.new(src, fast_length, macd_ma_type)
        slow_ma = Ma.new(src, slow_length, macd_ma_type)
        trend_ma = TrendMa.new(src, trend_length, trend_type)

        # ── MACD (manual — reuses fast/slow MA, нет двойного подсчёта) ───────
        macd_series = MutSeriesF.new(fast_ma[0] - slow_ma[0])
        signal_series = Ma.new(macd_series, signal_length, macd_ma_type)
        hist_series = MutSeriesF.new(macd_series[0] - signal_series[0])

        # ── ADX ─────────────────────────────────────────────────────────────
        adx_minus_di, adx_series, adx_plus_di = Adx.new(adx_length, adx_length)

        # ── ATR ─────────────────────────────────────────────────────────────
        atr_series = Atr.new(atr_length)
        atr_ma_series = Sma.new(atr_series, atr_length)

        # ── Pre-declare Color vars (Indie block-scoping rule) ─────────────────
        line_color: Color = color.BLUE
        trend_ma_color: Color = color.GRAY
        bar_trend_color: Color = color.BLUE
        bg_color: Color = color.TRANSPARENT
        ribbon_color: Color = color.GRAY(0.2)
        bull_marker_color: Color = color.TRANSPARENT
        bear_marker_color: Color = color.TRANSPARENT

        # ── Trend state ──────────────────────────────────────────────────────
        trend_rising = trend_ma[0] > trend_ma[1]
        trend_falling = trend_ma[0] < trend_ma[1]
        price_above = self.close[0] > trend_ma[0]
        price_below = self.close[0] < trend_ma[0]
        macd_pos = macd_series[0] > 0.0
        macd_neg = macd_series[0] < 0.0

        # ── Filter gates ─────────────────────────────────────────────────────
        adx_ok = (not use_adx) or (adx_series[0] > adx_threshold)
        atr_ok = (not use_atr) or (atr_series[0] > atr_ma_series[0])
        mtf_bull_ok = (not use_mtf) or (self._htf_close[0] > self._htf_sma[0])
        mtf_bear_ok = (not use_mtf) or (self._htf_close[0] < self._htf_sma[0])

        # ── Trend MA color ───────────────────────────────────────────────────
        if trend_rising:
            trend_ma_color = color.GREEN
        else:
            trend_ma_color = color.RED

        # ── Fast/slow line color ─────────────────────────────────────────────
        if fast_ma[0] > slow_ma[0] and trend_rising and self.close[0] > slow_ma[0]:
            line_color = color.GREEN
        elif fast_ma[0] < slow_ma[0] and trend_falling and self.close[0] < slow_ma[0]:
            line_color = color.RED
        else:
            line_color = color.BLUE

        # ── Bar color ────────────────────────────────────────────────────────
        if (self.close[0] > fast_ma[0] and self.close[0] > slow_ma[0]
                and self.close[0] > trend_ma[0] and slow_ma[0] > slow_ma[1]):
            bar_trend_color = color.GREEN
        elif (self.close[0] < fast_ma[0] and self.close[0] < slow_ma[0]
                and self.close[0] < trend_ma[0] and slow_ma[0] < slow_ma[1]):
            bar_trend_color = color.RED
        else:
            bar_trend_color = color.BLUE

        # ── 5-level trend regime (background + ribbon color) ──────────────────
        if price_above and trend_rising and macd_pos and adx_ok:
            bg_color = color.GREEN(0.15)      # Strong Uptrend
            ribbon_color = color.GREEN(0.25)
        elif price_above and trend_rising:
            bg_color = color.GREEN(0.07)      # Weak Uptrend
            ribbon_color = color.GREEN(0.15)
        elif price_below and trend_falling and macd_neg and adx_ok:
            bg_color = color.RED(0.15)        # Strong Downtrend
            ribbon_color = color.RED(0.25)
        elif price_below and trend_falling:
            bg_color = color.RED(0.07)        # Weak Downtrend
            ribbon_color = color.RED(0.15)
        else:
            bg_color = color.TRANSPARENT      # Neutral
            ribbon_color = color.GRAY(0.2)

        if not show_background:
            bg_color = color.TRANSPARENT

        # ── Entry signals ────────────────────────────────────────────────────
        bull_signal = (
            cross_over(hist_series, 0.0)
            and macd_pos
            and fast_ma[0] > slow_ma[0]
            and price_above
            and adx_ok
            and atr_ok
            and mtf_bull_ok
        )
        bear_signal = (
            cross_under(hist_series, 0.0)
            and macd_neg
            and fast_ma[0] < slow_ma[0]
            and price_below
            and adx_ok
            and atr_ok
            and mtf_bear_ok
        )

        if bull_signal:
            bull_marker_color = color.GREEN
        if bear_signal:
            bear_marker_color = color.RED

        # ── MA visibility toggle ─────────────────────────────────────────────
        fast_ma_val = fast_ma[0] if show_mas else nan
        slow_ma_val = slow_ma[0] if show_mas else nan
        trend_ma_val = trend_ma[0] if show_mas else nan

        return (
            plot.Line(fast_ma_val, color=line_color),
            plot.Line(slow_ma_val, color=line_color),
            plot.Line(trend_ma_val, color=trend_ma_color),
            plot.Fill(color=ribbon_color),
            plot.Background(color=bg_color),
            plot.BarColor(bar_trend_color if show_bar_color else None),
            plot.Marker(value=self.low[0], color=bull_marker_color),
            plot.Marker(value=self.high[0], color=bear_marker_color),
        )
