from __future__ import annotations

from dataclasses import dataclass, field

from focus import FocusConfig
from growth import GrowthPolicy
from runner.experiment_config import ConnectionScaffold, LayerScaffold, RegionScaffold

SECOND_PORT_FALLBACK_DIGITS: dict[str, int] = {
    "Q": 1,
    "W": 2,
    "E": 3,
    "R": 4,
    "T": 5,
    "Y": 6,
    "U": 7,
    "I": 8,
    "O": 9,
    "P": 0,
}

@dataclass
class InteractiveProfile:
    name: str
    description: str
    rank: int
    input_ports: list[str]
    regions: list[RegionScaffold]
    focus_config: FocusConfig = field(default_factory=FocusConfig)
    default_normalize_inputs: str = "normalized"
    default_output: str = "compact"
    default_trace: bool = False
    second_port_fallback_keys: dict[str, int] = field(default_factory=lambda: dict(SECOND_PORT_FALLBACK_DIGITS))

    def supports_dual_input(self) -> bool:
        return len(self.input_ports) >= 2


def list_builtin_profiles() -> list[str]:
    return ["single_scalar", "dual_scalar", "growth_probe"]


def load_builtin_profile(profile_name: str) -> InteractiveProfile:
    normalized_profile_name = str(profile_name or "single_scalar").strip().lower()
    builders = {
        "single_scalar": build_single_scalar_profile,
        "dual_scalar": build_dual_scalar_profile,
        "growth_probe": build_growth_probe_profile,
    }
    if normalized_profile_name not in builders:
        raise ValueError(f"Unsupported interactive profile: {profile_name}")
    return builders[normalized_profile_name]()


def build_single_scalar_profile() -> InteractiveProfile:
    region_scaffold = RegionScaffold(
        name="main",
        layers=[
            LayerScaffold(excitatory_count=1, inhibitory_count=0, modulatory_count=0),
            LayerScaffold(excitatory_count=1, inhibitory_count=0, modulatory_count=0),
        ],
        input_bindings={"a": [0]},
        connections=[
            ConnectionScaffold(source_layer_index=0, dest_layer_index=1, probability=1.0, feedback=False),
        ],
    )
    return InteractiveProfile(
        name="single_scalar",
        description="One scalar input port driving a simple two-layer scalar scaffold.",
        rank=1,
        input_ports=["a"],
        regions=[region_scaffold],
        default_normalize_inputs="normalized",
        default_output="compact",
        default_trace=False,
    )


def build_dual_scalar_profile() -> InteractiveProfile:
    region_scaffold = RegionScaffold(
        name="main",
        layers=[
            LayerScaffold(excitatory_count=2, inhibitory_count=0, modulatory_count=0),
            LayerScaffold(excitatory_count=1, inhibitory_count=0, modulatory_count=0),
        ],
        input_bindings={"a": [0], "b": [0]},
        connections=[
            ConnectionScaffold(source_layer_index=0, dest_layer_index=1, probability=1.0, feedback=False),
        ],
    )
    return InteractiveProfile(
        name="dual_scalar",
        description="Two scalar input ports with retained state across ticks and a shared downstream path.",
        rank=1,
        input_ports=["a", "b"],
        regions=[region_scaffold],
        default_normalize_inputs="normalized",
        default_output="compact",
        default_trace=False,
    )


def build_growth_probe_profile() -> InteractiveProfile:
    region_scaffold = RegionScaffold(
        name="growth_probe",
        layers=[
            LayerScaffold(excitatory_count=1, inhibitory_count=0, modulatory_count=0, neuron_limit=3),
        ],
        input_bindings={"a": [0]},
        post_build_hook=configure_growth_probe_region,
    )
    return InteractiveProfile(
        name="growth_probe",
        description="A small scaffold biased toward visible slot, neuron, and spillover-layer growth.",
        rank=1,
        input_ports=["a"],
        regions=[region_scaffold],
        default_normalize_inputs="normalized",
        default_output="verbose",
        default_trace=False,
    )


def configure_growth_probe_region(region_obj) -> None:
    region_layers = region_obj.get_layers()
    if not region_layers:
        return
    seed_layer = region_layers[0]
    seed_layer.neuron_limit = 3
    for neuron_obj in seed_layer.get_neurons():
        neuron_obj.slot_cfg.slot_limit = 1
        neuron_obj.slot_limit = 1
        neuron_obj.slot_cfg.fallback_growth_threshold = 1
        neuron_obj.slot_cfg.neuron_growth_cooldown_ticks = 0
        neuron_obj.slot_cfg.growth_enabled = True
        neuron_obj.slot_cfg.neuron_growth_enabled = True
        neuron_obj.slot_cfg.layer_growth_enabled = True
    region_obj.set_growth_policy(
        GrowthPolicy(
            enable_layer_growth=True,
            max_total_layers=8,
            avg_slots_threshold=2.0,
            percent_neurons_at_cap_threshold=100.0,
            layer_cooldown_ticks=0,
            new_layer_excitatory_count=2,
            wire_probability=1.0,
        )
    )
