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

# Read the price data from the Excel file
price_data = pd.read_excel('Momentum II.xlsx', index_col=0, parse_dates=True)

# Calculate the 120 trading day momentum for each stock
momentum = price_data.pct_change(120)

def get_long_short_positions(df, long_percentile, short_percentile):
    long_criteria = df > np.percentile(df.dropna(), long_percentile)
    short_criteria = df < np.percentile(df.dropna(), short_percentile)
    return long_criteria, short_criteria

# Loop through each day and stock, and calculate long and short positions
portfolio = pd.DataFrame(index=momentum.index, columns=momentum.columns)

for date, row in momentum.iterrows():
    long_pos, short_pos = get_long_short_positions(row, 85, 15)
    portfolio.loc[date, long_pos] = 1
    portfolio.loc[date, short_pos] = -1

# Calculate the daily P&L for long and short positions separately
daily_returns = price_data.pct_change().shift(-1)
long_pnl = (portfolio * daily_returns).where(portfolio == 1).sum(axis=1)
short_pnl = (portfolio * -daily_returns).where(portfolio == -1).sum(axis=1)

# Calculate the total daily P&L by adding the long and short P&L
daily_pnl = long_pnl + short_pnl

# Calculate the cumulative P&L by summing up daily PnL values
cumulative_pnl = daily_pnl.cumsum()

# Plot the cumulative P&L curve
fig, ax = plt.subplots(figsize=(10, 6))

ax.plot(cumulative_pnl)
ax.set_title('Cumulative P&L Curve')
ax.set_xlabel('Date')
ax.set_ylabel('Portfolio Value')
ax.grid()

plt.tight_layout()
plt.show()

with pd.ExcelWriter('trade_details.xlsx') as writer:
    daily_pnl.to_excel(writer, sheet_name='Daily PnL')
    cumulative_pnl.to_excel(writer, sheet_name='Cumulative PnL')
    portfolio.to_excel(writer, sheet_name='Portfolio Positions')