# Core data handling
import yfinance as yf
import pandas as pd
import numpy as np

# Book data (pip install "git+https://github.com/dppalomar/pob.git#subdirectory=python")
from pob_python import SP500_stocks_2015to2020, cryptos_2017to2021_daily

# Statistical analysis
from scipy.stats import norm, skew, kurtosis
import arch

# Visualization
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme(style="darkgrid")
import statsmodels.api as sm
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf

# S&P 500 data
sp500 = yf.download('^GSPC', start='2007-01-01', end='2022-11-04')
sp500_prices = sp500['Close']

# Bitcoin data 
btc = yf.download('BTC-USD', start='2017-01-01', end='2022-11-04')
btc_prices = btc['Close']

# Plot S&P 500 price
fig, ax = plt.subplots(figsize=(12, 6))
np.log(sp500_prices).plot(ax=ax)

# Plot Bitcoin price
fig, ax = plt.subplots(figsize=(12, 6))
np.log(btc_prices).plot(ax=ax)

# S&P 500 returns
sp500_returns = np.log(sp500_prices).diff().dropna()

# Bitcoin returns
btc_returns = np.log(btc_prices).diff().dropna()

def plot_returns(returns, title):
    fig, ax = plt.subplots(figsize=(12,6))
    ax.plot(returns, linewidth=0.5)
    ax.set_title(title)
    ax.set_xlabel('Date')
    ax.set_ylabel('Log Return')
    plt.show()

plot_returns(sp500_returns, 'S&P 500 Daily Log Returns')

plot_returns(btc_returns, 'Bitcoin Daily Log Returns')

def analyze_distribution(returns, asset_name):
    print(f"{asset_name} Distribution Properties:")
    print(f"Skewness: {skew(returns).item():.4f}")  # Fixed line
    print(f"Excess Kurtosis: {kurtosis(returns, fisher=False).item():.4f}") 
    
    fig, ax = plt.subplots(1,2, figsize=(15,5))
    
    # Histogram with normal fit
    sns.histplot(returns, kde=False, ax=ax[0], stat='density')
    x = np.linspace(returns.min(), returns.max(), 100)
    ax[0].plot(x, norm.pdf(x, returns.mean(), 0.6*returns.std()))
    ax[0].set_title('Return Distribution')
    
    # Q-Q Plot
    sm.graphics.qqplot(returns.squeeze(), line='45', fit=True, ax=ax[1])
    ax[1].set_xlim(-4, 4)
    ax[1].set_title('Q-Q Plot')
    
    plt.tight_layout()
    plt.show()

analyze_distribution(sp500_returns, 'S&P 500')

analyze_distribution(btc_returns, 'Bitcoin')

def plot_autocorrelation(returns, title):
    fig, (ax1, ax2) = plt.subplots(2,1, figsize=(12,8))
    plot_acf(returns, ax=ax1, lags=40)
    plot_pacf(returns, ax=ax2, lags=40)
    fig.suptitle(title)
    plt.tight_layout()
    plt.show()

plot_autocorrelation(sp500_returns, 'S&P 500 ACF/PACF')

plot_autocorrelation(btc_returns, 'Bitcoin ACF/PACF')

def estimate_volatility(returns):
    model = arch.arch_model(returns, vol='Garch', p=1, q=1, rescale=False)
    result = model.fit(disp='off')
    return result.conditional_volatility

sp500_vol = estimate_volatility(sp500_returns)
btc_vol = estimate_volatility(btc_returns)

def plot_volatility(returns, vol, title):
    fig, ax = plt.subplots(figsize=(12,6))
    ax.plot(returns, alpha=0.5, label='Returns')
    ax.plot(vol, color='red', label='Conditional Volatility')
    ax.set_title(title)
    ax.legend()
    plt.show()

plot_volatility(sp500_returns, sp500_vol, 'S&P 500 Volatility Clustering')

plot_volatility(btc_returns, btc_vol, 'Bitcoin Volatility Clustering')

plot_autocorrelation(abs(sp500_returns), 'S&P 500 ACF/PACF of absolute value of returns')

plot_autocorrelation(abs(btc_returns), 'Bitcoin ACF/PACF of absolute value of returns')

from pob_python import SP500_stocks_2015to2020, cryptos_2017to2021_hourly

def plot_correlation_matrix(prices, title):
    returns = np.log(prices).diff().dropna()
    corr_matrix = returns.corr()
    
    plt.figure(figsize=(12,10))
    sns.heatmap(corr_matrix, annot=False, cmap='viridis',
                xticklabels=True, yticklabels=True)
    plt.title(title)
    plt.show()

# Example 40 random stocks
data = SP500_stocks_2015to2020.sample(n=40, axis='columns')[-200:]
data.head()

plot_correlation_matrix(data, 'Stock Correlation Matrix')

# Example 40 random cryptos
data = cryptos_2017to2021_hourly.sample(n=40, axis='columns')[-200:]
data.head()

plot_correlation_matrix(data, 'Crypto Correlation Matrix')