Repository: CrepeGoat/FEHnt

Issue:
reverse orb iteration
Currently the program starts with a known orb count and iterates through states of *remaining orbs* down to zero.

The code can be rewritten to:
- iterate over orbs *spent*; start from 0 and iterate *upwards*
- aggregate results (%-chance of x targets pulled, (for finite targets) avg. orbs spent) on every orbs-spent milestone and yield them 

Discussion:

Patch:
diff --git a/fehnt/core.py b/fehnt/core.py
index 157f216..be7324e 100644
--- a/fehnt/core.py
+++ b/fehnt/core.py
@@ -64,7 +64,7 @@ class OutcomeCalculator:
                 alt_session.prob(self.event_details)
                 / session.prob(self.event_details)
             )
-            self.states[StateStruct(event, alt_session)] += prob * alt_subprob
+            self.push_state(event, alt_session, prob * alt_subprob)
 
             prob *= (1-alt_subprob)
 
@@ -82,15 +82,15 @@ class OutcomeCalculator:
             self.push_outcome(event, prob)
             return
 
-        self.states[StateStruct(event, SessionState(
+        self.push_state(event, SessionState(
             prob_tier=event.dry_streak // SUMMONS_PER_SESSION,
             stone_summons=sf.Series(0, index=tuple(Colors)),
             stone_presences=sf.Series(True, index=tuple(Colors)),
-        ))] += prob
+        ), prob)
 
     def branch_event(self, event, session, prob, stone_choice):
         """Split session into all potential following sub-sessions."""
-        orb_count = event.orb_count - stone_cost(session.stone_summons.sum()-1)
+        orbs_spent = event.orbs_spent + stone_cost(session.stone_summons.sum()-1)
 
         choice_starpool_probs = (self.event_details
                                  .pool_probs(session.prob_tier)
@@ -126,24 +126,34 @@ class OutcomeCalculator:
                          (event.targets_pulled, 1-prob_success))
 
             for targets_pulled, subsubprob in pulls:
-                new_event = EventState(orb_count, dry_streak, targets_pulled)
+                new_event = EventState(orbs_spent, dry_streak, targets_pulled)
+                new_prob = total_prob * subsubprob
 
-                self.states[StateStruct(new_event, session)] += (
-                    total_prob * subsubprob
-                )
+                self.push_state(new_event, session, new_prob)
+
+    def push_state(self, event, session, prob):
+        self.states[StateStruct(event, session)] += prob
+        self.push_outcome(event, prob)
+
+    def pull_outcome(self, event, prob):
+        """Remove a given probabilistic outcome to the recorded results."""
+        result = ResultState(event.orbs_spent, event.targets_pulled)
+        self.outcomes[result] -= prob
+        if self.outcomes[result] == Fraction:
+            del self.outcomes[result]
 
     def push_outcome(self, event, prob):
         """Add a given probabilistic outcome to the recorded results."""
-        result = ResultState(event.orb_count, event.targets_pulled)
+        result = ResultState(event.orbs_spent, event.targets_pulled)
         self.outcomes[result] += prob
 
-    def __call__(self, no_of_orbs):
-        """Calculate the summoning probabilities."""
+    def __iter__(self):
+        """Iterate the summoning probabilities at every orb milestone."""
         self.states = DefaultSortedDict()
-        self.outcomes = defaultdict(Fraction, [])
+        self.outcomes = DefaultSortedDict()
 
         self.init_new_session(
-            EventState(no_of_orbs, 0, 0*self.summoner.targets), Fraction(1)
+            EventState(0, 0, 0*self.summoner.targets), Fraction(1)
         )
 
         def iter_states():
@@ -154,17 +164,23 @@ class OutcomeCalculator:
         for i, ((event, session), prob) in enumerate(iter_states()):
             self.callback("  state no.:", i)
             self.callback("  no. of states in queue:", len(self.states))
-            self.callback('  orbs left:', event.orb_count)
+            self.callback("  no. of outcomes:", len(self.outcomes))
+            self.callback('  orbs spent:', event.orbs_spent)
 
             if session.stone_summons.sum() == SUMMONS_PER_SESSION:
                 self.callback('completed summoning session')
+                self.pull_outcome(event, prob)
                 self.init_new_session(event, prob)
                 continue
 
-            if event.orb_count < stone_cost(session.stone_summons.sum()):
-                self.callback('out of orbs')
-                self.push_outcome(event, prob)
-                continue
+            if self.outcomes and (
+                event.orbs_spent + stone_cost(session.stone_summons.sum())
+                > self.outcomes.peekitem(-1)[0].orbs_spent
+            ):
+                yield (self.outcomes.peekitem(-1)[0].orbs_spent, self.outcomes)
+
+            # Needs to pull *after* yielding outcomes
+            self.pull_outcome(event, prob)
 
             stone_choice_sequence = self.summoner.stone_choice_sequence(
                 event.targets_pulled,
@@ -177,7 +193,7 @@ class OutcomeCalculator:
                 event, session, prob, stone_choice_sequence
             )
 
-        return self.outcomes
+            assert sum(self.outcomes.values()) == 1
 
 
 def condense_results(results):
diff --git a/fehnt/core_utils.py b/fehnt/core_utils.py
index a52860c..18967b8 100644
--- a/fehnt/core_utils.py
+++ b/fehnt/core_utils.py
@@ -5,7 +5,7 @@ import static_frame as sf
 from fehnt.core_defs import Colors, stone_cost, SUMMONS_PER_SESSION
 
 
-EventState = namedtuple('EventState', 'orb_count dry_streak targets_pulled')
+EventState = namedtuple('EventState', 'orbs_spent dry_streak targets_pulled')
 
 
 class SessionState(
@@ -28,7 +28,6 @@ class SessionState(
         ].sum()
 
 
-
 class StateStruct(namedtuple('_', 'event session')):
     """
     Represents a unique state in summoning.
@@ -49,7 +48,7 @@ class StateStruct(namedtuple('_', 'event session')):
             next_stone_cost = stone_cost(self.session.stone_summons.sum())
 
         return (
-            self.event.orb_count - next_stone_cost,
+            -(self.event.orbs_spent + next_stone_cost),
             self.session.stone_presences.sum(),
         )
 
@@ -58,7 +57,7 @@ class StateStruct(namedtuple('_', 'event session')):
         return self._obj_func() < other._obj_func()
 
 
-ResultState = namedtuple('ResultState', 'orb_count targets_pulled')
+ResultState = namedtuple('ResultState', 'orbs_spent targets_pulled')
 
 
 def nCkarray(*k_values):
diff --git a/fehnt/runner.py b/fehnt/runner.py
index fd15112..ec5ece0 100644
--- a/fehnt/runner.py
+++ b/fehnt/runner.py
@@ -32,9 +32,12 @@ def run():
         (StarPools.x5_STAR_FOCUS, Colors.RED, 1),
     )
 
-    outcome_probs = OutcomeCalculator(
+    outcomes = OutcomeCalculator(
         event_details=EventDetails.make_standard(pool_counts),
         summoner=ColorHuntSummoner(target_pool_counts),
-    )(no_of_orbs=10)
+    )
+    for orbs_spent, outcome_probs in outcomes:
+        if orbs_spent == 10:
+            break
     for state, prob in condense_results(outcome_probs):
         print("{}: {:%}".format(state, float(prob)))


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
index c1b6b5d..7d727fe 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -34,31 +34,34 @@ def test_run():
         (StarPools.x5_STAR_FOCUS, Colors.RED, 1),
     )
 
-    outcome_probs = OutcomeCalculator(
+    outcomes = OutcomeCalculator(
         event_details=EventDetails.make_standard(pool_counts),
         summoner=ColorHuntSummoner(target_pool_counts),
-    )(no_of_orbs=10)
+    )
+    for orbs_spent, outcome_probs in outcomes:
+        if orbs_spent == 10:
+            break
 
     assert outcome_probs[5, make_pool_counts(
         (StarPools.x5_STAR_FOCUS, Colors.RED, 1)
     )] == Fraction(53331461890648098721, 2529964800000000000000)
 
-    assert outcome_probs[1, make_pool_counts(
+    assert outcome_probs[9, make_pool_counts(
         (StarPools.x5_STAR_FOCUS, Colors.RED, 1)
     )] == Fraction(1803428037031254581, 158122800000000000000)
 
-    assert outcome_probs[1, make_pool_counts(
+    assert outcome_probs[9, make_pool_counts(
         (StarPools.x5_STAR_FOCUS, Colors.RED, 0)
     )] == Fraction(8492342626380177821929, 19567696500000000000000)
 
-    assert outcome_probs[0, make_pool_counts(
+    assert outcome_probs[10, make_pool_counts(
         (StarPools.x5_STAR_FOCUS, Colors.RED, 1)
     )] == Fraction(
         35632905572024260525584546495453381641543009,
         3168357335173324800000000000000000000000000000
     )
 
-    assert outcome_probs[0, make_pool_counts(
+    assert outcome_probs[10, make_pool_counts(
         (StarPools.x5_STAR_FOCUS, Colors.RED, 0)
     )] == Fraction(
         1654738849167994100960528289363746618358456991,