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

# Define tickers
stock = 'NVDA'
sector_etf = 'SMH'
lookback = 365  # days

# Download daily data
stock_df = yf.download(stock, period=f'{lookback}d', progress=False)
sector_df = yf.download(sector_etf, period=f'{lookback}d', progress=False)

# Calculate technical indicators
stock_df['SMA20'] = stock_df['Close'].rolling(window=20).mean()
stock_df['SMA50'] = stock_df['Close'].rolling(window=50).mean()
stock_df['SMA200'] = stock_df['Close'].rolling(window=200).mean()
stock_df['DailyRet'] = stock_df['Close'].pct_change()
stock_df['Vol20'] = stock_df['DailyRet'].rolling(window=20).std() * np.sqrt(252)  # annualized

# Correlation with sector ETF
merged = pd.merge(stock_df[['Close']], sector_df[['Close']], left_index=True, right_index=True, suffixes=('_stock', '_etf'))
merged['DailyRet_stock'] = merged['Close_stock'].pct_change()
merged['DailyRet_etf'] = merged['Close_etf'].pct_change()
corr = merged['DailyRet_stock'].rolling(window=60).corr(merged['DailyRet_etf'])

# Plot price with SMAs
plt.figure(figsize=(12,6))
plt.plot(stock_df['Close'], label=f'{stock} Close')
plt.plot(stock_df['SMA20'], label='SMA20')
plt.plot(stock_df['SMA50'], label='SMA50')
plt.plot(stock_df['SMA200'], label='SMA200')
plt.title(f'{stock} Price with Moving Averages')
plt.legend()
plt.tight_layout()
plt.savefig('nvda_ma_chart.png')

# Plot volatility
plt.figure(figsize=(12,4))
plt.plot(stock_df['Vol20'], label='20‑Day Vol (annualized)')
plt.title(f'{stock} 20‑Day Volatility')
plt.legend()
plt.tight_layout()
plt.savefig('nvda_vol_chart.png')

# Plot correlation
plt.figure(figsize=(12,4))
plt.plot(corr, label='60‑Day Correlation with SMH')
plt.title(f'{stock} vs SMH Correlation')
plt.legend()
plt.tight_layout()
plt.savefig('nvda_corr_chart.png')

print('Charts generated: nvda_ma_chart.png, nvda_vol_chart.png, nvda_corr_chart.png')