from __future__ import annotations

from focus import FocusConfig
from growth import GrowthPolicy
from region import Region

from .experiment_config import InputSpaceConfig, RegionScaffold


def build_regions_from_scaffold(
    region_scaffolds: list[RegionScaffold],
    input_space: InputSpaceConfig,
    focus_config: FocusConfig | None = None,
) -> list[Region]:
    return [
        build_region_from_scaffold(region_scaffold, input_space, focus_config=focus_config)
        for region_scaffold in region_scaffolds
    ]


def build_region_from_scaffold(
    region_scaffold: RegionScaffold,
    input_space: InputSpaceConfig,
    focus_config: FocusConfig | None = None,
) -> Region:
    region_obj = Region(region_scaffold.name)
    region_obj.set_focus_config(focus_config)
    for layer_scaffold in region_scaffold.layers:
        layer_index = region_obj.add_layer(
            excitatory_count=layer_scaffold.excitatory_count,
            inhibitory_count=layer_scaffold.inhibitory_count,
            modulatory_count=layer_scaffold.modulatory_count,
        )
        if layer_scaffold.neuron_limit is not None:
            region_obj.get_layers()[layer_index].neuron_limit = int(layer_scaffold.neuron_limit)

    for connection_scaffold in region_scaffold.connections:
        region_obj.connect_layers(
            connection_scaffold.source_layer_index,
            connection_scaffold.dest_layer_index,
            probability=connection_scaffold.probability,
            feedback=connection_scaffold.feedback,
        )

    for port_name, layer_indices in region_scaffold.input_bindings.items():
        _bind_input_space(region_obj, input_space, port_name, list(layer_indices))

    if region_scaffold.growth_policy:
        region_obj.set_growth_policy(_build_growth_policy(region_scaffold.growth_policy))

    if region_scaffold.post_build_hook is not None:
        region_scaffold.post_build_hook(region_obj)
    return region_obj


def _bind_input_space(region_obj: Region, input_space: InputSpaceConfig, port_name: str, layer_indices: list[int]) -> None:
    if input_space.uses_scalar_port_inputs():
        region_obj.bind_input(port_name, layer_indices)
        return

    if input_space.rank == 2:
        region_obj.bind_input_2d(
            port_name,
            height=int(input_space.shape[0]),
            width=int(input_space.shape[1]),
            gain=float(input_space.gain),
            epsilon_fire=float(input_space.epsilon_fire),
            attach_layers=layer_indices,
        )
        return

    region_obj.bind_input_nd(
        port_name,
        shape=list(input_space.shape),
        gain=float(input_space.gain),
        epsilon_fire=float(input_space.epsilon_fire),
        layer_indices=layer_indices,
    )


def _build_growth_policy(growth_policy_payload: dict) -> GrowthPolicy:
    return GrowthPolicy(
        enable_layer_growth=bool(growth_policy_payload.get("enable_layer_growth", False)),
        max_total_layers=int(growth_policy_payload.get("max_total_layers", 8)),
        avg_slots_threshold=float(growth_policy_payload.get("avg_slots_threshold", 2.0)),
        percent_neurons_at_cap_threshold=float(
            growth_policy_payload.get("percent_neurons_at_cap_threshold", 100.0)
        ),
        layer_cooldown_ticks=int(growth_policy_payload.get("layer_cooldown_ticks", 0)),
        new_layer_excitatory_count=int(growth_policy_payload.get("new_layer_excitatory_count", 2)),
        wire_probability=float(growth_policy_payload.get("wire_probability", 1.0)),
    )
