"""V1.2 GPU feature engine.

This module intentionally leaves ``gpu_feature_engine_v1`` unchanged.  Only
Supertrend uses the V1.2 ordered, on-device RMA seed implementation.
"""
from __future__ import annotations

from typing import Any

import torch

from gpu_feature_engine_v1 import _empty, _true_range, evaluate_batch as _v1_evaluate_batch


def rma_ordered_seed(values: torch.Tensor, period: int) -> torch.Tensor:
    """RMA with the historical oracle's left-to-right seed association."""
    out = _empty(values)
    valid = torch.nonzero(~torch.isnan(values), as_tuple=False).flatten()
    if valid.numel() < period:
        return out
    seed_bar = int(valid[period - 1].item())
    total = torch.zeros((), dtype=values.dtype, device=values.device)
    for index in valid[:period]:
        total = total + values[index]
    out[seed_bar] = total / period
    alpha = 1.0 / period
    for index in range(seed_bar + 1, values.numel()):
        out[index] = out[index - 1] if torch.isnan(values[index]) else alpha * values[index] + (1.0 - alpha) * out[index - 1]
    return out


def supertrend_trace(close: torch.Tensor, high: torch.Tensor, low: torch.Tensor, period: int, multiplier: float) -> dict[str, torch.Tensor]:
    """Return complete V1.2 Supertrend state, including branch predicates."""
    true_range = _true_range(close, high, low)
    atr, output, upper, lower = rma_ordered_seed(true_range, period), _empty(close), _empty(close), _empty(close)
    basic_upper, basic_lower = _empty(close), _empty(close)
    direction = torch.zeros_like(close, dtype=torch.int8)
    initialized = torch.zeros_like(close, dtype=torch.bool)
    prior_close_le_upper = torch.zeros_like(close, dtype=torch.bool)
    prior_close_ge_lower = torch.zeros_like(close, dtype=torch.bool)
    active_long = torch.zeros_like(close, dtype=torch.bool)
    close_ge_lower = torch.zeros_like(close, dtype=torch.bool)
    close_le_upper = torch.zeros_like(close, dtype=torch.bool)
    output_uses_lower = torch.zeros_like(close, dtype=torch.bool)
    direction_transition = torch.zeros_like(close, dtype=torch.bool)
    for index in range(period, close.numel()):
        midpoint = (high[index] + low[index]) / 2.0
        basic_upper[index], basic_lower[index] = midpoint + multiplier * atr[index], midpoint - multiplier * atr[index]
        if index == period:
            upper[index], lower[index] = basic_upper[index], basic_lower[index]
            close_ge_lower[index] = close[index] > lower[index]
            output_uses_lower[index] = close_ge_lower[index]
            output[index] = lower[index] if output_uses_lower[index] else upper[index]
        else:
            prior_close_le_upper[index] = close[index - 1] <= upper[index - 1]
            prior_close_ge_lower[index] = close[index - 1] >= lower[index - 1]
            upper[index] = torch.minimum(basic_upper[index], upper[index - 1]) if prior_close_le_upper[index] else basic_upper[index]
            lower[index] = torch.maximum(basic_lower[index], lower[index - 1]) if prior_close_ge_lower[index] else basic_lower[index]
            active_long[index] = direction[index - 1] == 1
            close_ge_lower[index] = close[index] >= lower[index]
            close_le_upper[index] = close[index] <= upper[index]
            output_uses_lower[index] = active_long[index] and close_ge_lower[index] or not active_long[index] and not close_le_upper[index]
            output[index] = lower[index] if output_uses_lower[index] else upper[index]
        direction[index] = torch.where(close[index] > output[index], 1, -1)
        initialized[index] = True
        if index > period:
            direction_transition[index] = direction[index] != direction[index - 1]
    return {"output": output, "true_range": true_range, "atr": atr, "basic_upper": basic_upper, "basic_lower": basic_lower, "upper": upper, "lower": lower, "direction": direction, "initialized": initialized, "prior_close_le_upper": prior_close_le_upper, "prior_close_ge_lower": prior_close_ge_lower, "active_long": active_long, "close_ge_lower": close_ge_lower, "close_le_upper": close_le_upper, "output_uses_lower": output_uses_lower, "direction_transition": direction_transition}


def evaluate_batch(request: dict[str, Any], close: torch.Tensor, high: torch.Tensor, low: torch.Tensor, volume: torch.Tensor) -> dict[str, torch.Tensor]:
    """Evaluate Batch01, replacing only Supertrend requests with V1.2 logic."""
    other = {"requests": [item for item in request["requests"] if int(item["indicator_id"]) != 19]}
    outputs = _v1_evaluate_batch(other, close, high, low, volume) if other["requests"] else {}
    for item in request["requests"]:
        if int(item["indicator_id"]) == 19:
            outputs[str(item["request_id"])] = supertrend_trace(close, high, low, int(item["period"]), float(item["p1"]))["output"]
    return outputs