add calculations that determine the percent gain and loss from using this algorithm

1 files changed, 40 insertions, 3 deletions

diff --git a/analysis.py b/analysis.py
index 2b5ed59..644fe93 100644
--- a/analysis.py
+++ b/analysis.py
@@ -68,21 +68,21 @@ def interpolate_intersection(intersection_indices, timestamps, prices1, prices2)
 """
 Returns the indices of where two arrays' values intersects
 """
-def find_intersections(prices1, prices2):
+def find_intersections(prices1, prices2, offset=0):
     if len(prices1) != len(prices2):
         print("ERROR IN find_intersections: len of arrs not the same")
         return []
     prev_p1 = prices1[0]
     prev_p2 = prices2[0]
     intersection_indices = set()
-    for i in range(1, len(prices1)):
+    for i in range(1 + offset, len(prices1)):
         next_p1 = prices1[i]
         next_p2 = prices2[i]
         # if the sign (negative to positive) changes, then there was an intersection between these pts
         sub_prev = prev_p1 - prev_p2
         sub_next = next_p1 - next_p2
 
-        if (sub_prev > 0 and sub_next < 0) or (sub_prev < 0 and sub_next > 0): # TODO, consider on the 0 case
+        if (sub_prev > 0 and sub_next < 0) or (sub_prev < 0 and sub_next > 0):
             intersection_indices.add((i-1, i))
         
         if sub_next == 0:
@@ -92,3 +92,40 @@ def find_intersections(prices1, prices2):
         prev_p2 = next_p2
     
     return intersection_indices
+
+def calculate_profit(buy_line, sell_line, prices, timestamps, offset=0, starting_money=10000):
+    if len(buy_line) != len(sell_line):
+        print("ERROR IN find_intersections: len of arrs not the same")
+        return []
+    is_bought = False
+    curr_money = 10000
+    shares_owned = 0
+    buy_info = []
+    sell_info = []
+    for i in range(offset, len(buy_line)):
+        current_b1 = buy_line[i]
+        current_sl = sell_line[i]
+        # if the sign is positive, we want to hold, if it's negative, we want to sell
+        sign_signal = current_b1 - current_sl
+
+        if sign_signal > 0:
+            if not is_bought:
+                # buy the stock
+                shares_owned = curr_money / prices[i]
+                curr_money = 0
+                buy_info.append((timestamps[i], prices[i], i))
+            is_bought = True
+        if sign_signal < 0:
+            if is_bought:
+                # selling the stock
+                curr_money = prices[i] * shares_owned
+                shares_owned = 0
+                sell_info.append((timestamps[i], prices[i], i))
+            is_bought = False
+
+    # TODO: consider end interval
+    total_assets = prices[-1] * shares_owned + curr_money
+    percent_gain = (total_assets - starting_money) / starting_money
+    return (percent_gain, total_assets, buy_info, sell_info)
+
+    
\ No newline at end of file