from world.base.neural import FlexibleNeuralNetwork
from config.constants import MAX_VELOCITY, MAX_ROTATIONAL_VELOCITY
from world.behavioral import BehavioralModel


class CellBrain(BehavioralModel):
    """
    Enhanced CellBrain using a flexible neural network with input normalization.
    """

    def __init__(self, neural_network=None, input_ranges=None):
        super().__init__()

        # Define input and output keys
        self.input_keys = ['distance', 'angle', 'current_speed', 'current_angular_velocity']
        self.output_keys = ['linear_acceleration', 'angular_acceleration']

        # Initialize inputs and outputs
        self.inputs = {key: 0.0 for key in self.input_keys}
        self.outputs = {key: 0.0 for key in self.output_keys}

        # Set input ranges for normalization
        default_ranges = {
            'distance': (0, 50),
            'angle': (-180, 180),
            'current_speed': (-MAX_VELOCITY, MAX_VELOCITY),
            'current_angular_velocity': (-MAX_ROTATIONAL_VELOCITY, MAX_ROTATIONAL_VELOCITY)
        }
        self.input_ranges = input_ranges if input_ranges is not None else default_ranges

        # Use provided network or create new one
        if neural_network is None:
            self.neural_network = FlexibleNeuralNetwork(
                input_size=len(self.input_keys),
                output_size=len(self.output_keys)
            )
        else:
            self.neural_network = neural_network

    def _normalize_input(self, key, value):
        min_val, max_val = self.input_ranges.get(key, (0.0, 1.0))
        # Avoid division by zero
        if max_val == min_val:
            return 0.0
        # Normalize to [-1, 1]
        return 2 * (value - min_val) / (max_val - min_val) - 1

    def tick(self, input_data) -> dict:
        """
        Process inputs through neural network and produce outputs.

        :param input_data: Dictionary containing input values
        :return: Dictionary with output values
        """
        # Update internal input state
        for key in self.input_keys:
            self.inputs[key] = input_data.get(key, 0.0)

        # Normalize inputs
        input_array = [self._normalize_input(key, self.inputs[key]) for key in self.input_keys]

        # Process through neural network
        output_array = self.neural_network.forward(input_array)

        # Map outputs back to dictionary
        self.outputs = {
            key: output_array[i] if i < len(output_array) else 0.0
            for i, key in enumerate(self.output_keys)
        }

        return self.outputs.copy()

    def mutate(self, mutation_rate=0.1):
        """
        Create a mutated copy of this CellBrain.

        :param mutation_rate: Rate of mutation for the neural network
        :return: New CellBrain with mutated neural network
        """
        mutated_network = self.neural_network.mutate(mutation_rate)
        return CellBrain(neural_network=mutated_network, input_ranges=self.input_ranges.copy())

    def get_network_info(self):
        """Get information about the underlying neural network."""
        return self.neural_network.get_structure_info()

    def __repr__(self):
        inputs = {key: round(value, 5) for key, value in self.inputs.items()}
        outputs = {key: round(value, 5) for key, value in self.outputs.items()}
        network_info = self.get_network_info()

        return (f"CellBrain(inputs={inputs}, outputs={outputs}, "
                f"network_layers={network_info['layer_sizes']}, "
                f"connections={network_info['total_connections']})")