marketsim.py

import yaml
import os
import time
from typing import Union, Dict
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

from stockwrapper import StockWrapper
from strategy import Strategy
from ml.ql.agent import IntradayStockMarketAgent, StockMarketAgent


class MarketSimulator:
    def __init__(
        self, strategy: Union[Strategy, StockMarketAgent], stock_wrapper: StockWrapper
    ):
        self.strategy = strategy
        self.stock_wrapper = stock_wrapper

        self.train_end_idx = self.strategy.env.unwrapped.discrete_observations.index[-1]
        self.test_start_idx = self.stock_wrapper.data.index[
            self.stock_wrapper.data.index.get_loc(self.train_end_idx) + 1
        ]

        self.train_data = self.stock_wrapper.data.loc[: self.train_end_idx].copy()
        self.test_data = self.stock_wrapper.data.loc[self.test_start_idx :].copy()

        self.train_gain = None
        self.test_gain = None

        if isinstance(strategy, StockMarketAgent):
            self.discrete_observations = (
                self.stock_wrapper.generate_discrete_observations(
                    self.strategy.env.unwrapped.indicators
                )
            )

        self.directory = Path(
            f"reports/marketsim_{stock_wrapper.symbol}_{int(time.time())}"
        )

        self._run()

    def _run(self):
        if isinstance(self.strategy, Strategy):
            self.stock_wrapper.data.concat_indicators(self.strategy.indicators)
            self.stock_wrapper.data["state"] = (
                self.stock_wrapper.data.apply(self.strategy.predict, axis=1)
                .shift(fill_value=0)
                .ffill()
            )
        elif isinstance(self.strategy, StockMarketAgent):
            state = pd.Series(data=0, index=self.stock_wrapper.data.index)

            if isinstance(self.strategy, IntradayStockMarketAgent):
                for index, observations in self.discrete_observations.iterrows():
                    state[index] = self.strategy.env.unwrapped.action_mapping[
                        self.strategy.predict(index.time(), observations.to_numpy())
                    ]
            else:
                for index, observations in self.discrete_observations.iterrows():
                    state[index] = self.strategy.env.unwrapped.action_mapping[
                        self.strategy.predict(observations.to_numpy())
                    ]

            state = state.shift(fill_value=0).ffill()
            self.train_data["state"] = state.loc[: self.train_end_idx]
            self.test_data["state"] = state.loc[self.test_start_idx :]

        self.train_data["percent_gain"] = MarketSimulator.cumulative_gain(
            self.train_data["close"],
            self.train_data["state"],
        )

        self.test_data["percent_gain"] = MarketSimulator.cumulative_gain(
            self.test_data["close"],
            self.test_data["state"],
        )

        self.train_gain = self.train_data["percent_gain"].iloc[-1]
        self.test_gain = self.test_data["percent_gain"].iloc[-1]

    def plot_close_vs_gain(self, save=True):
        def plot_single_cvg_chart(ax, ax_idx, data, label):
            close = ax[ax_idx].plot(
                data.index,
                data["close"],
                color="orange",
                label=self.stock_wrapper.symbol,
                alpha=0.7,
            )

            ax[ax_idx].tick_params(axis="y", labelcolor="orange")
            ax[ax_idx].set_ylabel("$", rotation=0, labelpad=20, color="orange")
            ax2 = ax[ax_idx].twinx()

            gain = ax2.plot(
                data.index,
                data["percent_gain"] * 100,
                color="dodgerblue",
                label=data["percent_gain"].name.title().replace("_", " "),
                alpha=0.7,
            )

            ax2.tick_params(axis="y", labelcolor="dodgerblue")
            ax2.set_ylabel("%", rotation=0, labelpad=20, color="dodgerblue")
            time_format = mdates.DateFormatter("%m-%y")
            ax[ax_idx].xaxis.set_major_formatter(time_format)

            lines = [close[0], gain[0]]
            labels = [line.get_label() for line in lines]
            ax[ax_idx].set_xlabel("Time")
            ax[ax_idx].legend(lines, labels)
            ax[ax_idx].set_title(f"Adjusted Close VS Percent Gain ({label})")

        fig, ax = plt.subplots(1, 2, figsize=(20, 8))

        plot_single_cvg_chart(ax, 0, self.train_data, "Train")
        plot_single_cvg_chart(ax, 1, self.test_data, "Test")

        if save:
            fig.savefig(self.directory / "close_vs_gain.png", format="png")
            plt.close()
        else:
            plt.plot()

    def max_drawdown(self, period: str):
        if period is None:
            drawdown = (
                self.test_data["percent_gain"] - self.test_data["percent_gain"].cummax()
            )
        else:
            drawdown = (
                self.test_data["percent_gain"]
                .resample(period)
                .agg(lambda x: (x - x.cummax()).max())
            )

        min_idx = drawdown.idxmin()
        return drawdown[min_idx] / (
            self.test_data["percent_gain"].cummax().loc[min_idx] + 1
        )

    def _unsampled_test_points(self) -> int:
        train_observations = self.discrete_observations.loc[: self.train_end_idx]
        test_observations = self.discrete_observations.loc[self.test_start_idx :]

        diff = test_observations.merge(
            train_observations,
            how="left",
            indicator=True,
        )
        return diff[diff["_merge"] == "left_only"].shape[0]

    @staticmethod
    def cumulative_gain(close: pd.Series, state: pd.Series):
        return (1 + (close.pct_change() * state)).cumprod() - 1

    def generate_report(self, save=True):
        report_data = {
            "report": {
                "stock data": {
                    "ticker": self.stock_wrapper.symbol,
                    "train start date": self.train_data.index[0].date(),
                    "test start date": self.test_data.index[0].date(),
                    "end date": self.test_data.index[-1].date(),
                    "frequency": self.stock_wrapper.period,
                },
                "strategy performance": {
                    "total gain": round(
                        float(self.test_gain) * 100,
                        4,
                    ),
                    "maximum drawdown": {
                        "daily": float(round(self.max_drawdown("D") * 100, 4)),
                        "overall": float(round(self.max_drawdown(None) * 100, 4)),
                    },
                },
            }
        }

        if isinstance(self.strategy, StockMarketAgent):
            report_data["report"].update(
                {
                    "q-learner performance": {
                        "episodes": self.strategy.n_episodes,
                        "state coverage": round(
                            float(
                                len(self.strategy.q_values)
                                / (
                                    self.discrete_observations.max()
                                    - self.discrete_observations.min()
                                    + 1
                                ).prod()
                            )
                            * 100,
                            4,
                        ),
                        "unsampled test points": self._unsampled_test_points(),
                        "train gain": round(
                            float(self.train_gain) * 100,
                            4,
                        ),
                        "test gain": round(
                            float(self.test_gain) * 100,
                            4,
                        ),
                        "temporal difference": round(
                            float(np.mean(self.strategy.temporal_difference[-1])), 3
                        ),
                    }
                }
            )

            report_data["report"]["indicators"] = self.strategy.env.unwrapped.indicators
        else:
            report_data["report"]["indicators"] = self.strategy.indicators

        if save:
            self.directory.mkdir(parents=True, exist_ok=True)
            with open(self.directory / "report.yaml", "w+") as fd:
                yaml.dump(report_data, fd, default_flow_style=False)
        else:
            print(yaml.dump(report_data, default_flow_style=False))

        self.plot_close_vs_gain(save=save)

        if isinstance(self.strategy, StockMarketAgent):
            self.strategy.plot_training_data(self.directory, save=save)