"""Genetic Programming implementation for OpenClaw trading system.

This module provides GP capabilities for evolving trading strategy trees
with various node types, operators, and tree manipulation functions.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Callable, Dict, Generic, List, Optional, Protocol, Tuple, TypeVar, Union

import numpy as np

from openclaw.utils.logging import get_logger


class NodeType(Enum):
    """Types of nodes in GP trees."""

    # Terminal nodes (leaves)
    PRICE = auto()
    VOLUME = auto()
    INDICATOR = auto()
    CONSTANT = auto()
    PARAMETER = auto()

    # Function nodes (internal)
    ADD = auto()
    SUB = auto()
    MUL = auto()
    DIV = auto()
    GT = auto()  # Greater than
    LT = auto()  # Less than
    EQ = auto()  # Equal
    AND = auto()
    OR = auto()
    NOT = auto()
    IF = auto()  # If-then-else
    MIN = auto()
    MAX = auto()
    ABS = auto()
    MOVING_AVG = auto()
    RSI = auto()
    MACD = auto()


# Arity for each node type
NODE_ARITIES: Dict[NodeType, int] = {
    # Terminals (0 arity)
    NodeType.PRICE: 0,
    NodeType.VOLUME: 0,
    NodeType.INDICATOR: 0,
    NodeType.CONSTANT: 0,
    NodeType.PARAMETER: 0,
    # Binary operators (2 arity)
    NodeType.ADD: 2,
    NodeType.SUB: 2,
    NodeType.MUL: 2,
    NodeType.DIV: 2,
    NodeType.GT: 2,
    NodeType.LT: 2,
    NodeType.EQ: 2,
    NodeType.AND: 2,
    NodeType.OR: 2,
    NodeType.MIN: 2,
    NodeType.MAX: 2,
    NodeType.MOVING_AVG: 2,  # period, value
    NodeType.MACD: 2,  # fast, slow
    # Unary operators (1 arity)
    NodeType.NOT: 1,
    NodeType.ABS: 1,
    NodeType.RSI: 1,  # period
    # Ternary operators (3 arity)
    NodeType.IF: 3,  # condition, then, else
}


@dataclass
class Node:
    """Node in a GP tree.

    Attributes:
        node_type: Type of node
        value: Value for terminal nodes (constant value, parameter name, etc.)
        children: Child nodes
        depth: Cached depth of this node
    """

    node_type: NodeType
    value: Any = None
    children: List["Node"] = field(default_factory=list)
    depth: int = 0

    def __post_init__(self) -> None:
        """Update depth after initialization."""
        self._update_depth()

    def _update_depth(self) -> None:
        """Recalculate depth."""
        if not self.children:
            self.depth = 0
        else:
            self.depth = 1 + max(c.depth for c in self.children)

    def copy(self) -> Node:
        """Create a deep copy of this node and all children."""
        return Node(
            node_type=self.node_type,
            value=self.value,
            children=[c.copy() for c in self.children],
            depth=self.depth,
        )

    def get_depth(self) -> int:
        """Get the depth of this node in the tree."""
        return self.depth

    def get_size(self) -> int:
        """Get the number of nodes in the subtree rooted at this node."""
        return 1 + sum(c.get_size() for c in self.children)

    def is_terminal(self) -> bool:
        """Check if this is a terminal node."""
        return NODE_ARITIES[self.node_type] == 0

    def is_function(self) -> bool:
        """Check if this is a function node."""
        return NODE_ARITIES[self.node_type] > 0

    def get_all_nodes(self) -> List[Node]:
        """Get all nodes in this subtree (breadth-first)."""
        nodes = [self]
        queue = [self]
        while queue:
            current = queue.pop(0)
            for child in current.children:
                nodes.append(child)
                queue.append(child)
        return nodes

    def get_node_at_index(self, index: int) -> Optional[Node]:
        """Get node at given index (breadth-first traversal)."""
        nodes = self.get_all_nodes()
        return nodes[index] if 0 <= index < len(nodes) else None

    def replace_node_at_index(self, index: int, new_node: Node) -> bool:
        """Replace node at given index with new node."""
        if index == 0:
            # Cannot replace root through this method
            return False

        nodes = self.get_all_nodes()
        if not (0 <= index < len(nodes)):
            return False

        target = nodes[index]

        # Find parent of target
        for node in nodes:
            if target in node.children:
                idx = node.children.index(target)
                node.children[idx] = new_node
                self._update_depth_recursive()
                return True

        return False

    def _update_depth_recursive(self) -> None:
        """Recursively update depths from this node."""
        self._update_depth()
        for child in self.children:
            child._update_depth_recursive()

    def evaluate(self, context: Dict[str, Any]) -> Any:
        """Evaluate this node with given context.

        Args:
            context: Dictionary containing market data and parameters

        Returns:
            Evaluation result
        """
        return self._evaluate_node(self, context)

    def _evaluate_node(self, node: Node, context: Dict[str, Any]) -> Any:
        """Recursively evaluate a node."""
        # Terminal nodes
        if node.node_type == NodeType.CONSTANT:
            return node.value

        elif node.node_type == NodeType.PRICE:
            price_type = node.value or "close"
            return context.get(f"price_{price_type}", 0.0)

        elif node.node_type == NodeType.VOLUME:
            return context.get("volume", 0.0)

        elif node.node_type == NodeType.INDICATOR:
            return context.get(f"indicator_{node.value}", 0.0)

        elif node.node_type == NodeType.PARAMETER:
            return context.get(f"param_{node.value}", 0.0)

        # Function nodes
        if node.node_type == NodeType.ADD:
            return self._evaluate_node(node.children[0], context) + self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.SUB:
            return self._evaluate_node(node.children[0], context) - self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.MUL:
            return self._evaluate_node(node.children[0], context) * self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.DIV:
            denom = self._evaluate_node(node.children[1], context)
            return self._evaluate_node(node.children[0], context) / denom if denom != 0 else 0.0

        elif node.node_type == NodeType.GT:
            return self._evaluate_node(node.children[0], context) > self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.LT:
            return self._evaluate_node(node.children[0], context) < self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.EQ:
            return abs(self._evaluate_node(node.children[0], context) - self._evaluate_node(node.children[1], context)) < 1e-10

        elif node.node_type == NodeType.AND:
            return self._evaluate_node(node.children[0], context) and self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.OR:
            return self._evaluate_node(node.children[0], context) or self._evaluate_node(node.children[1], context)

        elif node.node_type == NodeType.NOT:
            return not self._evaluate_node(node.children[0], context)

        elif node.node_type == NodeType.IF:
            if self._evaluate_node(node.children[0], context):
                return self._evaluate_node(node.children[1], context)
            else:
                return self._evaluate_node(node.children[2], context)

        elif node.node_type == NodeType.MIN:
            return min(self._evaluate_node(node.children[0], context), self._evaluate_node(node.children[1], context))

        elif node.node_type == NodeType.MAX:
            return max(self._evaluate_node(node.children[0], context), self._evaluate_node(node.children[1], context))

        elif node.node_type == NodeType.ABS:
            return abs(self._evaluate_node(node.children[0], context))

        elif node.node_type == NodeType.MOVING_AVG:
            period = int(self._evaluate_node(node.children[0], context))
            prices = context.get("price_history", [])
            if len(prices) >= period and period > 0:
                return np.mean(prices[-period:])
            return prices[-1] if prices else 0.0

        elif node.node_type == NodeType.RSI:
            period = int(self._evaluate_node(node.children[0], context))
            prices = context.get("price_history", [])
            if len(prices) < period or period < 2:
                return 50.0

            deltas = np.diff(prices[-period:])
            gains = np.where(deltas > 0, deltas, 0)
            losses = np.where(deltas < 0, -deltas, 0)

            avg_gain = np.mean(gains) if gains.size > 0 else 0
            avg_loss = np.mean(losses) if losses.size > 0 else 0

            if avg_loss == 0:
                return 100.0
            rs = avg_gain / avg_loss
            return 100.0 - (100.0 / (1.0 + rs))

        return 0.0

    def __repr__(self) -> str:
        """String representation of node."""
        if self.is_terminal():
            if self.value is not None:
                return f"{self.node_type.name}({self.value})"
            return self.node_type.name
        return f"{self.node_type.name}({len(self.children)})"


@dataclass
class TreeChromosome:
    """Chromosome for genetic programming.

    Attributes:
        root: Root node of the tree
        fitness: Fitness value
        generation: Generation when created
        max_depth: Maximum allowed depth
    """

    root: Node
    fitness: float = 0.0
    generation: int = 0
    max_depth: int = 10

    def copy(self) -> TreeChromosome:
        """Create a deep copy."""
        return TreeChromosome(
            root=self.root.copy(),
            fitness=self.fitness,
            generation=self.generation,
            max_depth=self.max_depth,
        )

    def get_depth(self) -> int:
        """Get depth of the tree."""
        return self.root.get_depth()

    def get_size(self) -> int:
        """Get number of nodes in the tree."""
        return self.root.get_size()

    def evaluate(self, context: Dict[str, Any]) -> Any:
        """Evaluate the tree."""
        return self.root.evaluate(context)


@dataclass
class GPConfig:
    """Configuration for Genetic Programming.

    Attributes:
        population_size: Number of trees in population
        max_generations: Maximum number of generations
        max_depth: Maximum tree depth
        min_depth: Minimum tree depth for initialization
        crossover_rate: Probability of crossover
        mutation_rate: Probability of mutation
        elitism_rate: Percentage of elite trees to preserve
        tournament_size: Size of tournament for selection
    """

    population_size: int = 100
    max_generations: int = 500
    max_depth: int = 10
    min_depth: int = 2
    crossover_rate: float = 0.9
    mutation_rate: float = 0.1
    elitism_rate: float = 0.05
    tournament_size: int = 5

    def __post_init__(self) -> None:
        """Validate configuration."""
        if self.population_size < 10:
            raise ValueError("Population size must be at least 10")
        if self.max_depth < self.min_depth:
            raise ValueError("Max depth must be >= min depth")


class GeneticProgramming:
    """Genetic Programming for evolving trading strategy trees.

    This class implements GP with subtree crossover, point mutation,
    and tournament selection for evolving tree-based strategies.

    Args:
        config: GP configuration
        fitness_func: Function to evaluate tree fitness
        terminal_set: List of terminal node types
        function_set: List of function node types
    """

    def __init__(
        self,
        config: GPConfig,
        fitness_func: Callable[[TreeChromosome], float],
        terminal_set: Optional[List[NodeType]] = None,
        function_set: Optional[List[NodeType]] = None,
    ):
        self.config = config
        self.fitness_func = fitness_func
        self.terminal_set = terminal_set or [
            NodeType.CONSTANT,
            NodeType.PRICE,
            NodeType.INDICATOR,
        ]
        self.function_set = function_set or [
            NodeType.ADD,
            NodeType.SUB,
            NodeType.MUL,
            NodeType.DIV,
            NodeType.GT,
            NodeType.LT,
            NodeType.IF,
            NodeType.MIN,
            NodeType.MAX,
        ]
        self.population: List[TreeChromosome] = []
        self.generation = 0
        self.best_chromosome: Optional[TreeChromosome] = None

        self.logger = get_logger("evolution.gp")
        self.logger.info(f"GeneticProgramming initialized: max_depth={config.max_depth}")

    def _create_random_node(self, depth: int, max_depth: int, force_terminal: bool = False) -> Node:
        """Create a random node with given constraints."""
        if force_terminal or depth >= max_depth:
            node_type = np.random.choice(self.terminal_set)
            if node_type == NodeType.CONSTANT:
                return Node(node_type=node_type, value=np.random.uniform(-10, 10))
            elif node_type == NodeType.PRICE:
                return Node(node_type=node_type, value=np.random.choice(["open", "high", "low", "close"]))
            elif node_type == NodeType.INDICATOR:
                return Node(node_type=node_type, value=np.random.choice(["rsi", "macd", "bb_upper", "bb_lower"]))
            return Node(node_type=node_type)
        else:
            node_type = np.random.choice(self.function_set)
            arity = NODE_ARITIES[node_type]
            children = [self._create_random_node(depth + 1, max_depth) for _ in range(arity)]
            return Node(node_type=node_type, children=children)

    def _grow_tree(self, max_depth: int) -> Node:
        """Grow a tree using the grow method."""
        return self._create_random_node(0, max_depth)

    def _full_tree(self, depth: int, max_depth: int) -> Node:
        """Grow a tree using the full method (all leaves at max depth)."""
        if depth >= max_depth:
            return self._create_random_node(depth, max_depth, force_terminal=True)

        node_type = np.random.choice(self.function_set)
        arity = NODE_ARITIES[node_type]
        children = [self._full_tree(depth + 1, max_depth) for _ in range(arity)]
        return Node(node_type=node_type, children=children)

    def _ramped_half_and_half(self) -> Node:
        """Initialize using ramped half-and-half method."""
        max_d = np.random.randint(self.config.min_depth, self.config.max_depth + 1)
        if np.random.random() < 0.5:
            return self._grow_tree(max_d)
        else:
            return self._full_tree(0, max_d)

    def initialize(self) -> None:
        """Initialize population with ramped half-and-half."""
        self.population = []
        for _ in range(self.config.population_size):
            root = self._ramped_half_and_half()
            chromosome = TreeChromosome(
                root=root,
                max_depth=self.config.max_depth,
                generation=0,
            )
            chromosome.fitness = self.fitness_func(chromosome)
            self.population.append(chromosome)

        self._sort_population()
        self.best_chromosome = self.population[0]
        self.logger.info(f"Population initialized with {len(self.population)} trees")

    def _sort_population(self) -> None:
        """Sort population by fitness."""
        self.population.sort(key=lambda x: x.fitness, reverse=True)

    def _tournament_select(self) -> TreeChromosome:
        """Select a chromosome using tournament selection."""
        contestants = np.random.choice(self.population, self.config.tournament_size, replace=False)
        return max(contestants, key=lambda x: x.fitness)

    def _subtree_crossover(
        self, parent1: TreeChromosome, parent2: TreeChromosome
    ) -> Tuple[TreeChromosome, TreeChromosome]:
        """Perform subtree crossover between two parents."""
        if np.random.random() > self.config.crossover_rate:
            return parent1.copy(), parent2.copy()

        child1 = parent1.copy()
        child2 = parent2.copy()

        # Get random crossover points
        nodes1 = child1.root.get_all_nodes()
        nodes2 = child2.root.get_all_nodes()

        if len(nodes1) < 2 or len(nodes2) < 2:
            return child1, child2

        # Select crossover points (excluding root)
        point1_idx = np.random.randint(1, len(nodes1))
        point2_idx = np.random.randint(1, len(nodes2))

        # Perform swap
        # Find parents of selected nodes
        for node1 in nodes1:
            if nodes1[point1_idx] in node1.children:
                idx1 = node1.children.index(nodes1[point1_idx])
                for node2 in nodes2:
                    if nodes2[point2_idx] in node2.children:
                        idx2 = node2.children.index(nodes2[point2_idx])

                        # Swap subtrees
                        temp = node1.children[idx1]
                        node1.children[idx1] = node2.children[idx2]
                        node2.children[idx2] = temp

                        # Update depths
                        child1.root._update_depth_recursive()
                        child2.root._update_depth_recursive()

                        # Check depth constraints
                        if child1.get_depth() > self.config.max_depth:
                            return parent1.copy(), parent2.copy()
                        if child2.get_depth() > self.config.max_depth:
                            return parent1.copy(), parent2.copy()

                        return child1, child2

        return child1, child2

    def _point_mutation(self, chromosome: TreeChromosome) -> TreeChromosome:
        """Perform point mutation on a chromosome."""
        if np.random.random() > self.config.mutation_rate:
            return chromosome.copy()

        mutated = chromosome.copy()
        nodes = mutated.root.get_all_nodes()

        if not nodes:
            return mutated

        # Select random node to mutate
        node_idx = np.random.randint(len(nodes))
        node = nodes[node_idx]

        if node.is_terminal():
            # Replace with new terminal
            node.node_type = np.random.choice(self.terminal_set)
            if node.node_type == NodeType.CONSTANT:
                node.value = np.random.uniform(-10, 10)
            elif node.node_type in (NodeType.PRICE, NodeType.INDICATOR):
                node.value = None
        else:
            # Replace with new function of same arity
            current_arity = NODE_ARITIES[node.node_type]
            compatible_functions = [
                nt for nt in self.function_set if NODE_ARITIES[nt] == current_arity
            ]
            if compatible_functions:
                node.node_type = np.random.choice(compatible_functions)

        mutated.root._update_depth_recursive()
        return mutated

    def _subtree_mutation(self, chromosome: TreeChromosome) -> TreeChromosome:
        """Perform subtree mutation."""
        if np.random.random() > self.config.mutation_rate:
            return chromosome.copy()

        mutated = chromosome.copy()
        nodes = mutated.root.get_all_nodes()

        if len(nodes) < 2:
            return mutated

        # Select random non-root node to replace
        node_idx = np.random.randint(1, len(nodes))

        # Create new subtree
        new_subtree = self._grow_tree(self.config.max_depth // 2)

        # Replace
        if mutated.root.replace_node_at_index(node_idx, new_subtree):
            mutated.root._update_depth_recursive()

        return mutated

    def _create_next_generation(self) -> None:
        """Create next generation."""
        new_population: List[TreeChromosome] = []

        # Elitism
        elite_count = max(1, int(self.config.elitism_rate * self.config.population_size))
        elite = [c.copy() for c in self.population[:elite_count]]
        new_population.extend(elite)

        # Create offspring
        while len(new_population) < self.config.population_size:
            parent1 = self._tournament_select()
            parent2 = self._tournament_select()

            child1, child2 = self._subtree_crossover(parent1, parent2)
            child1 = self._point_mutation(child1)
            child1.generation = self.generation
            child1.fitness = self.fitness_func(child1)

            child2 = self._point_mutation(child2)
            child2.generation = self.generation
            child2.fitness = self.fitness_func(child2)

            new_population.append(child1)
            if len(new_population) < self.config.population_size:
                new_population.append(child2)

        self.population = new_population
        self._sort_population()
        self.best_chromosome = self.population[0]

    def step(self) -> float:
        """Execute one generation step.

        Returns:
            Best fitness value
        """
        if not self.population:
            raise ValueError("Population not initialized. Call initialize() first.")

        self.generation += 1
        self._create_next_generation()

        return self.population[0].fitness

    def run(self, callback: Optional[Callable[[int, float], None]] = None) -> TreeChromosome:
        """Run the genetic programming algorithm.

        Args:
            callback: Optional callback (generation, best_fitness) -> None

        Returns:
            Best tree chromosome found
        """
        if not self.population:
            self.initialize()

        self.logger.info(f"GP started for {self.config.max_generations} generations")

        for gen in range(self.config.max_generations):
            best_fitness = self.step()

            if callback:
                callback(gen, best_fitness)

            if gen % 50 == 0 or gen < 5:
                avg_size = np.mean([c.get_size() for c in self.population])
                self.logger.info(
                    f"Generation {gen}: best={best_fitness:.4f}, avg_size={avg_size:.1f}"
                )

        self.logger.info(f"GP completed: best_fitness={self.best_chromosome.fitness:.4f}")
        return self.best_chromosome

    def simplify_tree(self, chromosome: TreeChromosome) -> TreeChromosome:
        """Simplify a tree by removing redundant nodes.

        Args:
            chromosome: Tree to simplify

        Returns:
            Simplified tree
        """
        simplified = chromosome.copy()
        self._simplify_node(simplified.root)
        simplified.root._update_depth_recursive()
        return simplified

    def _simplify_node(self, node: Node) -> None:
        """Recursively simplify a node."""
        # Simplify children first
        for child in node.children:
            self._simplify_node(child)

        # Simplify arithmetic with constants
        if node.node_type in (NodeType.ADD, NodeType.SUB, NodeType.MUL, NodeType.DIV):
            if all(c.node_type == NodeType.CONSTANT for c in node.children):
                # Both children are constants, compute result
                values = [c.value for c in node.children]
                result = 0.0
                if node.node_type == NodeType.ADD:
                    result = values[0] + values[1]
                elif node.node_type == NodeType.SUB:
                    result = values[0] - values[1]
                elif node.node_type == NodeType.MUL:
                    result = values[0] * values[1]
                elif node.node_type == NodeType.DIV and values[1] != 0:
                    result = values[0] / values[1]

                node.node_type = NodeType.CONSTANT
                node.value = result
                node.children = []

        # Remove redundant operations
        if node.node_type == NodeType.ADD:
            # x + 0 = x
            for i, child in enumerate(node.children):
                if child.node_type == NodeType.CONSTANT and abs(child.value) < 1e-10:
                    # Replace node with other child
                    other = node.children[1 - i]
                    node.node_type = other.node_type
                    node.value = other.value
                    node.children = other.children.copy()
                    break

    def get_statistics(self) -> Dict[str, Any]:
        """Get statistics about the population."""
        if not self.population:
            return {}

        sizes = [c.get_size() for c in self.population]
        depths = [c.get_depth() for c in self.population]
        fitnesses = [c.fitness for c in self.population]

        return {
            "generation": self.generation,
            "best_fitness": float(np.max(fitnesses)),
            "avg_fitness": float(np.mean(fitnesses)),
            "avg_tree_size": float(np.mean(sizes)),
            "max_tree_size": int(np.max(sizes)),
            "avg_tree_depth": float(np.mean(depths)),
            "max_tree_depth": int(np.max(depths)),
        }
