# -*- coding: utf-8 -*-
"""Generate the simulated returns that the post loads, returns.csv, and check their momentum.

Six years of daily returns for 40 stocks. A stock's daily return is a market move shared by all
stocks, the stock's own part of the expected return and noise. The own part persists from day
to day, which is the momentum the post's strategy looks for. Run the script once; the post only
reads the file it writes.
"""
import numpy as np
import pandas as pd

# A fixed seed gives the same draws on every run, so the file and the post's numbers reproduce.
rng = np.random.default_rng(42)

N_STOCKS = 40
N_DAYS = 252 * 6                  # six years of 252 trading days
MARKET_MEAN = 0.0005              # the market move's mean, 0.05% a day, my choice
MARKET_SD = 0.010                 # the market move's standard deviation, 1% a day, my choice
NOISE_SD = 0.012                  # each stock's daily noise, 1.2%, my choice
PHI = 0.997                       # share of yesterday's own part kept today, my choice
# The spread of the own parts across stocks, 0.05% a day, sized so that the momentum matches
# the evidence: expected_decile_spread() below checks it.
OWN_PART_SD = 0.0005

dates = pd.bdate_range("2018-01-02", periods=N_DAYS)   # business days, Monday to Friday
tickers = []
for number in range(N_STOCKS):
    tickers.append(f"STK{number:02d}")                 # STK00 to STK39; :02d pads to two digits

# One market move per day, shared by all 40 stocks
market = rng.normal(MARKET_MEAN, MARKET_SD, N_DAYS)

# The own part starts as a normal draw across stocks. Each day it keeps PHI of yesterday's value
# and adds an independent normal draw with a standard deviation of sqrt(1 - PHI**2) times
# OWN_PART_SD. The two parts are independent, so their variances add up to
# PHI**2 * OWN_PART_SD**2 + (1 - PHI**2) * OWN_PART_SD**2 = OWN_PART_SD**2, and the spread across
# stocks stays at OWN_PART_SD.
own_part = np.zeros((N_DAYS, N_STOCKS))
own_part[0] = rng.normal(0, OWN_PART_SD, N_STOCKS)
for day in range(1, N_DAYS):
    fresh_draw = rng.normal(0, OWN_PART_SD * np.sqrt(1 - PHI**2), N_STOCKS)
    own_part[day] = PHI * own_part[day - 1] + fresh_draw

# Noise, independent across stocks and days
noise = rng.normal(0, NOISE_SD, (N_DAYS, N_STOCKS))

# market[:, None] turns the 1,512 market moves into one column, so each day's move is added to
# all 40 stocks of that day
returns = pd.DataFrame(market[:, None] + own_part + noise, index=dates, columns=tickers)
returns.index.name = "date"
returns.to_csv("returns.csv")
print(f"wrote returns.csv: {N_STOCKS} stocks x {N_DAYS} days")


def expected_decile_spread(n_stocks=400_000, seed=1):
    """The expected monthly return of the tenth of stocks with the highest return over the past
    12 months minus the tenth with the lowest, over the three months after the ranking.

    This is the comparison in Jegadeesh and Titman (1993), Table I, for a 12-month ranking and a
    3-month holding period. The function simulates the own part and noise of n_stocks stocks for
    252 days and ranks them on their compounded return. It leaves out the market move, which
    every stock shares, so it cancels in the difference. For the three months after the ranking,
    it uses the expected own part, PHI**k times the own part on the ranking day, k days later,
    instead of drawing it, so the only randomness left is in the ranking. It has its own random
    generator, so returns.csv does not depend on it.
    """
    check_rng = np.random.default_rng(seed)
    own = check_rng.normal(0, OWN_PART_SD, n_stocks)
    log_growth = np.zeros(n_stocks)
    for day in range(252):
        own = PHI * own + check_rng.normal(0, OWN_PART_SD * np.sqrt(1 - PHI**2), n_stocks)
        # log1p(r) is the log of 1 + r; adding these up and taking the exponential compounds
        # the daily returns, so log_growth ranks the stocks as their compounded return does
        log_growth += np.log1p(own + check_rng.normal(0, NOISE_SD, n_stocks))
    # argsort lists the stocks from the lowest past return to the highest
    order = np.argsort(log_growth)
    tenth = n_stocks // 10
    own_gap = own[order[-tenth:]].mean() - own[order[:tenth]].mean()
    # the expected own parts over the next 63 days add up to own_gap times the sum of PHI**k
    # for k = 1 to 63; dividing by 3 gives the average month
    phi_sum = 0.0
    for k in range(1, 64):
        phi_sum += PHI**k
    return own_gap * phi_sum / 3


print(f"expected monthly return of the top tenth minus the bottom tenth over three months: "
      f"{expected_decile_spread():.2%} (Jegadeesh and Titman 1993: 1.31%)")