import warnings
warnings.filterwarnings('ignore')
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yfinance as yf

tickers = ['AAPL','MSFT','GOOGL','AMZN','TSLA']
ohlcv = yf.download(tickers, period='1y', auto_adjust=True, progress=False)
ohlcv = ohlcv.stack(future_stack=True).rename_axis(['Date','Ticker']).reset_index()
ohlcv.columns.name = None
print(ohlcv.shape)
print(ohlcv.head(3))

# Beginner Example 1: How does AAPL price change by month?
aapl = ohlcv[ohlcv['Ticker'] == 'AAPL'].copy()
aapl['Month'] = pd.to_datetime(aapl['Date']).dt.month
monthly_means = aapl.groupby('Month')['Close'].mean().round(2)
print(monthly_means)

# Beginner Example 2: Plotting monthly average price for Apple
monthly_means.plot(kind='bar', title='AAPL Average Closing Price by Month')
plt.ylabel('Average Closing Price (USD)')
plt.xlabel('Month')
plt.show()

# Beginner Example 3: Add month names for clarity
month_names = {1: 'Jan', 2: 'Feb', 3: 'Mar', 4: 'Apr', 5: 'May', 6: 'Jun',
               7: 'Jul', 8: 'Aug', 9: 'Sep', 10: 'Oct', 11: 'Nov', 12: 'Dec'}
monthly_means.index = monthly_means.index.map(month_names)
monthly_means.plot(kind='bar', color='skyblue', title='AAPL Average Price by Calendar Month')
plt.ylabel('Average Closing Price (USD)')
plt.show()

# Intermediate Example 1: Compute average monthly returns instead of price
aapl['Return'] = aapl['Close'].pct_change()
aapl['Month'] = pd.to_datetime(aapl['Date']).dt.month
monthly_returns = aapl.groupby('Month')['Return'].mean().mul(100).round(3)
monthly_returns.index = monthly_returns.index.map(month_names)
print(monthly_returns)

# Intermediate Example 2: Visualize mean monthly returns for AAPL
monthly_returns.plot(kind='bar', color='orange', title='AAPL Mean Monthly Returns (%)')
plt.ylabel('Mean Monthly Return (%)')
plt.xlabel('Month')
plt.axhline(0, color='gray', linestyle='--')
plt.show()

# Intermediate Example 3: Compare monthly returns of all five stocks
monthly_returns_all = (ohlcv.assign(Month = pd.to_datetime(ohlcv['Date']).dt.month)
    .groupby(['Ticker','Month'])['Close'].apply(lambda x: x.pct_change().mean()*100).unstack(0).round(2))
monthly_returns_all.index = monthly_returns_all.index.map(month_names)
print(monthly_returns_all)

# Intermediate Example 4: Visualize all stocks' monthly returns in one chart
monthly_returns_all.plot(kind='bar', figsize=(12,5), title='Mean Monthly Returns by Ticker (%)')
plt.ylabel('Mean Monthly Return (%)')
plt.xlabel('Month')
plt.axhline(0, color='gray', linestyle='--', linewidth=1)
plt.legend(title='Ticker')
plt.tight_layout()
plt.show()

# Intermediate Example 5: Which month has the highest average return overall?
monthly_returns_all['Average'] = monthly_returns_all.mean(axis=1)
best_month = monthly_returns_all['Average'].idxmax()
print(f'Highest average return month: {best_month}')

# Advanced Example 1: Analyzing holiday season effects (Nov-Dec)
holiday_months = ['Nov', 'Dec']
holiday_returns = monthly_returns_all.loc[holiday_months]
print('Mean returns for Nov/Dec by ticker:')
print(holiday_returns)

# Advanced Example 2: Test statistical significance of seasonality
from scipy.stats import f_oneway
returns_by_month = [aapl[aapl['Month']==m]['Return'].dropna() for m in range(1,13)]
fstat, pval = f_oneway(*returns_by_month)
print(f'ANOVA F-stat: {fstat:.4f}, p-value: {pval:.4f}')
if pval < 0.05:
    print('Significant monthly return differences detected!')
else:
    print('No significant difference in monthly returns.')

# Advanced Example 3: Heatmap for visualizing returns by ticker and month
import seaborn as sns
plt.figure(figsize=(9,5))
sns.heatmap(monthly_returns_all.T, annot=True, cmap='RdYlGn', center=0)
plt.title('Monthly Mean Returns Heatmap (%)')
plt.xlabel('Month')
plt.ylabel('Ticker')
plt.show()

# Advanced Example 4: Rolling 3-month seasonal returns
def three_month_rolling(df):
    df = df.sort_values('Date')
    df['3mo_ret'] = df['Close'].pct_change(periods=63)
    df['Month'] = pd.to_datetime(df['Date']).dt.month
    return df

aapl_3mo = three_month_rolling(aapl)
three_monthly = aapl_3mo.groupby('Month')['3mo_ret'].mean().mul(100).round(2)
three_monthly.index = three_monthly.index.map(month_names)
print(three_monthly)

# Error Handling Example 1: What if data for a month is missing?
missing_months = set(range(1,13)) - set(aapl['Month'])
print(f'Missing data for months: {missing_months}')

# Error Handling Example 2: Handling zeros or NaNs in returns
bad_rows = aapl[aapl['Return'].isna()]
print('Rows with missing return:', bad_rows)

# Best Practice 1: Encapsulate seasonality logic as a reusable function
def get_monthly_returns(df, ticker):
    sub = df[df['Ticker']==ticker].copy()
    sub['Return'] = sub['Close'].pct_change()
    sub['Month'] = pd.to_datetime(sub['Date']).dt.month
    return sub.groupby('Month')['Return'].mean().mul(100).round(2)

msft_monthly_returns = get_monthly_returns(ohlcv, 'MSFT')
print(msft_monthly_returns)

# Best Practice 2: Visualize uncertainty using error bars
means = aapl.groupby('Month')['Return'].mean().mul(100)
stds = aapl.groupby('Month')['Return'].std().mul(100)
plt.bar(month_names.values(), means, yerr=stds, capsize=6, alpha=0.6, color='slateblue')
plt.title('AAPL Mean Monthly Return with Standard Deviation')
plt.ylabel('Monthly Return (%)')
plt.xlabel('Month')
plt.axhline(0, color='gray', linestyle='dashed')
plt.show()

# End-to-End Example: Identify the best and worst months for TSLA over the last year
def annotate_extremes(monthly_ret):
    best = monthly_ret.idxmax()
    worst = monthly_ret.idxmin()
    print(f'Best month for TSLA: {month_names[best]} ({monthly_ret[best]:.2f}%)')
    print(f'Worst month for TSLA: {month_names[worst]} ({monthly_ret[worst]:.2f}%)')

tsla_ret = get_monthly_returns(ohlcv, 'TSLA')
annotate_extremes(tsla_ret)