# Run this cell first to install required libraries
!pip install yfinance scikit-learn matplotlib seaborn pandas numpy --quiet

!pip install yfinance

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

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import (
    classification_report, confusion_matrix,
    ConfusionMatrixDisplay, roc_curve, roc_auc_score
)
from sklearn.preprocessing import StandardScaler

import warnings
warnings.filterwarnings('ignore')

# Plot style
plt.rcParams['figure.figsize'] = (10, 5)
plt.rcParams['axes.spines.top'] = False
plt.rcParams['axes.spines.right'] = False

print('All libraries imported successfully!')

# Download AAPL data
ticker = 'AAPL'
df = yf.download(ticker, start='2022-01-01', end='2024-12-31', auto_adjust=True)

# Flatten multi-level columns if present
if isinstance(df.columns, pd.MultiIndex):
    df.columns = df.columns.get_level_values(0)

print(f'Shape: {df.shape}')
print(f'Date range: {df.index[0].date()} → {df.index[-1].date()}')
df.head()

# 📊 Visualise the closing price
fig, ax = plt.subplots(figsize=(12, 4))
ax.plot(df.index, df['Close'], color='steelblue', linewidth=1.5)
ax.set_title(f'{ticker} Closing Price (2022–2024)', fontsize=14, fontweight='bold')
ax.set_ylabel('Price (USD)')
ax.set_xlabel('')
plt.tight_layout()
plt.show()

df.head(10)

df['MA5']  = df['Close'].rolling(window=5).mean()
df['MA20'] = df['Close'].rolling(window=20).mean()

# Plot
fig, ax = plt.subplots(figsize=(12, 4))
ax.plot(df.index, df['Close'], label='Close',  color='lightgray', linewidth=1)
ax.plot(df.index, df['MA5'],   label='MA5',    color='steelblue', linewidth=1.5)
ax.plot(df.index, df['MA20'],  label='MA20',   color='tomato',    linewidth=1.5)
ax.set_title('Closing Price with Moving Averages', fontsize=13, fontweight='bold')
ax.legend()
plt.tight_layout()
plt.show()

# A new column MA5, MA20
df.head(25)

# Now we build RSI

def compute_rsi(series, period=14):
    delta = series.diff()
    gain  = delta.clip(lower=0)
    loss  = -delta.clip(upper=0)
    avg_gain = gain.rolling(window=period).mean()
    avg_loss = loss.rolling(window=period).mean()
    rs  = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

df['RSI'] = compute_rsi(df['Close'])

# Plot
fig, ax = plt.subplots(figsize=(12, 3))
ax.plot(df.index, df['RSI'], color='darkorange', linewidth=1.2)
ax.axhline(70, color='red',   linestyle='--', linewidth=0.8, label='Overbought (70)')
ax.axhline(30, color='green', linestyle='--', linewidth=0.8, label='Oversold (30)')
ax.set_title('RSI (14-day)', fontsize=13, fontweight='bold')
ax.set_ylim(0, 100)
ax.legend()
plt.tight_layout()
plt.show()

# We see that there is a new column - RSI
df.head(20)

df.head(10)

df.head(20)

# Drop rows with NaN (from rolling windows)
df_clean = df[['RSI', 'MA5', 'MA20', 'Return', 'Volume_change', 'Return_lag1', 'Target']].dropna()


print(f'Dataset size after cleaning: {df_clean.shape[0]} rows')
print(f"\nTarget distribution:")
counts = df_clean['Target'].value_counts()
print(f"  UP   (1): {counts[1]} days ({counts[1]/len(df_clean)*100:.1f}%)")
print(f"  DOWN (0): {counts[0]} days ({counts[0]/len(df_clean)*100:.1f}%)")

df_clean.head(5)

# Quick look at descriptive stats for few of our features

df_clean[['RSI', 'MA5', 'MA20', 'Return']].describe().round(3)

# Step 1 — Features and label
X = df_clean[['RSI', 'MA5', 'MA20', 'Volume_change', 'Return_lag1']]
y = df_clean['Target']

# Step 2 — Train / test split (80% train, 20% test, time-ordered)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, shuffle=False  # shuffle=False preserves time order
)

print(f'Training samples : {len(X_train)}')
print(f'Test samples     : {len(X_test)}')

# Lets look at few records of train

X_train.iloc[0:3]
# X_test.iloc[0:3]

# compare actual vs scaled values
print(list(X_train.iloc[0]))
print((X_train_scaled[0]))

# Step 4 — Train logistic regression
model = LogisticRegression(class_weight='balanced',random_state=42)
model.fit(X_train_scaled, y_train)

print('✅ Model trained!')
print(f'\nTrain accuracy : {model.score(X_train_scaled, y_train):.3f}')
print(f'Test  accuracy : {model.score(X_test_scaled,  y_test):.3f}')

features = ['RSI', 'MA5', 'MA20', 'Volume_change', 'Return_lag1']
coefs    = model.coef_[0]

coef_df = pd.DataFrame({'Feature': features, 'Coefficient': coefs})
coef_df['Direction'] = coef_df['Coefficient'].apply(lambda x: '↑ Bullish' if x > 0 else '↓ Bearish')

print(f'Intercept (β₀): {model.intercept_[0]:.4f}\n')
print(coef_df.to_string(index=False))

# Bar chart
colors = ['tomato' if c < 0 else 'steelblue' for c in coefs]
fig, ax = plt.subplots(figsize=(7, 3))
ax.barh(features, coefs, color=colors)
ax.axvline(0, color='black', linewidth=0.8)
ax.set_title('Logistic Regression Coefficients', fontsize=13, fontweight='bold')
ax.set_xlabel('Coefficient value')
plt.tight_layout()
plt.show()

z = np.linspace(-8, 8, 300)
sigmoid = 1 / (1 + np.exp(-z))

fig, ax = plt.subplots(figsize=(8, 4))

ax.plot(z, sigmoid,
        color='steelblue',
        linewidth=2.5,
        label='Sigmoid σ(z)'
)

ax.axhline(0.5, color='tomato', linestyle='--', linewidth=1, label='Decision boundary (0.5)')
ax.axvline(0,   color='gray',   linestyle='--', linewidth=0.8)
ax.fill_between(z, 0.5, sigmoid, where=(sigmoid > 0.5), alpha=0.1, color='green', label='Predict UP')
ax.fill_between(z, sigmoid, 0.5, where=(sigmoid < 0.5), alpha=0.1, color='red',   label='Predict DOWN')
ax.set_title('The Sigmoid Function', fontsize=13, fontweight='bold')
ax.set_xlabel('z  =  β₀ + β·x')
ax.set_ylabel('P(y = 1 | x)')
ax.legend()
plt.tight_layout()
plt.show()

# This line already exists in your notebook — y_prob is created here
y_prob = model.predict_proba(X_test_scaled)[:, 1]

# NOTE: Play around with this number in class
y_pred = (y_prob > 0.45).astype(int)

cm = confusion_matrix(y_test, y_pred)
print("confusion matrix:\n", cm)

disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=['DOWN (0)', 'UP (1)'])

fig, ax = plt.subplots(figsize=(5, 4))
disp.plot(ax=ax, colorbar=False, cmap='Blues')
ax.set_title('Confusion Matrix (Test Set)', fontsize=13, fontweight='bold')
plt.tight_layout()
plt.show()

print(classification_report(y_test, y_pred, target_names=['DOWN (0)', 'UP (1)']))

y_prob = model.predict_proba(X_test_scaled)[:, 1]
fpr, tpr, _ = roc_curve(y_test, y_prob)
auc = roc_auc_score(y_test, y_prob)

fig, ax = plt.subplots(figsize=(6, 5))
ax.plot(fpr, tpr, color='steelblue', linewidth=2, label=f'Logistic Regression (AUC = {auc:.3f})')
ax.plot([0, 1], [0, 1], 'k--', linewidth=1, label='Random guess (AUC = 0.5)')
ax.set_xlabel('False Positive Rate')
ax.set_ylabel('True Positive Rate')
ax.set_title('ROC Curve', fontsize=13, fontweight='bold')
ax.legend()
plt.tight_layout()
plt.show()

print(X_train.iloc[0])        # first row
print(y_train[0])

# --- Change these values and re-run! ---
rsi_today  =  53    # RSI between 0–100  
ma5_today  = 161   # 5-day average price
ma20_today = 166    # 20-day average price
volume_change_today = -0.35
return_lag1_today = 0.069

# ---------------------------------------

new_data = np.array([[rsi_today, ma5_today, ma20_today, volume_change_today, return_lag1_today]])
new_data_scaled = scaler.transform(new_data)

probability = model.predict_proba(new_data_scaled)[0][1]
prediction  = model.predict(new_data_scaled)[0]

label = 'UP' if prediction == 1 else 'DOWN'
# print(f'Input  → RSI={rsi_today}, MA5={ma5_today}, MA20={ma20_today}')
print(f'Input  → {new_data}')

print(f'P(UP)  → {probability:.3f}')
print(f'Prediction → {label}')