"""Automated operator invention via LLM-guided proposal and validation.

Uses an LLM to propose novel operator definitions (as NumPy functions),
validates them in a sandboxed environment, and checks for differentiation
from existing operators and IC contribution.
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np

from src.factorminer.agent.llm_interface import LLMProvider
from src.factorminer.core.types import OPERATOR_REGISTRY, OperatorSpec

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class ProposedOperator:
    """A single operator proposal generated by the LLM.

    Attributes
    ----------
    name : str
        Canonical name for the operator (e.g. ``"ExpDecayDiff"``).
    arity : int
        Number of expression children (1 = unary, 2 = binary).
    description : str
        Short human-readable description.
    numpy_code : str
        Python source defining a function called ``compute``.
        The function signature must accept (M, T)-shaped ndarrays and
        return an (M, T)-shaped ndarray.
    param_names : tuple of str
        Names of extra numeric parameters.
    param_defaults : dict
        Default value for each parameter.
    param_ranges : dict
        Valid (inclusive) range for each parameter.
    rationale : str
        Why this operator might be useful for alpha factor construction.
    based_on : list of str
        Existing operators that inspired this proposal.
    """

    name: str
    arity: int
    description: str
    numpy_code: str
    param_names: Tuple[str, ...] = ()
    param_defaults: Dict[str, float] = field(default_factory=dict)
    param_ranges: Dict[str, Tuple[float, float]] = field(default_factory=dict)
    rationale: str = ""
    based_on: List[str] = field(default_factory=list)


@dataclass
class ValidationResult:
    """Result of validating a proposed operator.

    Attributes
    ----------
    valid : bool
        True if the operator passed all validation checks.
    error : str
        Error message if validation failed; empty string on success.
    output_shape_ok : bool
        True if the operator output has the correct (M, T) shape.
    nan_ratio : float
        Fraction of NaN values in the operator output.
    differentiates_from_existing : bool
        True if the operator is sufficiently different from all existing operators.
    ic_contribution : float
        Information coefficient of a simple factor using this operator.
    """

    valid: bool
    error: str = ""
    output_shape_ok: bool = False
    nan_ratio: float = 1.0
    differentiates_from_existing: bool = False
    ic_contribution: float = 0.0


# ---------------------------------------------------------------------------
# Sandbox security: allowed names in exec()
# ---------------------------------------------------------------------------

_SAFE_GLOBALS: Dict[str, Any] = {
    "np": np,
    "numpy": np,
    "__builtins__": {},
}

# Explicitly blocked tokens in submitted code.  If any of these appear in the
# raw source string, the code is rejected *before* exec().
_BLOCKED_TOKENS: Tuple[str, ...] = (
    "import ",
    "__import__",
    "os.",
    "sys.",
    "subprocess",
    "open(",
    "exec(",
    "eval(",
    "compile(",
    "getattr(",
    "setattr(",
    "delattr(",
    "globals(",
    "locals(",
    "__class__",
    "__subclasses__",
    "__bases__",
    "__mro__",
    "breakpoint(",
    "exit(",
    "quit(",
)


# ---------------------------------------------------------------------------
# OperatorInventor
# ---------------------------------------------------------------------------

class OperatorInventor:
    """Proposes and validates new operators using an LLM.

    Parameters
    ----------
    llm_provider : LLMProvider
        LLM backend used to generate proposals.
    data_tensor : np.ndarray
        Shape ``(M, T, F)`` -- M stocks, T time steps, F features.
        Used as test data for validation.
    returns : np.ndarray
        Shape ``(M, T)`` -- forward returns for IC measurement.
    max_proposals_per_round : int
        Maximum number of proposals to request per LLM call.
    """

    _SYSTEM_PROMPT = (
        "You are an expert in quantitative finance operator design. "
        "Your task is to invent novel numerical operators that transform "
        "stock market time-series data (shape: stocks x time) into alpha "
        "signals. Each operator is a pure NumPy function.\n\n"
        "RULES:\n"
        "1. Each operator must define a function called `compute`.\n"
        "2. The function receives ndarray inputs of shape (M, T) and must "
        "return an ndarray of shape (M, T).\n"
        "3. You may ONLY use numpy (imported as `np`). No other imports.\n"
        "4. Handle NaN values gracefully (use np.nan, np.nanmean, etc.).\n"
        "5. Avoid division by zero -- use np.where or add epsilon.\n"
        "6. Operators should be economically meaningful for alpha factor "
        "construction.\n"
        "7. Do NOT use os, sys, subprocess, open, exec, eval, or any "
        "filesystem/network access.\n"
    )

    def __init__(
        self,
        llm_provider: LLMProvider,
        data_tensor: np.ndarray,
        returns: np.ndarray,
        max_proposals_per_round: int = 5,
    ) -> None:
        if data_tensor.ndim != 3:
            raise ValueError(
                f"data_tensor must be 3-D (M, T, F), got shape {data_tensor.shape}"
            )
        if returns.ndim != 2:
            raise ValueError(
                f"returns must be 2-D (M, T), got shape {returns.shape}"
            )
        self.llm = llm_provider
        self.data_tensor = data_tensor
        self.returns = returns
        self.max_proposals = max_proposals_per_round
        # Pre-compute test slices for validation
        self._test_x = data_tensor[:, :, 0]  # first feature as default input
        self._test_y = (
            data_tensor[:, :, 1] if data_tensor.shape[2] > 1 else data_tensor[:, :, 0]
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def propose_operators(
        self,
        existing_operators: Dict[str, OperatorSpec],
        successful_patterns: Optional[List[str]] = None,
    ) -> List[ProposedOperator]:
        """Ask the LLM to propose new operators.

        Parameters
        ----------
        existing_operators : dict
            Mapping of name -> OperatorSpec for already-registered operators.
        successful_patterns : list of str, optional
            Natural-language descriptions of patterns that have worked well.

        Returns
        -------
        list of ProposedOperator
        """
        successful_patterns = successful_patterns or []
        prompt = self._build_proposal_prompt(existing_operators, successful_patterns)
        logger.info(
            "Requesting %d operator proposals from %s",
            self.max_proposals,
            self.llm.provider_name,
        )
        raw = self.llm.generate(
            system_prompt=self._SYSTEM_PROMPT,
            user_prompt=prompt,
            temperature=0.9,
            max_tokens=8192,
        )
        proposals = self._parse_proposals(raw)
        logger.info("Parsed %d proposals from LLM output", len(proposals))
        return proposals

    def validate_operator(self, proposal: ProposedOperator) -> ValidationResult:
        """Validate a proposed operator through compilation, execution, and IC check.

        Parameters
        ----------
        proposal : ProposedOperator

        Returns
        -------
        ValidationResult
        """
        # Step 1: compile safely
        fn = self._compile_safely(proposal.numpy_code)
        if fn is None:
            return ValidationResult(valid=False, error="Compilation failed or blocked code detected")

        # Step 2: execute on test data, check output shape
        try:
            if proposal.arity == 1:
                output = fn(self._test_x)
            elif proposal.arity >= 2:
                output = fn(self._test_x, self._test_y)
            else:
                output = fn(self._test_x)
        except Exception as exc:
            return ValidationResult(valid=False, error=f"Execution error: {exc}")

        if not isinstance(output, np.ndarray):
            return ValidationResult(valid=False, error="Output is not an ndarray")

        M, T = self._test_x.shape
        if output.shape != (M, T):
            return ValidationResult(
                valid=False,
                error=f"Shape mismatch: expected ({M}, {T}), got {output.shape}",
                output_shape_ok=False,
            )

        # Step 3: NaN ratio
        nan_ratio = float(np.isnan(output).sum()) / output.size if output.size > 0 else 1.0
        if nan_ratio > 0.5:
            return ValidationResult(
                valid=False,
                error=f"NaN ratio too high: {nan_ratio:.2%}",
                output_shape_ok=True,
                nan_ratio=nan_ratio,
            )

        # Step 4: differentiation from existing operators
        differentiates = self._check_differentiation(fn, proposal)

        # Step 5: IC contribution
        ic = self._measure_ic_contribution(fn, proposal)

        valid = differentiates  # must differentiate; IC is informational
        error = "" if valid else "Too correlated with existing operators (r > 0.9)"

        return ValidationResult(
            valid=valid,
            error=error,
            output_shape_ok=True,
            nan_ratio=nan_ratio,
            differentiates_from_existing=differentiates,
            ic_contribution=ic,
        )

    # ------------------------------------------------------------------
    # Internal: prompt building & parsing
    # ------------------------------------------------------------------

    def _build_proposal_prompt(
        self,
        existing_ops: Dict[str, OperatorSpec],
        patterns: List[str],
    ) -> str:
        """Format the user prompt for operator proposals."""
        lines: List[str] = []

        # Existing operators summary
        lines.append("## EXISTING OPERATORS (do NOT duplicate these)")
        for name, spec in sorted(existing_ops.items()):
            lines.append(f"- {name} (arity={spec.arity}): {spec.description}")

        # Successful patterns
        if patterns:
            lines.append("\n## SUCCESSFUL PATTERNS (build on these themes)")
            for p in patterns:
                lines.append(f"  * {p}")

        # Request
        lines.append(f"\n## REQUEST")
        lines.append(
            f"Propose exactly {self.max_proposals} new operators. "
            f"For each operator, output a JSON object with these fields:"
        )
        lines.append(
            '  {"name": "OpName", "arity": 1, "description": "...", '
            '"numpy_code": "def compute(x):\\n    ...", '
            '"param_names": [], "param_defaults": {}, "param_ranges": {}, '
            '"rationale": "...", "based_on": ["ExistingOp1"]}'
        )
        lines.append(
            "\nOutput each JSON object on a separate line, preceded by "
            "the line number (e.g., '1. {...}')."
        )
        lines.append(
            "\nIMPORTANT: The `numpy_code` field must define a function "
            "called `compute` that accepts ndarray(s) of shape (M, T) "
            "and returns an ndarray of shape (M, T). Use only numpy (as np)."
        )
        return "\n".join(lines)

    def _parse_proposals(self, raw: str) -> List[ProposedOperator]:
        """Parse LLM output into ProposedOperator objects."""
        proposals: List[ProposedOperator] = []

        # Try to find JSON objects in the text
        # Pattern: optional number prefix, then a JSON object
        json_pattern = re.compile(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}')
        matches = json_pattern.findall(raw)

        for match in matches:
            try:
                data = json.loads(match)
            except json.JSONDecodeError:
                logger.debug("Failed to parse JSON: %s", match[:100])
                continue

            name = data.get("name", "")
            if not name:
                continue

            # Normalize numpy_code: replace \\n with actual newlines
            numpy_code = data.get("numpy_code", "")
            if "\\n" in numpy_code:
                numpy_code = numpy_code.replace("\\n", "\n")

            proposal = ProposedOperator(
                name=name,
                arity=int(data.get("arity", 1)),
                description=data.get("description", ""),
                numpy_code=numpy_code,
                param_names=tuple(data.get("param_names", [])),
                param_defaults=data.get("param_defaults", {}),
                param_ranges={
                    k: tuple(v) for k, v in data.get("param_ranges", {}).items()
                },
                rationale=data.get("rationale", ""),
                based_on=data.get("based_on", []),
            )
            proposals.append(proposal)

            if len(proposals) >= self.max_proposals:
                break

        return proposals

    # ------------------------------------------------------------------
    # Internal: safe compilation & validation helpers
    # ------------------------------------------------------------------

    def _compile_safely(self, code: str) -> Optional[Callable]:
        """Compile operator code in a restricted sandbox.

        SECURITY: Only numpy is available. No filesystem, network,
        subprocess, or introspection access is permitted.

        Parameters
        ----------
        code : str
            Python source defining a function called ``compute``.

        Returns
        -------
        Callable or None
            The compiled ``compute`` function, or None if compilation
            failed or the code was rejected for security reasons.
        """
        # Pre-scan for blocked tokens
        code_lower = code.lower()
        for token in _BLOCKED_TOKENS:
            if token.lower() in code_lower:
                logger.warning("Blocked token '%s' found in operator code", token)
                return None

        # Restricted exec
        safe_ns: Dict[str, Any] = dict(_SAFE_GLOBALS)
        try:
            exec(code, safe_ns)  # noqa: S102 -- intentional sandboxed exec
        except Exception as exc:
            logger.warning("Operator compilation failed: %s", exc)
            return None

        fn = safe_ns.get("compute")
        if fn is None or not callable(fn):
            logger.warning("No callable 'compute' found in operator code")
            return None

        return fn

    def _check_differentiation(self, fn: Callable, proposal: ProposedOperator) -> bool:
        """Check that the operator output is not too correlated with existing operators.

        Computes the new operator on test data and correlates the result with
        outputs from a sample of existing operators.  If max |correlation| > 0.9,
        the operator is considered redundant.

        Returns
        -------
        bool
            True if the operator differentiates sufficiently.
        """
        try:
            if proposal.arity == 1:
                new_output = fn(self._test_x)
            else:
                new_output = fn(self._test_x, self._test_y)
        except Exception:
            return False

        new_flat = new_output.flatten()
        valid_mask = ~np.isnan(new_flat)
        if valid_mask.sum() < 10:
            return False

        # Compare against a sample of existing operator implementations
        from src.factorminer.operators.registry import OPERATOR_REGISTRY as RUNTIME_REG

        sample_names = list(RUNTIME_REG.keys())[:20]  # sample for efficiency
        for name in sample_names:
            entry = RUNTIME_REG.get(name)
            if entry is None:
                continue
            spec, np_fn, _ = entry
            if np_fn is None:
                continue
            try:
                if spec.arity == 1:
                    existing_output = np_fn(self._test_x)
                elif spec.arity == 2:
                    existing_output = np_fn(self._test_x, self._test_y)
                else:
                    continue
            except Exception:
                continue

            if existing_output.shape != new_output.shape:
                continue

            ex_flat = existing_output.flatten()
            both_valid = valid_mask & ~np.isnan(ex_flat)
            if both_valid.sum() < 10:
                continue

            corr = np.corrcoef(new_flat[both_valid], ex_flat[both_valid])[0, 1]
            if abs(corr) > 0.9:
                logger.info(
                    "Proposed operator '%s' too correlated with '%s' (r=%.3f)",
                    proposal.name,
                    name,
                    corr,
                )
                return False

        return True

    def _measure_ic_contribution(self, fn: Callable, proposal: ProposedOperator) -> float:
        """Measure the Information Coefficient of a simple factor using the operator.

        Constructs a basic factor as CsRank(NewOp(data)) and computes
        the rank IC against forward returns.

        Returns
        -------
        float
            Mean rank IC across time steps.
        """
        try:
            if proposal.arity == 1:
                raw = fn(self._test_x)
            else:
                raw = fn(self._test_x, self._test_y)
        except Exception:
            return 0.0

        if raw.shape != self.returns.shape:
            return 0.0

        # Cross-sectional rank at each time step
        M, T = raw.shape
        ranked = np.full_like(raw, np.nan)
        for t in range(T):
            col = raw[:, t]
            valid = ~np.isnan(col)
            if valid.sum() < 3:
                continue
            order = col[valid].argsort().argsort()
            ranked[valid, t] = order / (valid.sum() - 1)  # percentile rank

        # Rank IC per time step
        ics: List[float] = []
        for t in range(T):
            factor_col = ranked[:, t]
            ret_col = self.returns[:, t]
            both_valid = ~np.isnan(factor_col) & ~np.isnan(ret_col)
            if both_valid.sum() < 5:
                continue
            corr = np.corrcoef(factor_col[both_valid], ret_col[both_valid])[0, 1]
            if not np.isnan(corr):
                ics.append(corr)

        if not ics:
            return 0.0
        return float(np.mean(ics))
