3 files changed, 119 insertions, 13 deletions

diff --git a/src/balancer/blc_damage_type.erl b/src/balancer/blc_damage_type.erl
index 340b4ff..feac276 100644
--- a/src/balancer/blc_damage_type.erl
+++ b/src/balancer/blc_damage_type.erl
@@ -99,7 +99,11 @@ generate_entries_from_score (TargetScore, SortedRatios) ->
          SortedRatios
       ),
 
-   Base = (TargetScore / (Distribution / 100)),
+   Base =
+      case Distribution of
+         0 -> 0;
+         _ -> (TargetScore / (Distribution / 100))
+      end,
 
    UnderperformingEntries =
       lists:map
@@ -110,8 +114,19 @@ generate_entries_from_score (TargetScore, SortedRatios) ->
 
    MissingScore = (TargetScore - compute_score(UnderperformingEntries)),
 
-   case (MissingScore >= 0) of
-      true -> apply_score_modifier(MissingScore, 1, UnderperformingEntries);
-      false ->
-         apply_score_modifier((-1 * MissingScore), -1, UnderperformingEntries)
+   case Base of
+      0 -> [];
+      _ ->
+         case (MissingScore >= 0) of
+            true ->
+               apply_score_modifier(MissingScore, 1, UnderperformingEntries);
+
+            false ->
+               apply_score_modifier
+               (
+                  (-1 * MissingScore),
+                  -1,
+                  UnderperformingEntries
+               )
+         end
    end.
diff --git a/src/balancer/blc_distribution.erl b/src/balancer/blc_distribution.erl
index 2e8a141..a305dd4 100644
--- a/src/balancer/blc_distribution.erl
+++ b/src/balancer/blc_distribution.erl
@@ -10,7 +10,8 @@
 -export
 (
    [
-      generate/2
+      generate/2,
+      generate/3
    ]
 ).
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
@@ -50,7 +51,12 @@ generate_internals (N, CurrentResult, Sequence) ->
 -spec generate (non_neg_integer(), 0..100) -> list(list(0..100)).
 generate (0, _Step) -> [];
 generate (Elements, Step) ->
-   Sequence = lists:seq(0, 100, Step),
+   generate(Elements, 0, Step).
+
+-spec generate (non_neg_integer(), 0..100, 0..100) -> list(list(0..100)).
+generate (0, _Min, _Step) -> [];
+generate (Elements, Min, Step) ->
+   Sequence = lists:seq(Min, 100, Step),
    generate_internals
    (
       (Elements - 1),
diff --git a/src/balancer/struct/blc_armor.erl b/src/balancer/struct/blc_armor.erl
index 0c642e1..4d94a88 100644
--- a/src/balancer/struct/blc_armor.erl
+++ b/src/balancer/struct/blc_armor.erl
@@ -77,7 +77,9 @@
    [
       new/1,
       get_remaining_points/1,
-      generate/1
+      generate/3,
+      finalize/1,
+      export/1
    ]
 ).
 
@@ -327,15 +329,15 @@ increase_defense_score_for (GivenPoints, Armor) ->
 -spec get_remaining_points (type()) -> non_neg_integer().
 get_remaining_points (Armor) -> Armor#proto_armor.remaining_points.
 
--spec generate (non_neg_integer()) -> list(type()).
-generate (Step) ->
+-spec generate (0..100, 0..100, 0..100) -> list(type()).
+generate (AttributeMin, AttributeStep, ElementStep) ->
    MaxDamageModifier = get_max_attribute_ratio(?ATTRIBUTE_DAMAGE_MODIFIER),
    MaxDodgeChance = get_max_attribute_ratio(?ATTRIBUTE_DODGE_CHANCE),
    MaxMovementPoints = get_max_attribute_ratio(?ATTRIBUTE_MOVEMENT_POINTS),
    MaxHealth = get_max_attribute_ratio(?ATTRIBUTE_HEALTH),
    MaxDefenseScore = get_max_attribute_ratio(?ATTRIBUTE_DEFENSE_SCORE),
 
-   Distributions = blc_distribution:generate(5, Step),
+   Distributions = blc_distribution:generate(5, AttributeMin, AttributeStep),
    ValidDistributions =
       lists:filtermap
       (
@@ -439,9 +441,92 @@ generate (Step) ->
                   PointsUsed
                ),
 
-            FinalArmor
+            lists:foldl
+            (
+               fun ({Attribute, _}, Armor) ->
+                  NewArmor =
+                     case Attribute of
+                        ?ATTRIBUTE_DEFENSE_SCORE ->
+                           increase_defense_score_for
+                           (
+                              Armor#proto_armor.remaining_points,
+                              Armor
+                           );
+
+                        _ -> increase_attribute_by(Attribute, 1, Armor)
+                     end,
+
+                  case NewArmor of
+                     {ok, NextArmor} -> NextArmor;
+                     _ -> Armor
+                  end
+               end,
+               FinalArmor,
+               lists:sort
+               (
+                  fun ({_AttributeA, ScoreA}, {_AttributeB, ScoreB}) ->
+                     (ScoreA > ScoreB)
+                  end,
+                  lists:map
+                  (
+                     fun (Attribute) ->
+                        case Attribute of
+                           ?ATTRIBUTE_DEFENSE_SCORE ->
+                              {
+                                 ?ATTRIBUTE_DEFENSE_SCORE,
+                                 FinalArmor#proto_armor.defense_score
+                              };
+
+                           _ ->
+                              {
+                                 Attribute,
+                                 shr_omnimods:get_attribute_modifier
+                                 (
+                                    Attribute,
+                                    FinalArmor#proto_armor.omnimods
+                                 )
+                              }
+                        end
+                     end,
+                     [
+                        ?ATTRIBUTE_DEFENSE_SCORE,
+                        ?ATTRIBUTE_DODGE_CHANCE,
+                        ?ATTRIBUTE_HEALTH,
+                        ?ATTRIBUTE_DAMAGE_MODIFIER,
+                        ?ATTRIBUTE_MOVEMENT_POINTS
+                     ]
+                  )
+               )
+            )
          end,
          ValidDistributions
       ),
 
-   BaseArmors.
+   shr_lists_util:product
+   (
+      fun (Armor, ElementDistribution) ->
+         set_defense_coefficients(ElementDistribution, Armor)
+      end,
+      BaseArmors,
+      lists:map
+      (
+         fun ([A, B, C]) ->
+            [
+               {?DAMAGE_TYPE_SLASH, A},
+               {?DAMAGE_TYPE_PIERCE, B},
+               {?DAMAGE_TYPE_BLUNT, C}
+            ]
+         end,
+         blc_distribution:generate(3, ElementStep)
+      )
+   ).
+
+-spec export (type()) -> list().
+export (Armor) -> shr_omnimods:export(Armor#proto_armor.omnimods).
+
+-spec finalize (type()) -> {shr_omnimods:type(), non_neg_integer()}.
+finalize (Armor) ->
+   {
+      Armor#proto_armor.omnimods,
+      Armor#proto_armor.remaining_points
+   }.