from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

from debug import trace as grownet_trace
from metrics import RegionMetrics

from .experiment_config import ExperimentConfig
from .input_loaders import expand_tick_payloads, flatten_tensor_payload, load_input_schedule
from .scaffold_builder import build_regions_from_scaffold


@dataclass(frozen=True)
class ExperimentRunSummary:
    delivered_events: int
    total_slots: int
    total_synapses: int
    layer_count: int
    neuron_count: int


def execute_experiment(
    experiment_config: ExperimentConfig,
    trace_config,
    writer,
    emit_plain_summary: bool,
) -> ExperimentRunSummary:
    loaded_schedule = load_input_schedule(experiment_config)
    tick_payloads = expand_tick_payloads(experiment_config, loaded_schedule, trace_config.ticks)
    regions = build_regions_from_scaffold(
        experiment_config.regions,
        experiment_config.input_space,
        focus_config=experiment_config.focus,
    )

    latest_metrics = RegionMetrics()
    for tick_ordinal, tick_payload in enumerate(tick_payloads, start=1):
        _begin_tick(experiment_config, tick_payload)
        latest_metrics = execute_tick(regions, experiment_config, tick_payload)
        grownet_trace.end_tick(latest_metrics)
        if emit_plain_summary and not trace_config.trace_enabled:
            writer.write(_format_summary_line(experiment_config.name, tick_ordinal, latest_metrics) + "\n")
            writer.flush()

    summary = ExperimentRunSummary(
        delivered_events=int(latest_metrics.delivered_events),
        total_slots=int(latest_metrics.total_slots),
        total_synapses=int(latest_metrics.total_synapses),
        layer_count=sum(len(region_obj.get_layers()) for region_obj in regions),
        neuron_count=sum(len(layer_obj.get_neurons()) for region_obj in regions for layer_obj in region_obj.get_layers()),
    )
    if not trace_config.benchmark_enabled:
        write_metrics_file(experiment_config, summary)
    return summary


def execute_tick(regions, experiment_config: ExperimentConfig, tick_payload) -> RegionMetrics:
    aggregate_metrics = RegionMetrics()
    for region_obj in regions:
        if experiment_config.input_space.uses_scalar_port_inputs():
            region_payload = {
                port_name: float(port_value)
                for port_name, port_value in dict(tick_payload).items()
                if port_name in region_obj.input_edges
            }
            region_metrics = region_obj.tick_inputs(region_payload)
        elif experiment_config.input_space.rank == 2:
            if experiment_config.focus.enabled:
                region_metrics = region_obj.tick_2d_focused(experiment_config.input_space.primary_port(), tick_payload)
            else:
                region_metrics = region_obj.tick_2d(experiment_config.input_space.primary_port(), tick_payload)
        else:
            flat_values, shape_values = flatten_tensor_payload(tick_payload)
            region_metrics = region_obj.tick_nd(
                experiment_config.input_space.primary_port(),
                flat_values,
                shape_values,
            )
        aggregate_metrics.inc_delivered_events(region_metrics.delivered_events)
        aggregate_metrics.add_slots(region_metrics.total_slots)
        aggregate_metrics.add_synapses(region_metrics.total_synapses)
    return aggregate_metrics


def write_metrics_file(experiment_config: ExperimentConfig, run_summary: ExperimentRunSummary) -> None:
    metrics_path = experiment_config.output.metrics_file
    if metrics_path is None:
        return
    metrics_path = _resolve_metrics_path(metrics_path)
    metrics_path.parent.mkdir(parents=True, exist_ok=True)
    metrics_payload = {
        "experiment": experiment_config.name,
        "delivered_events": run_summary.delivered_events,
        "total_slots": run_summary.total_slots,
        "total_synapses": run_summary.total_synapses,
        "layer_count": run_summary.layer_count,
        "neuron_count": run_summary.neuron_count,
    }
    with metrics_path.open("w", encoding="utf-8") as metrics_handle:
        json.dump(metrics_payload, metrics_handle, indent=2, sort_keys=True)
        metrics_handle.write("\n")


def _begin_tick(experiment_config: ExperimentConfig, tick_payload) -> None:
    if experiment_config.input_space.uses_scalar_port_inputs():
        if isinstance(tick_payload, dict) and len(tick_payload) == 1:
            port_name = next(iter(tick_payload.keys()))
            port_value = float(next(iter(tick_payload.values())))
            grownet_trace.begin_tick(experiment_config.name, port_name, port_value, use_default_signal_id=True)
            return
        grownet_trace.begin_tick(experiment_config.name, "inputs", None, use_default_signal_id=True)
        return

    grownet_trace.begin_tick(
        experiment_config.name,
        experiment_config.input_space.primary_port(),
        None,
        use_default_signal_id=True,
    )


def _format_summary_line(experiment_name: str, tick_ordinal: int, metrics: RegionMetrics) -> str:
    return (
        f"experiment={experiment_name} tick={tick_ordinal} "
        f"delivered_events={int(metrics.delivered_events)} "
        f"total_slots={int(metrics.total_slots)} "
        f"total_synapses={int(metrics.total_synapses)}"
    )


def _resolve_metrics_path(metrics_path: Path) -> Path:
    if (metrics_path.exists() and metrics_path.is_dir()) or metrics_path.suffix == "":
        return metrics_path / "metrics.json"
    return metrics_path
