"""
Unified operator mapping system for FactorDB

This module provides a unified operator registry that maps operators from different
factor mining frameworks to a common set of operations.
"""

import numpy as np
import pandas as pd
from typing import Dict, Any, Callable, List, Optional, Union
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)


@dataclass
class OperatorInfo:
    """Information about an operator"""
    name: str  # Canonical name
    func: Callable  # Implementation function
    args: List[str]  # Argument names
    description: str  # Description
    aliases: List[str]  # Alternative names from different frameworks


class UnifiedOperatorRegistry:
    """
    Unified operator registry that maps operators from different factor mining
    frameworks to a common set of operations.

    Supports mappings from:
    - op1.py (LLM Factor Mining System - CamelCase)
    - op2.py (MAS Factor Mining - lowercase with ts_/cs_ prefixes)
    - op3.h (C++ GP Factor Mining)
    """

    def __init__(self):
        """Initialize unified operator registry"""
        self._operators: Dict[str, OperatorInfo] = {}
        self._alias_map: Dict[str, str] = {}  # alias -> canonical name
        self._register_all_operators()

    def _register_all_operators(self):
        """Register all operators with their aliases"""
        self._register_unary_operators()
        self._register_binary_operators()
        self._register_ts_operators()
        self._register_cs_operators()
        self._register_conditional_operators()

    def _register_operator(self, info: OperatorInfo):
        """Register an operator and its aliases"""
        self._operators[info.name] = info
        self._alias_map[info.name.lower()] = info.name
        for alias in info.aliases:
            self._alias_map[alias.lower()] = info.name

    def _register_unary_operators(self):
        """Register unary operators (single input)"""
        operators = [
            OperatorInfo(
                name="abs",
                func=lambda x: np.abs(x),
                args=["x"],
                description="Absolute value",
                aliases=["Abs", "abs_op"]
            ),
            OperatorInfo(
                name="sign",
                func=lambda x: np.sign(x),
                args=["x"],
                description="Sign function: 1 if x > 0, -1 if x < 0, 0 if x = 0",
                aliases=["Sign"]
            ),
            OperatorInfo(
                name="log",
                func=lambda x: np.log(np.clip(np.abs(x), 1e-10, None)),
                args=["x"],
                description="Natural logarithm",
                aliases=["Log"]
            ),
            OperatorInfo(
                name="neg",
                func=lambda x: -x,
                args=["x"],
                description="Negation",
                aliases=["Neg", "opp"]
            ),
            OperatorInfo(
                name="inv",
                func=lambda x: 1.0 / (x + np.where(x >= 0, 1e-10, -1e-10)),
                args=["x"],
                description="Inverse (1/x)",
                aliases=["Inv"]
            ),
            OperatorInfo(
                name="sqrt",
                func=lambda x: np.sqrt(np.clip(x, 0, None)),
                args=["x"],
                description="Square root",
                aliases=["Sqrt"]
            ),
            OperatorInfo(
                name="square",
                func=lambda x: x ** 2,
                args=["x"],
                description="Square",
                aliases=["Square"]
            ),
            OperatorInfo(
                name="sigmoid",
                func=lambda x: 1 / (1 + np.exp(-np.clip(x, -500, 500))),
                args=["x"],
                description="Sigmoid function",
                aliases=["Sigmoid"]
            ),
            OperatorInfo(
                name="tanh",
                func=lambda x: np.tanh(x),
                args=["x"],
                description="Hyperbolic tangent",
                aliases=["Tanh"]
            ),
        ]
        for op in operators:
            self._register_operator(op)

    def _register_binary_operators(self):
        """Register binary operators (two inputs)"""
        operators = [
            OperatorInfo(
                name="add",
                func=lambda x, y: x + y,
                args=["x", "y"],
                description="Addition",
                aliases=["Add"]
            ),
            OperatorInfo(
                name="sub",
                func=lambda x, y: x - y,
                args=["x", "y"],
                description="Subtraction",
                aliases=["Sub"]
            ),
            OperatorInfo(
                name="mul",
                func=lambda x, y: x * y,
                args=["x", "y"],
                description="Multiplication",
                aliases=["Mul"]
            ),
            OperatorInfo(
                name="div",
                func=self._safe_div,
                args=["x", "y"],
                description="Division (with epsilon for safety)",
                aliases=["Div"]
            ),
            OperatorInfo(
                name="max",
                func=lambda x, y: np.maximum(x, y),
                args=["x", "y"],
                description="Element-wise maximum of two arrays",
                aliases=["Max2", "max_op"]
            ),
            OperatorInfo(
                name="min",
                func=lambda x, y: np.minimum(x, y),
                args=["x", "y"],
                description="Element-wise minimum of two arrays",
                aliases=["Min2", "min_op"]
            ),
            OperatorInfo(
                name="power",
                func=lambda x, y: np.power(np.clip(np.abs(x), 1e-10, None), y),
                args=["x", "y"],
                description="Power (x^y)",
                aliases=["Power"]
            ),
        ]
        for op in operators:
            self._register_operator(op)

    def _register_ts_operators(self):
        """Register time-series operators"""
        operators = [
            OperatorInfo(
                name="ts_mean",
                func=self._ts_mean,
                args=["x", "window"],
                description="Rolling mean over window",
                aliases=["Mean", "ts_mean"]
            ),
            OperatorInfo(
                name="ts_std",
                func=self._ts_std,
                args=["x", "window"],
                description="Rolling standard deviation over window",
                aliases=["Std", "ts_std"]
            ),
            OperatorInfo(
                name="ts_var",
                func=self._ts_var,
                args=["x", "window"],
                description="Rolling variance over window",
                aliases=["Var", "ts_var"]
            ),
            OperatorInfo(
                name="ts_max",
                func=self._ts_max,
                args=["x", "window"],
                description="Rolling maximum over window",
                aliases=["Max", "ts_max"]
            ),
            OperatorInfo(
                name="ts_min",
                func=self._ts_min,
                args=["x", "window"],
                description="Rolling minimum over window",
                aliases=["Min", "ts_min"]
            ),
            OperatorInfo(
                name="ts_sum",
                func=self._ts_sum,
                args=["x", "window"],
                description="Rolling sum over window",
                aliases=["Sum", "ts_sum"]
            ),
            OperatorInfo(
                name="ts_median",
                func=self._ts_median,
                args=["x", "window"],
                description="Rolling median over window",
                aliases=["Median"]
            ),
            OperatorInfo(
                name="ts_delta",
                func=self._ts_delta,
                args=["x", "window"],
                description="Difference from window periods ago",
                aliases=["Delta", "ts_delta"]
            ),
            OperatorInfo(
                name="ts_delay",
                func=self._ts_delay,
                args=["x", "window"],
                description="Value from window periods ago",
                aliases=["Delay", "Ref", "ts_delay"]
            ),
            OperatorInfo(
                name="ts_return",
                func=self._ts_return,
                args=["x", "window"],
                description="Percentage return over window periods",
                aliases=["Return", "ts_return"]
            ),
            OperatorInfo(
                name="ts_slope",
                func=self._ts_slope,
                args=["x", "window"],
                description="Linear regression slope over window",
                aliases=["Slope"]
            ),
            OperatorInfo(
                name="ts_corr",
                func=self._ts_corr,
                args=["x", "y", "window"],
                description="Rolling correlation between x and y over window",
                aliases=["Corr", "ts_corr", "ts_pears_corr"]
            ),
            OperatorInfo(
                name="ts_cov",
                func=self._ts_cov,
                args=["x", "y", "window"],
                description="Rolling covariance between x and y over window",
                aliases=["Cov", "ts_cov"]
            ),
            OperatorInfo(
                name="ts_ema",
                func=self._ts_ema,
                args=["x", "window"],
                description="Exponential moving average",
                aliases=["EMA", "ts_decay_exp"]
            ),
            OperatorInfo(
                name="ts_wma",
                func=self._ts_wma,
                args=["x", "window"],
                description="Weighted moving average (linear decay)",
                aliases=["WMA", "ts_decay_linear"]
            ),
            OperatorInfo(
                name="ts_skew",
                func=self._ts_skew,
                args=["x", "window"],
                description="Rolling skewness over window",
                aliases=["Skew", "ts_skew"]
            ),
            OperatorInfo(
                name="ts_kurt",
                func=self._ts_kurt,
                args=["x", "window"],
                description="Rolling kurtosis over window",
                aliases=["Kurt", "ts_kurt"]
            ),
            OperatorInfo(
                name="ts_rank",
                func=self._ts_rank,
                args=["x", "window"],
                description="Time-series rank within window",
                aliases=["Rank_ts", "ts_rank"]
            ),
            OperatorInfo(
                name="ts_zscore",
                func=self._ts_zscore,
                args=["x", "window"],
                description="Rolling z-score over window",
                aliases=["Zscore", "ts_zscore"]
            ),
            OperatorInfo(
                name="ts_argmax",
                func=self._ts_argmax,
                args=["x", "window"],
                description="Position of maximum value within window",
                aliases=["ArgMax", "ts_argmax"]
            ),
            OperatorInfo(
                name="ts_argmin",
                func=self._ts_argmin,
                args=["x", "window"],
                description="Position of minimum value within window",
                aliases=["ArgMin", "ts_argmin"]
            ),
            OperatorInfo(
                name="ts_prod",
                func=self._ts_prod,
                args=["x", "window"],
                description="Rolling product over window",
                aliases=["ts_product", "ts_prod"]
            ),
            OperatorInfo(
                name="ts_quantile",
                func=self._ts_quantile,
                args=["x", "window", "q"],
                description="Rolling quantile over window",
                aliases=["Quantile", "ts_quantile"]
            ),
        ]
        for op in operators:
            self._register_operator(op)

    def _register_cs_operators(self):
        """Register cross-sectional operators"""
        operators = [
            OperatorInfo(
                name="cs_rank",
                func=self._cs_rank,
                args=["x"],
                description="Cross-sectional rank (percentile)",
                aliases=["Rank", "cs_rank"]
            ),
            OperatorInfo(
                name="cs_zscore",
                func=self._cs_zscore,
                args=["x"],
                description="Cross-sectional z-score normalization",
                aliases=["cs_zscore"]
            ),
            OperatorInfo(
                name="cs_demean",
                func=self._cs_demean,
                args=["x"],
                description="Cross-sectional demeaning",
                aliases=["cs_demean"]
            ),
            OperatorInfo(
                name="cs_scale",
                func=self._cs_scale,
                args=["x"],
                description="Cross-sectional scaling by absolute sum",
                aliases=["cs_scale"]
            ),
        ]
        for op in operators:
            self._register_operator(op)

    def _register_conditional_operators(self):
        """Register conditional operators"""
        operators = [
            OperatorInfo(
                name="if_else",
                func=lambda cond, x, y: np.where(cond, x, y),
                args=["cond", "x", "y"],
                description="Conditional selection: if cond then x else y",
                aliases=["if_else"]
            ),
            OperatorInfo(
                name="greater",
                func=lambda x, y: x > y,
                args=["x", "y"],
                description="Greater than comparison",
                aliases=["greater"]
            ),
            OperatorInfo(
                name="less",
                func=lambda x, y: x < y,
                args=["x", "y"],
                description="Less than comparison",
                aliases=["less"]
            ),
        ]
        for op in operators:
            self._register_operator(op)

    # Helper implementations
    @staticmethod
    def _safe_div(x, y):
        """Safe division with handling for near-zero denominators"""
        result = x / (y + np.where(np.abs(y) < 1e-10, np.sign(y + 1e-20) * 1e-10, 0))
        if isinstance(result, pd.DataFrame):
            result = result.replace([np.inf, -np.inf], np.nan)
        else:
            result = np.where(np.isinf(result), np.nan, result)
        return result

    @staticmethod
    def _ensure_int_window(window) -> int:
        """Ensure window is a valid positive integer"""
        try:
            window = int(window)
            return max(1, window)
        except (TypeError, ValueError):
            return 1

    @staticmethod
    def _ensure_series(x) -> pd.Series:
        """Ensure input is a pandas Series"""
        if isinstance(x, pd.Series):
            return x
        elif isinstance(x, np.ndarray):
            return pd.Series(x)
        elif isinstance(x, pd.DataFrame):
            # Take first column if DataFrame
            return x.iloc[:, 0]
        else:
            return pd.Series([x] if np.isscalar(x) else x)

    # Time-series operator implementations
    def _ts_mean(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).mean()

    def _ts_std(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=2).std()

    def _ts_var(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=2).var()

    def _ts_max(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).max()

    def _ts_min(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).min()

    def _ts_sum(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).sum()

    def _ts_median(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).median()

    def _ts_delta(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x - x.shift(window)

    def _ts_delay(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.shift(window)

    def _ts_return(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.pct_change(periods=window)

    def _ts_slope(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        def slope(arr):
            if len(arr) < 2:
                return np.nan
            y = arr.values if hasattr(arr, 'values') else arr
            x_vals = np.arange(len(y))
            mask = ~np.isnan(y)
            if mask.sum() < 2:
                return np.nan
            return np.polyfit(x_vals[mask], y[mask], 1)[0]
        return x.rolling(window=window, min_periods=2).apply(slope, raw=False)

    def _ts_corr(self, x, y, window) -> pd.Series:
        x = self._ensure_series(x)
        y = self._ensure_series(y)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=2).corr(y)

    def _ts_cov(self, x, y, window) -> pd.Series:
        x = self._ensure_series(x)
        y = self._ensure_series(y)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=2).cov(y)

    def _ts_ema(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.ewm(span=window, adjust=False, min_periods=1).mean()

    def _ts_wma(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        weights = np.arange(1, window + 1, dtype=float)
        def weighted_mean(s):
            n = len(s)
            w = weights[-n:]
            return np.sum(s * w) / w.sum()
        return x.rolling(window=window, min_periods=1).apply(weighted_mean, raw=True)

    def _ts_skew(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=3).skew()

    def _ts_kurt(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=4).kurt()

    def _ts_rank(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        def ts_rank(arr):
            return pd.Series(arr).rank(pct=True).iloc[-1]
        return x.rolling(window=window, min_periods=1).apply(ts_rank, raw=True)

    def _ts_zscore(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        mean = x.rolling(window=window, min_periods=1).mean()
        std = x.rolling(window=window, min_periods=1).std()
        return (x - mean) / (std + 1e-10)

    def _ts_argmax(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).apply(lambda arr: np.argmax(arr), raw=True)

    def _ts_argmin(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).apply(lambda arr: np.argmin(arr), raw=True)

    def _ts_prod(self, x, window) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        return x.rolling(window=window, min_periods=1).apply(np.prod, raw=True)

    def _ts_quantile(self, x, window, q=0.5) -> pd.Series:
        x = self._ensure_series(x)
        window = self._ensure_int_window(window)
        try:
            q_val = float(q)
        except (TypeError, ValueError):
            q_val = 0.5
        q_val = min(max(q_val, 0.0), 1.0)
        return x.rolling(window=window, min_periods=1).quantile(q_val)

    # Cross-sectional operator implementations
    @staticmethod
    def _cs_rank(x):
        """Cross-sectional rank"""
        if isinstance(x, pd.DataFrame):
            return x.rank(axis=1, pct=True)
        elif isinstance(x, pd.Series):
            return x.rank(pct=True)
        return pd.Series(x).rank(pct=True)

    @staticmethod
    def _cs_zscore(x):
        """Cross-sectional z-score"""
        if isinstance(x, pd.DataFrame):
            mean = x.mean(axis=1)
            std = x.std(axis=1)
            return x.sub(mean, axis=0).div(std + 1e-10, axis=0)
        elif isinstance(x, pd.Series):
            return (x - x.mean()) / (x.std() + 1e-10)
        return x

    @staticmethod
    def _cs_demean(x):
        """Cross-sectional demeaning"""
        if isinstance(x, pd.DataFrame):
            mean = x.mean(axis=1)
            return x.sub(mean, axis=0)
        elif isinstance(x, pd.Series):
            return x - x.mean()
        return x

    @staticmethod
    def _cs_scale(x):
        """Cross-sectional scaling by absolute sum"""
        if isinstance(x, pd.DataFrame):
            abs_sum = x.abs().sum(axis=1)
            return x.div(abs_sum + 1e-10, axis=0)
        elif isinstance(x, pd.Series):
            return x / (x.abs().sum() + 1e-10)
        return x

    def get_operator(self, name: str) -> Optional[OperatorInfo]:
        """Get operator by name or alias"""
        canonical = self._alias_map.get(name.lower())
        if canonical:
            return self._operators.get(canonical)
        return None

    def get_canonical_name(self, name: str) -> Optional[str]:
        """Get canonical operator name from an alias"""
        return self._alias_map.get(name.lower())

    def get_all_operators(self) -> Dict[str, OperatorInfo]:
        """Get all registered operators"""
        return self._operators

    def get_operator_func(self, name: str) -> Optional[Callable]:
        """Get operator function by name or alias"""
        op = self.get_operator(name)
        return op.func if op else None

    def has_operator(self, name: str) -> bool:
        """Check if an operator exists (by name or alias)"""
        return name.lower() in self._alias_map

    def list_operators(self) -> List[str]:
        """List all canonical operator names"""
        return list(self._operators.keys())

    def list_aliases(self, name: str) -> List[str]:
        """List all aliases for an operator"""
        op = self.get_operator(name)
        return op.aliases if op else []


# Global registry instance
_global_registry: Optional[UnifiedOperatorRegistry] = None


def get_operator_registry() -> UnifiedOperatorRegistry:
    """Get the global operator registry instance"""
    global _global_registry
    if _global_registry is None:
        _global_registry = UnifiedOperatorRegistry()
    return _global_registry
