///- --package .
:# network static =
    :< testing

:= one () : i32 = 256
:= input_count () : i32 = 8
:= hidden_count () : i32 = 6
:= output_count () : i32 = 3
:= attacker_unit () : i32 = 85
:= weight_count () : i32 = input_count () * hidden_count () + hidden_count () * output_count ()

:! weights = array i32

:== sensed = { health : i32; attackers : i32; target : i32; distance : i32; heal : i32; strike : i32; flees : i32 }

:! inputs = sensed

:| feature =
    | HealthFraction
    | Attackers
    | TargetHealth
    | TargetDistance
    | HealReady
    | StrikeReady
    | TargetFleeing
    | Bias

:| output =
    | HealOutput
    | EngageOutput
    | RetreatOutput

:! hidden = i32

:= feature_slot (f : feature) : i32 =
    ?= f ?|
    | HealthFraction => 0
    | Attackers => 1
    | TargetHealth => 2
    | TargetDistance => 3
    | HealReady => 4
    | StrikeReady => 5
    | TargetFleeing => 6
    | Bias => 7

:= output_slot (o : output) : i32 =
    ?= o ?|
    | HealOutput => 0
    | EngageOutput => 1
    | RetreatOutput => 2

:= input_weight_slot (hidden h) (f : feature) : i32 = h * input_count () + feature_slot f

:= output_weight_slot (o : output) (hidden h) : i32 =
    input_count () * hidden_count () + output_slot o * hidden_count () + h

:== decision = { heal : bool; engage : bool; retreat : bool }

:= filled (count : i32) (value : i32) : array i32 =
    := x : array i32 = create count
    :~ ^ fill (i : i32) : () =
        ? i < count ?>
            set x i value
            fill (i + 1)
    fill 0
    x

:= zero_weights () : weights = weights (filled (weight_count ()) 0)

:= weight (weights w) (i : i32) : i32 = index w i

:= feature_flag (b : bool) : i32 = ? b ?> one () |: 0

:= sensed_distance (distance : i32) : i32 = (? distance > 7 ?> 7 |: distance) * 32

:= flags_of (heal : bool) (strike : bool) (flees : bool) : i32 * i32 * i32 =
    feature_flag heal, feature_flag strike, feature_flag flees

:= sensed_of (h : i32) (a : i32) (t : i32) (d : i32) ((x, y, z) : i32 * i32 * i32) : sensed =
    {
        health = h; attackers = a * attacker_unit (); target = t; distance = sensed_distance d
        heal = x; strike = y; flees = z
    }

:= inputs_of
    (health : i32) (attackers : i32) (target : i32) (distance : i32) (heal : bool) (strike : bool) (flees : bool)
    : inputs
    =
    inputs (sensed_of health attackers target distance (flags_of heal strike flees))

:= term (w : weights) (j : i32) (i : i32) (value : i32) : i32 = weight w (j * input_count () + i) * value

:= sensed_terms (w : weights) ((inputs s) : inputs) (j : i32) : i32 =
    term w j 0 s::health + term w j 1 s::attackers + term w j 2 s::target + term w j 3 s::distance

:= flag_terms (w : weights) ((inputs s) : inputs) (j : i32) : i32 =
    term w j 4 s::heal + term w j 5 s::strike + term w j 6 s::flees + term w j 7 (one ())

:= hidden_unit (w : weights) (x : inputs) (j : i32) : i32 =
    := h = (sensed_terms w x j + flag_terms w x j) / one ()
    ? h > 0 ?> h |: 0

:= hidden_layer (w : weights) (x : inputs) : array i32 =
    := hidden = filled (hidden_count ()) 0
    :~ ^ activate (j : i32) : () =
        ? j < hidden_count () ?>
            set hidden j (hidden_unit w x j)
            activate (j + 1)
    activate 0
    hidden

:= output_weight_at (w : weights) (k : i32) (j : i32) : i32 =
    weight w (input_count () * hidden_count () + k * hidden_count () + j)

:~ ^ output_sum (w : weights) (hidden : array i32) (k : i32) (j : i32) (total : i32) : i32 =
    ? j >= hidden_count () ?> total |: output_sum w hidden k (j + 1) (total + output_weight_at w k j * index hidden j)

:= output_unit (w : weights) (hidden : array i32) (k : i32) : i32 = output_sum w hidden k 0 0 / one ()

:= decision_of (heal : i32) (engage : i32) (retreat : i32) : decision =
    { heal = heal > 0; engage = engage > 0 && engage >= retreat; retreat = retreat > 0 && retreat > engage }

:= scaled (heal : i32) (engage : i32) (retreat : i32) : decision =
    decision_of (heal / one ()) (engage / one ()) (retreat / one ())

:= added (w : weights) (j : i32) (a : i32) (k : i32) (acc : i32) : i32 = acc + output_weight_at w k j * a

:~ ^ outputs_from (w : weights) (x : inputs) (j : i32) (h : i32) (e : i32) (r : i32) : decision =
    ? j >= hidden_count () ?> scaled h e r
    |:
        := a = hidden_unit w x j
        outputs_from w x (j + 1) (added w j a 0 h) (added w j a 1 e) (added w j a 2 r)

:= decide (w : weights) (x : inputs) : decision = outputs_from w x 0 0 0 0

:= with_weight (i : i32) (value : i32) (weights w) : weights =
    set w i value
    weights w

:= with_input_weight (h : hidden) (f : feature) (value : i32) (w : weights) : weights =
    with_weight (input_weight_slot h f) value w

:= with_output_weight (o : output) (h : hidden) (value : i32) (w : weights) : weights =
    with_weight (output_weight_slot o h) value w

:= fraction_below (heal_below : i32) (max_health : i32) : i32 = (heal_below - 1) * one () / max_health + 1

:= with_heal_threshold (heal_below_fraction : i32) (w : weights) : weights =
    w
    |> with_input_weight (hidden 0) HealthFraction (0 - one ())
    |> with_input_weight (hidden 0) Bias heal_below_fraction
    |> with_output_weight HealOutput (hidden 0) (one ())

:= with_engaging (w : weights) : weights =
    w |> with_input_weight (hidden 1) Bias (one ()) |> with_output_weight EngageOutput (hidden 1) (one ())

:= threshold_healer (heal_below_fraction : i32) : weights =
    zero_weights () |> with_heal_threshold heal_below_fraction |> with_engaging

:= with_kite_trigger (w : weights) : weights =
    w |> with_input_weight (hidden 2) Attackers (one ()) |> with_input_weight (hidden 2) StrikeReady (0 - one ())

:= with_kiting (min_attackers : i32) (w : weights) : weights =
    w
    |> with_kite_trigger
    |> with_input_weight (hidden 2) Bias (0 - (min_attackers - 1) * attacker_unit ())
    |> with_output_weight RetreatOutput (hidden 2) (4 * one ())

:= clamp_weight (x : i32) : i32 = ? x < -2048 ?> -2048 |? x > 2048 ?> 2048 |: x

:== weight_search = { seed : u32; generations : i32; offspring : i32; spread : i32 }

:== scored = { w : weights; score : i32 }

:== weight_round = { s : weight_search; index : i32; width : i32; parent : scored }

:< keyed_rng

:| network_roll =
    | WeightNudge : i32

:+ purpose_code network_roll =
    \ r =>
        ?= r ?|
        | WeightNudge i => conv (1000 + i)

:= weight_key (r : weight_round) (child : i32) (i : i32) : keyed_rng::roll_key network_roll =
    := s = keyed_rng::seed r::s::seed
    keyed_rng::roll_key s (keyed_rng::tick (conv r::index)) (keyed_rng::entity (conv child)) (WeightNudge i)

:= weight_draw (r : weight_round) (child : i32) (i : i32) : i32 =
    conv (keyed_rng::roll_between 0 (conv (2 * r::width)) (weight_key r child i))

:~ ^ copy_mutated (r : weight_round) (child : i32) (next : array i32) (i : i32) : () =
    ? i < weight_count () ?>
        set next i (clamp_weight (weight r::parent::w i + weight_draw r child i - r::width))
        copy_mutated r child next (i + 1)

:= mutate (r : weight_round) (child : i32) : weights =
    := next = filled (weight_count ()) 0
    copy_mutated r child next 0
    weights next

:= fitter_weights (fitness : weights -> i32) (candidate : weights) (best : scored) : scored =
    := score = fitness candidate
    ? score > best::score ?> { w = candidate; score = score } |: best

:~ ^ offspring_from (fitness : weights -> i32) (r : weight_round) (child : i32) (best : scored) : scored =
    ? child >= r::s::offspring ?> best
    |: offspring_from fitness r (child + 1) (fitter_weights fitness (mutate r child) best)

:= next_weight_round (r : weight_round) (best : scored) : weight_round =
    { r ?| index = r::index + 1; width = (? r::width > 4 ?> r::width * 4 / 5 |: 4); parent = best }

:~ ^ weight_climb (fitness : weights -> i32) (r : weight_round) : weights =
    ? r::index >= r::s::generations ?> r::parent::w
    |: weight_climb fitness (next_weight_round r (offspring_from fitness r 0 r::parent))

:= evolve (fitness : weights -> i32) (s : weight_search) (start : weights) : weights =
    weight_climb fitness { s = s; index = 0; width = s::spread; parent = { w = start; score = fitness start } }

:? judge f : f -> weights -> i32

:= fitter_by \* f {judge}. (j : f) (candidate : weights) (best : scored) : scored =
    := score = judge j candidate
    ? score > best::score ?> { w = candidate; score = score } |: best

:~ ^ offspring_by \* f {judge}. (j : f) (r : weight_round) (child : i32) (best : scored) : scored =
    ? child >= r::s::offspring ?> best |: offspring_by j r (child + 1) (fitter_by j (mutate r child) best)

:~ ^ climb_by \* f {judge}. (j : f) (r : weight_round) : weights =
    ? r::index >= r::s::generations ?> r::parent::w |: climb_by j (next_weight_round r (offspring_by j r 0 r::parent))

:= evolve_by \* f {judge}. (j : f) (s : weight_search) (start : weights) : weights =
    climb_by j { s = s; index = 0; width = s::spread; parent = { w = start; score = judge j start } }

///> c
:# a_zero_network_holds =
    := ~health = 128i32
    := d = decide (zero_weights ()) (inputs_of health 1 200 1 !+ !+ !-)
    d::heal |> _assert_eq !-
    d::engage |> _assert_eq !-
    d::retreat |> _assert_eq !-

:= hand_wired_healer () : weights =
    zero_weights ()
    |> with_input_weight (hidden 0) HealthFraction (-256)
    |> with_input_weight (hidden 0) Bias 128
    |> with_output_weight HealOutput (hidden 0) 256

:= heals_at (w : weights) (health : i32) : bool = (decide w (inputs_of health 0 0 0 !+ !+ !-))::heal

///> c
:# a_hand_wired_network_heals_when_low =
    := ~health = 64i32
    := w = hand_wired_healer ()
    (heals_at w health, heals_at w (health + 100)) |> _assert_eq (!+, !-)

:= threshold_at (health : i32) (attackers : i32) : decision =
    decide (threshold_healer 100) (inputs_of health attackers 0 1 !+ !+ !-)

///> c
:# a_threshold_healer_heals_below_and_engages_above =
    := ~health = 99i32
    := above = threshold_at (health + 1) 0
    ((threshold_at health 0)::heal, above::heal, above::engage, (threshold_at (health + 1) 2)::retreat)
    |> _assert_eq (!+, !-, !+, !-)

:= kiter_at (health : i32) (attackers : i32) (strike_ready : bool) : decision =
    decide (threshold_healer 100 |> with_kiting 2) (inputs_of health attackers 0 1 !+ strike_ready !-)

///> c
:# kiting_retreats_only_while_the_strike_recovers =
    := ~health = 400i32
    := ready = kiter_at health 2 !+
    ((kiter_at health 2 !-)::retreat, ready::retreat, (kiter_at health 1 !-)::retreat, ready::engage)
    |> _assert_eq (!+, !-, !-, !+)

///> c
:# an_output_is_not_a_feature fails "Unification failure" =
    zero_weights () |> with_input_weight (hidden 0) HealOutput 1

///> c
:# inputs_are_not_weights fails "Unification failure" =
    decide (inputs_of 1 0 0 0 !+ !+ !-) (inputs_of 1 0 0 0 !+ !+ !-)

:~ ^ heal_score (w : weights) (h : i32) (total : i32) : i32 =
    ? h > 256 ?> total |: heal_score w (h + 16) (? (h < 64) = heals_at w h ?> total + 1 |: total)

:= quarter_healer_search (seed : u32) : weight_search = { seed = seed; generations = 30; offspring = 12; spread = 256 }

///> c
:# evolution_learns_to_heal_below_a_quarter =
    := ~seed = 5u32
    := score w = heal_score w 0 0
    score (evolve score (quarter_healer_search seed) (zero_weights ())) |> _assert_eq 17i32

///> wgsl
:# network_decides_on_gpu =
    := ~health = 64i32
    := w = hand_wired_healer ()
    ? heals_at w health && heals_at w (health + 100) = !- ?> 0i32 |: 1i32