"""Safe tensor utilities and GP operator AST for factor mining."""

from __future__ import annotations

import math
import random

import torch

WINDOWS = [3, 5, 8, 10, 15, 20, 30, 40, 60, 80, 120]
CLIP_ABS_VALUE = 1e6


def clean_tensor(x, clip=CLIP_ABS_VALUE):
    x = torch.nan_to_num(x, nan=0.0, posinf=clip, neginf=-clip)
    return torch.clamp(x, -clip, clip)


def safe_div(x, y, eps=1e-8):
    y_safe = torch.where(torch.abs(y) > eps, y, torch.ones_like(y))
    return clean_tensor(x / y_safe)


def safe_log(x):
    return torch.log(torch.clamp(torch.abs(x), min=1e-8))


def safe_sqrt(x):
    return torch.sqrt(torch.clamp(torch.abs(x), min=1e-8))


def ts_delay_raw(x, d):
    y = torch.roll(x, shifts=d, dims=1)
    y[:, :d] = float("nan")
    return y


def ts_future_raw(x, d):
    y = torch.roll(x, shifts=-d, dims=1)
    y[:, -d:] = float("nan")
    return y


def rolling_unfold(x, d):
    if x.shape[1] < d:
        return None
    return x.unfold(1, d, 1)


def rolling_pad(x, core, d):
    padding = torch.full((x.shape[0], d - 1), float("nan"), device=x.device)
    return torch.cat([padding, core], dim=1)


def nanmean_dim(u, dim, keepdim=False):
    valid = ~torch.isnan(u)
    safe = torch.where(valid, u, torch.zeros_like(u))
    count = valid.sum(dim=dim, keepdim=keepdim).float()
    return safe.sum(dim=dim, keepdim=keepdim) / torch.clamp(count, min=1.0)


def nanstd_dim(u, dim, keepdim=False):
    mean = nanmean_dim(u, dim=dim, keepdim=True)
    valid = ~torch.isnan(u)
    diff = torch.where(valid, u - mean, torch.zeros_like(u))
    count = valid.sum(dim=dim, keepdim=True).float()
    var = (diff ** 2).sum(dim=dim, keepdim=True) / torch.clamp(count - 1.0, min=1.0)
    std = torch.sqrt(var + 1e-8)
    return std if keepdim else std.squeeze(dim)


class Node:
    def evaluate(self, engine):
        raise NotImplementedError

    def clone(self):
        raise NotImplementedError

    def get_depth(self):
        raise NotImplementedError

    def get_nodes(self):
        raise NotImplementedError

    def get_size(self):
        return len(self.get_nodes())


class Terminal(Node):
    def __init__(self, feature_name):
        self.feature_name = feature_name

    def evaluate(self, engine):
        return engine.get_data(self.feature_name)

    def clone(self):
        return Terminal(self.feature_name)

    def get_depth(self):
        return 1

    def get_nodes(self):
        return [self]

    def __str__(self):
        return str(self.feature_name)


class Constant(Node):
    def __init__(self, value):
        self.value = float(value)

    def evaluate(self, engine):
        base = engine.get_data("收盘价")
        return torch.full_like(base, self.value)

    def clone(self):
        return Constant(self.value)

    def get_depth(self):
        return 1

    def get_nodes(self):
        return [self]

    def __str__(self):
        return f"{self.value:.4g}"


class Operator(Node):
    def __init__(self, name, *children):
        self.name = name
        self.children = list(children)

    def evaluate(self, engine):
        try:
            vals = [c.evaluate(engine) for c in self.children]
            return clean_tensor(self._compute(*vals))
        except Exception:
            return torch.zeros_like(engine.get_data("收盘价"))

    def _compute(self, *args):
        raise NotImplementedError

    def clone(self):
        return type(self)(*[c.clone() for c in self.children])

    def get_depth(self):
        return 1 + max(c.get_depth() for c in self.children)

    def get_nodes(self):
        nodes = [self]
        for c in self.children:
            nodes.extend(c.get_nodes())
        return nodes

    def __str__(self):
        args = ", ".join(str(c) for c in self.children)
        return f"{self.name}({args})"


class Add(Operator):
    def __init__(self, a, b):
        super().__init__("Add", a, b)

    def _compute(self, x, y):
        return x + y


class Sub(Operator):
    def __init__(self, a, b):
        super().__init__("Sub", a, b)

    def _compute(self, x, y):
        return x - y


class Mul(Operator):
    def __init__(self, a, b):
        super().__init__("Mul", a, b)

    def _compute(self, x, y):
        return x * y


class Div(Operator):
    def __init__(self, a, b):
        super().__init__("Div", a, b)

    def _compute(self, x, y):
        return safe_div(x, y)


class Max2(Operator):
    def __init__(self, a, b):
        super().__init__("Max", a, b)

    def _compute(self, x, y):
        return torch.maximum(x, y)


class Min2(Operator):
    def __init__(self, a, b):
        super().__init__("Min", a, b)

    def _compute(self, x, y):
        return torch.minimum(x, y)


class AbsOp(Operator):
    def __init__(self, a):
        super().__init__("Abs", a)

    def _compute(self, x):
        return torch.abs(x)


class Neg(Operator):
    def __init__(self, a):
        super().__init__("Neg", a)

    def _compute(self, x):
        return -x


class LogOp(Operator):
    def __init__(self, a):
        super().__init__("Log", a)

    def _compute(self, x):
        return safe_log(x)


class SqrtOp(Operator):
    def __init__(self, a):
        super().__init__("Sqrt", a)

    def _compute(self, x):
        return safe_sqrt(x)


class SignedPower(Operator):
    def __init__(self, a, power=2.0):
        super().__init__(f"SignedPower_{power}", a)
        self.power = power

    def _compute(self, x):
        return torch.sign(x) * torch.pow(torch.abs(x), self.power)

    def clone(self):
        return SignedPower(self.children[0].clone(), self.power)


class RankCS(Operator):
    def __init__(self, a):
        super().__init__("RankCS", a)

    def _compute(self, x):
        valid = ~torch.isnan(x)
        safe = torch.where(valid, x, torch.zeros_like(x))
        rank = safe.argsort(dim=0).argsort(dim=0).float()
        count = valid.sum(dim=0).float()
        out = rank / torch.clamp(count - 1.0, min=1.0)
        return torch.where(valid, out, torch.full_like(out, float("nan")))


class ZScoreCS(Operator):
    def __init__(self, a):
        super().__init__("ZScoreCS", a)

    def _compute(self, x):
        valid = ~torch.isnan(x)
        safe = torch.where(valid, x, torch.zeros_like(x))
        count = valid.sum(dim=0).float()
        mean = safe.sum(dim=0) / torch.clamp(count, min=1.0)
        diff = torch.where(valid, x - mean, torch.zeros_like(x))
        std = torch.sqrt((diff ** 2).sum(dim=0) / torch.clamp(count - 1.0, min=1.0)) + 1e-6
        return torch.where(valid, diff / std, torch.full_like(x, float("nan")))


class ScaleCS(Operator):
    def __init__(self, a):
        super().__init__("ScaleCS", a)

    def _compute(self, x):
        valid = ~torch.isnan(x)
        denom = torch.where(valid, torch.abs(x), torch.zeros_like(x)).sum(dim=0)
        return torch.where(valid, x / torch.clamp(denom, min=1e-6), torch.full_like(x, float("nan")))


def _ts_unary(op_name, x, d, fn):
    u = rolling_unfold(x, d)
    if u is None:
        return torch.zeros_like(x)
    return rolling_pad(x, fn(u), d)


class TsDelay(Operator):
    def __init__(self, a, d=5):
        super().__init__(f"TsDelay_{d}", a)
        self.d = d

    def _compute(self, x):
        return ts_delay_raw(x, self.d)

    def clone(self):
        return TsDelay(self.children[0].clone(), self.d)


class TsDelta(Operator):
    def __init__(self, a, d=5):
        super().__init__(f"TsDelta_{d}", a)
        self.d = d

    def _compute(self, x):
        return x - ts_delay_raw(x, self.d)

    def clone(self):
        return TsDelta(self.children[0].clone(), self.d)


class TsReturn(Operator):
    def __init__(self, a, d=5):
        super().__init__(f"TsReturn_{d}", a)
        self.d = d

    def _compute(self, x):
        return safe_div(x, ts_delay_raw(x, self.d)) - 1.0

    def clone(self):
        return TsReturn(self.children[0].clone(), self.d)


class TsMean(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsMean_{d}", a)
        self.d = d

    def _compute(self, x):
        return _ts_unary("TsMean", x, self.d, lambda u: nanmean_dim(u, 2))

    def clone(self):
        return TsMean(self.children[0].clone(), self.d)


class TsStd(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsStd_{d}", a)
        self.d = d

    def _compute(self, x):
        return _ts_unary("TsStd", x, self.d, lambda u: nanstd_dim(u, 2))

    def clone(self):
        return TsStd(self.children[0].clone(), self.d)


class TsRank(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsRank_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        valid = ~torch.isnan(u)
        safe = torch.where(valid, u, torch.zeros_like(u))
        rank = safe.argsort(dim=2).argsort(dim=2).float()
        count = valid.sum(dim=2).float()
        core = rank[:, :, -1] / torch.clamp(count - 1.0, min=1.0)
        core = torch.where(valid[:, :, -1], core, torch.full_like(core, float("nan")))
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsRank(self.children[0].clone(), self.d)


class TsDecayLinear(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsDecayLinear_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        weights = torch.arange(1, self.d + 1, device=x.device).float()
        weights = weights / weights.sum()
        valid = ~torch.isnan(u)
        safe = torch.where(valid, u, torch.zeros_like(u))
        return rolling_pad(x, (safe * weights).sum(dim=2), self.d)

    def clone(self):
        return TsDecayLinear(self.children[0].clone(), self.d)


class TsSlope(Operator):
    def __init__(self, a, d=20):
        super().__init__(f"TsSlope_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        t = torch.arange(self.d, device=x.device).float()
        t = t - t.mean()
        denom = (t ** 2).sum() + 1e-8
        mean_x = nanmean_dim(u, 2, keepdim=True)
        diff_x = torch.where(torch.isnan(u), torch.zeros_like(u), u - mean_x)
        return rolling_pad(x, (diff_x * t).sum(dim=2) / denom, self.d)

    def clone(self):
        return TsSlope(self.children[0].clone(), self.d)


class TsCorr(Operator):
    def __init__(self, a, b, d=20):
        super().__init__(f"TsCorr_{d}", a, b)
        self.d = d

    def _compute(self, x, y):
        ux, uy = rolling_unfold(x, self.d), rolling_unfold(y, self.d)
        if ux is None or uy is None:
            return torch.zeros_like(x)
        valid = ~torch.isnan(ux) & ~torch.isnan(uy)
        sx = torch.where(valid, ux, torch.zeros_like(ux))
        sy = torch.where(valid, uy, torch.zeros_like(uy))
        count = valid.sum(dim=2, keepdim=True).float()
        mx = sx.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        my = sy.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        dx = torch.where(valid, ux - mx, torch.zeros_like(ux))
        dy = torch.where(valid, uy - my, torch.zeros_like(uy))
        num = (dx * dy).sum(dim=2)
        den = torch.sqrt((dx ** 2).sum(dim=2)) * torch.sqrt((dy ** 2).sum(dim=2))
        core = torch.where(den > 1e-8, num / den, torch.full_like(num, float("nan")))
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsCorr(self.children[0].clone(), self.children[1].clone(), self.d)


class TsSum(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsSum_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        core = torch.where(~torch.isnan(u), u, torch.zeros_like(u)).sum(dim=2)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsSum(self.children[0].clone(), self.d)


class TsZScore(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsZScore_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        mean = nanmean_dim(u, 2)
        std = nanstd_dim(u, 2)
        core = (x[:, self.d - 1:] - mean) / (std + 1e-6)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsZScore(self.children[0].clone(), self.d)


class TsMin(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsMin_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        core = torch.min(torch.nan_to_num(u, nan=float("inf")), dim=2).values
        core = torch.where(torch.isinf(core), torch.full_like(core, float("nan")), core)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsMin(self.children[0].clone(), self.d)


class TsMax(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsMax_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        core = torch.max(torch.nan_to_num(u, nan=float("-inf")), dim=2).values
        core = torch.where(torch.isinf(core), torch.full_like(core, float("nan")), core)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsMax(self.children[0].clone(), self.d)


class TsArgMax(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsArgMax_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        core = torch.argmax(torch.nan_to_num(u, nan=float("-inf")), dim=2).float() / max(self.d - 1, 1)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsArgMax(self.children[0].clone(), self.d)


class TsArgMin(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsArgMin_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        core = torch.argmin(torch.nan_to_num(u, nan=float("inf")), dim=2).float() / max(self.d - 1, 1)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsArgMin(self.children[0].clone(), self.d)


class TsWMA(TsDecayLinear):
    def __init__(self, a, d=10):
        super().__init__(a, d)
        self.name = f"TsWMA_{d}"

    def clone(self):
        return TsWMA(self.children[0].clone(), self.d)


class TsEMA(Operator):
    def __init__(self, a, d=10):
        super().__init__(f"TsEMA_{d}", a)
        self.d = d

    def _compute(self, x):
        alpha = 2.0 / (self.d + 1.0)
        out = torch.empty_like(x)
        out[:, 0] = x[:, 0]
        for t in range(1, x.shape[1]):
            out[:, t] = alpha * torch.where(torch.isnan(x[:, t]), out[:, t - 1], x[:, t]) + (1 - alpha) * out[:, t - 1]
        out[:, : self.d - 1] = float("nan")
        return out

    def clone(self):
        return TsEMA(self.children[0].clone(), self.d)


class TsRSI(Operator):
    def __init__(self, a, d=14):
        super().__init__(f"TsRSI_{d}", a)
        self.d = d

    def _compute(self, x):
        diff = x - ts_delay_raw(x, 1)
        gain = torch.where(diff > 0, diff, torch.zeros_like(diff))
        loss = torch.where(diff < 0, -diff, torch.zeros_like(diff))
        ug, ul = rolling_unfold(gain, self.d), rolling_unfold(loss, self.d)
        if ug is None or ul is None:
            return torch.zeros_like(x)
        rs = safe_div(nanmean_dim(ug, 2), nanmean_dim(ul, 2) + 1e-8)
        return rolling_pad(x, (100.0 - 100.0 / (1.0 + rs)) / 100.0, self.d)

    def clone(self):
        return TsRSI(self.children[0].clone(), self.d)


class TsQuantile(Operator):
    def __init__(self, a, d=20, q=0.5):
        super().__init__(f"TsQuantile_{d}_{q}", a)
        self.d = d
        self.q = q

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        core = torch.quantile(torch.nan_to_num(u, nan=0.0), self.q, dim=2)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsQuantile(self.children[0].clone(), self.d, self.q)


class TsEntropy(Operator):
    def __init__(self, a, d=20):
        super().__init__(f"TsEntropy_{d}", a)
        self.d = d

    def _compute(self, x):
        u = rolling_unfold(x, self.d)
        if u is None:
            return torch.zeros_like(x)
        rank = torch.nan_to_num(u, nan=0.0).argsort(dim=2).argsort(dim=2).float()
        p = rank / torch.clamp(rank.sum(dim=2, keepdim=True), min=1e-6)
        core = -(p * torch.log(torch.clamp(p, min=1e-8))).sum(dim=2) / math.log(self.d)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsEntropy(self.children[0].clone(), self.d)


class TsCov(Operator):
    def __init__(self, a, b, d=20):
        super().__init__(f"TsCov_{d}", a, b)
        self.d = d

    def _compute(self, x, y):
        ux, uy = rolling_unfold(x, self.d), rolling_unfold(y, self.d)
        if ux is None or uy is None:
            return torch.zeros_like(x)
        valid = ~torch.isnan(ux) & ~torch.isnan(uy)
        sx = torch.where(valid, ux, torch.zeros_like(ux))
        sy = torch.where(valid, uy, torch.zeros_like(uy))
        count = valid.sum(dim=2, keepdim=True).float()
        mx = sx.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        my = sy.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        dx = torch.where(valid, ux - mx, torch.zeros_like(ux))
        dy = torch.where(valid, uy - my, torch.zeros_like(uy))
        core = (dx * dy).sum(dim=2) / torch.clamp(count.squeeze(2) - 1.0, min=1.0)
        return rolling_pad(x, core, self.d)

    def clone(self):
        return TsCov(self.children[0].clone(), self.children[1].clone(), self.d)


class TsRegBeta(Operator):
    def __init__(self, y_node, x_node, d=20):
        super().__init__(f"TsRegBeta_{d}", y_node, x_node)
        self.d = d

    def _compute(self, y, x):
        uy, ux = rolling_unfold(y, self.d), rolling_unfold(x, self.d)
        if uy is None or ux is None:
            return torch.zeros_like(y)
        valid = ~torch.isnan(uy) & ~torch.isnan(ux)
        sy = torch.where(valid, uy, torch.zeros_like(uy))
        sx = torch.where(valid, ux, torch.zeros_like(ux))
        count = valid.sum(dim=2, keepdim=True).float()
        my = sy.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        mx = sx.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        dy = torch.where(valid, uy - my, torch.zeros_like(uy))
        dx = torch.where(valid, ux - mx, torch.zeros_like(ux))
        var = (dx ** 2).sum(dim=2)
        beta = torch.where(var > 1e-8, (dx * dy).sum(dim=2) / var, torch.full_like(var, float("nan")))
        return rolling_pad(y, beta, self.d)

    def clone(self):
        return TsRegBeta(self.children[0].clone(), self.children[1].clone(), self.d)


class TsRegResidual(Operator):
    def __init__(self, y_node, x_node, d=20):
        super().__init__(f"TsRegResidual_{d}", y_node, x_node)
        self.d = d

    def _compute(self, y, x):
        uy, ux = rolling_unfold(y, self.d), rolling_unfold(x, self.d)
        if uy is None or ux is None:
            return torch.zeros_like(y)
        valid = ~torch.isnan(uy) & ~torch.isnan(ux)
        sy = torch.where(valid, uy, torch.zeros_like(uy))
        sx = torch.where(valid, ux, torch.zeros_like(ux))
        count = valid.sum(dim=2, keepdim=True).float()
        my = sy.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        mx = sx.sum(dim=2, keepdim=True) / torch.clamp(count, min=1.0)
        dy = torch.where(valid, uy - my, torch.zeros_like(uy))
        dx = torch.where(valid, ux - mx, torch.zeros_like(ux))
        beta = safe_div((dx * dy).sum(dim=2), (dx ** 2).sum(dim=2))
        alpha = my.squeeze(2) - beta * mx.squeeze(2)
        residual = y[:, self.d - 1 :] - (alpha + beta * x[:, self.d - 1 :])
        return rolling_pad(y, residual, self.d)

    def clone(self):
        return TsRegResidual(self.children[0].clone(), self.children[1].clone(), self.d)


TERMINALS = [
    "开盘价", "收盘价", "最高价", "最低价", "成交量", "成交额", "vwap",
    "return_1", "return_2", "return_4", "oc_return", "co_return",
    "hl_spread", "ho_gap", "lo_gap", "vwap_close_gap",
    "amount_per_volume", "log_volume", "log_amount",
]

CONSTANTS = [-20, -10, -5, -3, -2, -1, -0.5, -0.25, -0.1, 0.1, 0.25, 0.5, 1, 2, 3, 5, 10, 20]
UNARY_OPS = [AbsOp, Neg, LogOp, SqrtOp, RankCS, ZScoreCS, ScaleCS]
SAFE_BINARY_OPS = [Add, Sub, Mul, Max2, Min2]
RISKY_BINARY_OPS = [Div]
TS_UNARY_OPS = [
    TsDelay, TsDelta, TsReturn, TsMean, TsSum, TsStd, TsZScore, TsMin, TsMax,
    TsRank, TsArgMax, TsArgMin, TsDecayLinear, TsWMA, TsEMA, TsRSI, TsQuantile,
    TsEntropy, TsSlope,
]
SAFE_TS_BINARY_OPS = [TsCorr, TsCov]
RISKY_TS_BINARY_OPS = [TsRegBeta, TsRegResidual]
RISKY_OPERATOR_NAMES = {"Div", "TsRegResidual", "TsRegBeta", "TsCov", "TsCorr"}

_TS_UNARY_WITH_WINDOW = {
    TsDelay, TsDelta, TsReturn, TsMean, TsSum, TsStd, TsZScore, TsMin, TsMax,
    TsRank, TsArgMax, TsArgMin, TsDecayLinear, TsWMA, TsEMA, TsRSI, TsEntropy, TsSlope,
}
_TS_BINARY_WITH_WINDOW = {TsCorr, TsCov, TsRegBeta, TsRegResidual}


def make_unary(op_class, child, windows=None):
    windows = windows or WINDOWS
    if op_class in _TS_UNARY_WITH_WINDOW:
        return op_class(child, d=random.choice(windows))
    if op_class is TsQuantile:
        return TsQuantile(child, d=random.choice(windows), q=random.choice([0.2, 0.3, 0.5, 0.7, 0.8]))
    if op_class is SignedPower:
        return SignedPower(child, power=random.choice([0.5, 1.5, 2.0, 3.0]))
    return op_class(child)


def make_binary(op_class, left, right, windows=None):
    windows = windows or WINDOWS
    if op_class in _TS_BINARY_WITH_WINDOW:
        return op_class(left, right, d=random.choice(windows))
    return op_class(left, right)


def choose_binary_operator():
    r = random.random()
    if r < 0.68:
        return random.choice(SAFE_BINARY_OPS)
    if r < 0.88:
        return random.choice(SAFE_TS_BINARY_OPS)
    if r < 0.96:
        return random.choice(RISKY_BINARY_OPS)
    return random.choice(RISKY_TS_BINARY_OPS)


def random_terminal():
    if random.random() < 0.90:
        return Terminal(random.choice(TERMINALS))
    return Constant(random.choice(CONSTANTS))


def generate_random_tree(depth, max_depth, min_tree_nodes=3):
    if depth >= max_depth:
        return random_terminal()
    if depth > 1 and random.random() < (0.08 + 0.06 * depth):
        return random_terminal()

    r = random.random()
    if r < 0.35:
        child = generate_random_tree(depth + 1, max_depth, min_tree_nodes)
        op = random.choice(UNARY_OPS + TS_UNARY_OPS) if random.random() < 0.92 else SignedPower
        return make_unary(op, child)
    if r < 0.80:
        left = generate_random_tree(depth + 1, max_depth, min_tree_nodes)
        right = generate_random_tree(depth + 1, max_depth, min_tree_nodes)
        return make_binary(choose_binary_operator(), left, right)

    child = generate_random_tree(depth + 1, max_depth, min_tree_nodes)
    ts_op = random.choice([TsMean, TsStd, TsZScore, TsRank, TsDecayLinear, TsEMA, TsSlope])
    return random.choice([RankCS, ZScoreCS])(make_unary(ts_op, child))