diff --git a/backtester/backtester.py b/backtester/backtester.py
index a3c43a6..82016b5 100644
--- a/backtester/backtester.py
+++ b/backtester/backtester.py
@@ -1,10 +1,8 @@
 import pandas as pd
 import numpy as np
 import pyprind
-import seaborn as sns
-import matplotlib.pyplot as plt
 
-from .strategy import Strategy
+from .strategy import Strategy, Order
 from .datahandler import HistoricalOptionsData
 
 
@@ -105,18 +103,18 @@ class Backtest:
                ('totals',
                 'capital')] = (-df['totals']['cost'] * df['totals']['qty']).cumsum() + self._strategy.initial_capital
 
-        daily_df = df.groupby(('totals', 'date'))
-        daily_capital = daily_df.apply(lambda row: row['totals']['capital'].tail(1))
-        daily_returns = daily_capital.pct_change() * 100
+        daily_returns = df.groupby(('totals', 'date')).last()['totals']['capital'].pct_change() * 100
 
-        entries_mask = df.apply(lambda row: row['leg_1']['order'][2] == 'O', axis=1)
-        entries = df.loc[entries_mask]
-        exits = df.loc[~entries_mask]
+        first_leg = self._strategy.legs[0].name
+
+        entry_mask = df[first_leg].eval('(order == @Order.BTO) | (order == @Order.STO)')
+        entries = df.loc[entry_mask]
+        exits = df.loc[~entry_mask]
 
         costs = np.array([])
-        for contract in entries['leg_1']['contract']:
-            entry = entries.loc[entries['leg_1']['contract'] == contract]
-            exit_ = exits.loc[exits['leg_1']['contract'] == contract]
+        for contract in entries[first_leg]['contract']:
+            entry = entries.loc[entries[first_leg]['contract'] == contract]
+            exit_ = exits.loc[exits[first_leg]['contract'] == contract]
             try:
                 # Here we assume we are entering only once per contract (i.e both entry and exit_ have only one row)
                 costs = np.append(costs, (entry['totals']['cost'] * entry['totals']['qty']).values[0] +
@@ -152,18 +150,6 @@ class Backtest:
         strat = ['Strategy']
         summary = pd.DataFrame(data, stats, strat)
 
-        daily_returns = daily_returns[1:].reset_index(level=1, drop=True)
-        daily_returns_df = pd.DataFrame(data=daily_returns.groupby(daily_returns.index.year).apply(list).array,
-                                        index=daily_returns.index.year.unique(),
-                                        columns=[
-                                            'January', 'February', 'March', 'April', 'May', 'June', 'July', 'August',
-                                            'September', 'October', 'November', 'December'
-                                        ])
-
-        sns.heatmap(daily_returns_df, linewidth=0.5, annot=True, fmt='f', cmap='YlGnBu', cbar=False)
-        plt.title('Monthly returns heatmap (in percentage)')
-        plt.show()
-
         return summary
 
     def __repr__(self):
diff --git a/backtester/strategy/__init__.py b/backtester/strategy/__init__.py
index 7ec0bf8..1b111c3 100644
--- a/backtester/strategy/__init__.py
+++ b/backtester/strategy/__init__.py
@@ -1,3 +1,4 @@
 from .strategy import Strategy, Condition
 from .strategy_leg import StrategyLeg
 from .strangle import Strangle
+from .signal import *
diff --git a/backtester/strategy/strategy.py b/backtester/strategy/strategy.py
index 48a0015..50b340c 100644
--- a/backtester/strategy/strategy.py
+++ b/backtester/strategy/strategy.py
@@ -161,7 +161,7 @@ class Strategy:
             subset_df.rename(columns=fields, inplace=True)
 
             order = get_order(leg.direction, signal)
-            subset_df['order'] = order.name
+            subset_df['order'] = order
 
             # Change sign of cost for SELL orders
             if leg.direction == Direction.SELL:
@@ -241,7 +241,7 @@ class Strategy:
         candidates = inventory_leg[['contract']].merge(options, how='left', on='contract')
 
         order = get_order(direction, Signal.EXIT)
-        candidates['order'] = order.name
+        candidates['order'] = order
 
         # Change sign of cost for SELL orders
         if ~direction == Direction.SELL: