"""
Run a quick Omniglot stream through a single GrowNet layer.
No heavy frameworks: only numpy, pillow, requests, tqdm.
"""
import argparse
import time

import numpy as np
from tqdm import tqdm

from omniglot_numpy import omniglot_batches
from layer import Layer


def collect_layer_snapshot(layer: Layer) -> dict[str, float | int]:
    neurons = layer.get_neurons()
    all_weights = [weight_obj for neuron_obj in neurons for weight_obj in neuron_obj.slots.values()]
    total_slots = sum(len(neuron_obj.slots) for neuron_obj in neurons)
    total_synapses = sum(len(neuron_obj.get_outgoing()) for neuron_obj in neurons)
    fired_neurons = sum(1 for neuron_obj in neurons if bool(getattr(neuron_obj, "fired_last", False)))
    average_ema = float(np.mean([weight_obj.ema_rate for weight_obj in all_weights])) if all_weights else 0.0
    total_strength = float(sum(weight_obj.get_strength_value() for weight_obj in all_weights)) if all_weights else 0.0

    return {
        "fired_neurons": int(fired_neurons),
        "total_slots": int(total_slots),
        "total_synapses": int(total_synapses),
        "average_ema": float(average_ema),
        "total_strength": float(total_strength),
    }


def run_few_shot(
    shots: int = 5,
    num_batches: int = 100,
    num_excit: int = 50,
    num_inhib: int = 0,
    num_mod: int = 0,
    p_feedforward: float = 0.1,
    p_feedback: float = 0.01,
    log_every: int = 0,
    batch_stream=None,
):
    if batch_stream is None:
        batch_stream = omniglot_batches(shots, batch_classes=5)

    layer = Layer(
        excitatory_count=num_excit,
        inhibitory_count=num_inhib,
        modulatory_count=num_mod,
    )
    layer.wire_random_feedforward(probability=p_feedforward)
    layer.wire_random_feedback(probability=p_feedback)

    start_time = time.perf_counter()

    for batch_index in tqdm(range(num_batches), desc="Omniglot episodes"):
        images, _ = next(batch_stream)
        for pixel_value in images.ravel():
            layer.forward(float(pixel_value))

        if log_every > 0 and (batch_index + 1) % log_every == 0:
            layer_snapshot = collect_layer_snapshot(layer)
            print(
                f"[batch {batch_index + 1}] "
                f"fired={layer_snapshot['fired_neurons']} "
                f"slots={layer_snapshot['total_slots']} "
                f"synapses={layer_snapshot['total_synapses']} "
                f"ema={layer_snapshot['average_ema']:.3f} "
                f"strength_sum={layer_snapshot['total_strength']:.3f}"
            )


    duration = time.perf_counter() - start_time

    final_snapshot = collect_layer_snapshot(layer)
    summary_data = {
        "shots": int(shots),
        "num_batches": int(num_batches),
        "num_excit": int(num_excit),
        "num_inhib": int(num_inhib),
        "num_mod": int(num_mod),
        "p_feedforward": float(p_feedforward),
        "p_feedback": float(p_feedback),
        "total_slots": int(final_snapshot["total_slots"]),
        "total_synapses": int(final_snapshot["total_synapses"]),
        "average_ema": float(final_snapshot["average_ema"]),
        "fired_neurons_last_batch": int(final_snapshot["fired_neurons"]),
        "elapsed_seconds": float(duration),
    }

    print("\n=== Run summary ===")
    print(f"Shots per class          : {summary_data['shots']}")
    print(f"E/I/M neurons            : {summary_data['num_excit']}/{summary_data['num_inhib']}/{summary_data['num_mod']}")
    print(f"Feedforward / Feedback p : {summary_data['p_feedforward']:.3f} / {summary_data['p_feedback']:.3f}")
    print(f"Total slots              : {summary_data['total_slots']}")
    print(f"Total synapses           : {summary_data['total_synapses']}")
    print(f"Average firing EMA       : {summary_data['average_ema']:.4f}")
    print(f"Elapsed time             : {summary_data['elapsed_seconds']:.1f} s")

    return summary_data


def main() -> None:
    argument_parser = argparse.ArgumentParser()
    argument_parser.add_argument("--shots", type=int, default=5)
    argument_parser.add_argument("--batches", type=int, default=100)
    argument_parser.add_argument("--excit", type=int, default=50)
    argument_parser.add_argument("--inhib", type=int, default=0)
    argument_parser.add_argument("--mod", type=int, default=0)
    argument_parser.add_argument("--p_feedforward", type=float, default=0.10)
    argument_parser.add_argument("--p_feedback", type=float, default=0.01)
    argument_parser.add_argument("--log_every", type=int, default=0, help="Log layer state every N batches (0 = off)")

    cli_args = argument_parser.parse_args()

    run_few_shot(
        shots=cli_args.shots,
        num_batches=cli_args.batches,
        num_excit=cli_args.excit,
        num_inhib=cli_args.inhib,
        num_mod=cli_args.mod,
        p_feedforward=cli_args.p_feedforward,
        p_feedback=cli_args.p_feedback,
        log_every=cli_args.log_every,
    )

if __name__ == "__main__":
    main()
