train_dt.py

# ============================================================
# train_dt.py -- v5 (Kaggle data + 3 training windows)
# ============================================================
# Tests 3 training windows automatically:
#   Window A : Last 2 years  (2022-2023)
#   Window B : Last 5 years  (2019-2023)
#   Window C : Last 10 years (2014-2023)
# Picks the one with best test accuracy
# ============================================================

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import (accuracy_score, confusion_matrix,
                             classification_report)

# -- 1. Load features -----------------------------------------
df = pd.read_csv("nifty50_features.csv",
                 index_col="Date", parse_dates=True, dayfirst=True)
df.dropna(inplace=True)
print(f"Loaded {len(df)} samples")
print(f"Full date range : {df.index[0].date()} --> {df.index[-1].date()}")

FEATURES = ["MA_ratio", "Market_regime", "RSI",
            "Daily_Return", "Return_3d", "Return_5d",
            "Volatility_5", "Price_Position",
            "RSI_lag1", "Daily_Return_lag1", "MA_ratio_lag1"]

X = df[FEATURES]
y = df["Target"]

# -- 2. Define test set and 3 training windows ----------------
TEST_START = "2024-01-01"

X_test = X[X.index >= TEST_START]
y_test = y[y.index >= TEST_START]

# Three windows to compare
WINDOWS = {
    "2 years  (2022-2023)" : "2022-01-01",
    "5 years  (2019-2023)" : "2019-01-01",
    "10 years (2014-2023)" : "2014-01-01",
}

print(f"\nTest set : {len(X_test)} days "
      f"({X_test.index[0].date()} --> {X_test.index[-1].date()})")

# -- 3. Grid search parameters --------------------------------
DEPTHS     = [2, 3, 4, 5, 6, 7, 8]
MIN_LEAVES = [10, 20, 30, 50, 75, 100]

# -- 4. Run grid search for each window -----------------------
print("\n" + "=" * 65)
print("COMPARING 3 TRAINING WINDOWS")
print("=" * 65)

window_results = {}   # store best result per window

for window_name, train_start in WINDOWS.items():

    X_train = X[(X.index >= train_start) & (X.index < TEST_START)]
    y_train = y[(y.index >= train_start) & (y.index < TEST_START)]

    best_acc  = 0
    best_depth = 4
    best_leaf  = 30
    all_combos = []

    for depth in DEPTHS:
        for min_leaf in MIN_LEAVES:
            m = DecisionTreeClassifier(
                    max_depth=depth,
                    min_samples_leaf=min_leaf,
                    class_weight="balanced",
                    random_state=42)
            m.fit(X_train, y_train)
            tr  = accuracy_score(y_train, m.predict(X_train))
            te  = accuracy_score(y_test,  m.predict(X_test))

            # Skip one-class predictions
            if len(set(m.predict(X_test))) < 2:
                continue

            all_combos.append((depth, min_leaf, tr, te))

            if te > best_acc:
                best_acc   = te
                best_depth = depth
                best_leaf  = min_leaf

    window_results[window_name] = {
        "train_start" : train_start,
        "train_days"  : len(X_train),
        "best_acc"    : best_acc,
        "best_depth"  : best_depth,
        "best_leaf"   : best_leaf,
        "all_combos"  : all_combos,
    }

    print(f"\n  Window : {window_name}")
    print(f"  Train days : {len(X_train)}")
    print(f"  Best depth={best_depth}, min_leaf={best_leaf} "
          f"--> {best_acc*100:.1f}%")

# -- 5. Pick the best window ----------------------------------
best_window = max(window_results, key=lambda w: window_results[w]["best_acc"])
best        = window_results[best_window]

print(f"\n{'='*65}")
print(f"WINNER : {best_window}")
print(f"  depth={best['best_depth']}, "
      f"min_leaf={best['best_leaf']}, "
      f"accuracy={best['best_acc']*100:.1f}%")
print(f"{'='*65}")

# -- 6. Train final model on winning window -------------------
X_train_final = X[(X.index >= best["train_start"]) &
                  (X.index < TEST_START)]
y_train_final = y[(y.index >= best["train_start"]) &
                  (y.index < TEST_START)]

model = DecisionTreeClassifier(
            max_depth=best["best_depth"],
            min_samples_leaf=best["best_leaf"],
            class_weight="balanced",
            random_state=42)
model.fit(X_train_final, y_train_final)
pred = model.predict(X_test)

print(f"\nFINAL DECISION TREE")
print(f"  Window : {best_window}")
print(f"  Train  : {len(X_train_final)} days")
print(classification_report(y_test, pred,
                             target_names=["Down", "Up"],
                             zero_division=0))

# -- 7. Feature importance ------------------------------------
importance = pd.Series(model.feature_importances_,
                       index=FEATURES).sort_values(ascending=False)
print("FEATURE IMPORTANCE:")
for feat, val in importance.items():
    bar = "|" * int(val * 50)
    print(f"  {feat:<22} {val:.4f}  {bar}")

# -- 8. Progress tracker --------------------------------------
print("\n" + "=" * 55)
print("PROGRESS ACROSS ALL VERSIONS")
print("=" * 55)
print(f"  v1  Basic DT, raw features        : 45.4%  (all UP)")
print(f"  v2  Better features, depth tune   : 50.0%  (majority trap)")
print(f"  v3  class_weight + grid search    : 44.2%  (honest)")
print(f"  v4  +lags +COVID removed          : 49.6%")
print(f"  v5  Kaggle 33yr + 3 windows       : {best['best_acc']*100:.1f}%")

# -- 9. Plots -------------------------------------------------
fig = plt.figure(figsize=(18, 14))
fig.suptitle(
    f"NIFTY 50 -- Decision Tree v5\n"
    f"Winner: {best_window}  |  "
    f"depth={best['best_depth']}, min_leaf={best['best_leaf']}  |  "
    f"Accuracy={best['best_acc']*100:.1f}%",
    fontsize=13, fontweight="bold")
gs = gridspec.GridSpec(3, 2, figure=fig, hspace=0.45, wspace=0.35)

# Plot 1: Window comparison bar chart
ax1 = fig.add_subplot(gs[0, :])
names     = list(window_results.keys())
accs      = [window_results[w]["best_acc"] * 100 for w in names]
days      = [window_results[w]["train_days"] for w in names]
bar_colors = ["#2E7D32" if n == best_window else "#90CAF9" for n in names]
bars = ax1.bar(names, accs, color=bar_colors, width=0.4)
ax1.axhline(50, color="gray", linestyle="--",
            linewidth=1.5, label="Random chance (50%)")
ax1.set_ylim(40, 65)
ax1.set_title("Best Test Accuracy by Training Window\n"
              "(green = winner)")
ax1.set_ylabel("Test Accuracy %")
ax1.legend(fontsize=9)
ax1.grid(axis="y", alpha=0.3)
for bar, acc, day in zip(bars, accs, days):
    ax1.text(bar.get_x() + bar.get_width()/2,
             bar.get_height() + 0.3,
             f"{acc:.1f}%\n({day} days)",
             ha="center", fontsize=10, fontweight="bold")

# Plot 2: Grid search heatmap for winning window
ax2 = fig.add_subplot(gs[1, :])
grid_matrix = np.zeros((len(DEPTHS), len(MIN_LEAVES)))
for (d, ml, tr, te) in best["all_combos"]:
    if d in DEPTHS and ml in MIN_LEAVES:
        grid_matrix[DEPTHS.index(d)][MIN_LEAVES.index(ml)] = te * 100
im = ax2.imshow(grid_matrix, cmap="YlGn", aspect="auto", vmin=40, vmax=62)
ax2.set_xticks(range(len(MIN_LEAVES))); ax2.set_xticklabels(MIN_LEAVES)
ax2.set_yticks(range(len(DEPTHS)));    ax2.set_yticklabels(DEPTHS)
ax2.set_xlabel("min_samples_leaf")
ax2.set_ylabel("max_depth")
ax2.set_title(f"Grid Search Heatmap -- {best_window}\n"
              "Test Accuracy % per depth + min_leaf combination")
for i in range(len(DEPTHS)):
    for j in range(len(MIN_LEAVES)):
        val = grid_matrix[i][j]
        if val > 0:
            ax2.text(j, i, f"{val:.1f}", ha="center", va="center",
                     fontsize=9, fontweight="bold",
                     color="white" if val > 56 else "black")
plt.colorbar(im, ax=ax2, label="Test Accuracy %")

# Plot 3: Confusion matrix
ax3 = fig.add_subplot(gs[2, 0])
cm = confusion_matrix(y_test, pred)
ax3.imshow(cm, cmap="Blues")
ax3.set_xticks([0, 1]); ax3.set_yticks([0, 1])
ax3.set_xticklabels(["Pred Down", "Pred Up"])
ax3.set_yticklabels(["True Down", "True Up"])
for i in range(2):
    for j in range(2):
        ax3.text(j, i, cm[i, j], ha="center", va="center",
                 fontsize=20, fontweight="bold",
                 color="white" if cm[i, j] > cm.max()/2 else "black")
ax3.set_title(f"Confusion Matrix\n"
              f"Accuracy: {best['best_acc']*100:.1f}%")

# Plot 4: Rolling accuracy
ax4 = fig.add_subplot(gs[2, 1])
res        = X_test.copy()
res["Hit"] = (pred == y_test.values).astype(int)
rolling    = res["Hit"].rolling(30).mean() * 100
ax4.plot(rolling.index, rolling,
         color="#7B1FA2", linewidth=1.5)
ax4.axhline(50, color="gray", linestyle="--",
            linewidth=1, label="Random (50%)")
ax4.axhline(best["best_acc"] * 100, color="green",
            linestyle="--", linewidth=1,
            label=f"Overall ({best['best_acc']*100:.1f}%)")
ax4.fill_between(rolling.index, rolling, 50,
                 where=(rolling >= 50),
                 alpha=0.15, color="green")
ax4.fill_between(rolling.index, rolling, 50,
                 where=(rolling < 50),
                 alpha=0.15, color="red")
ax4.set_title("30-day Rolling Accuracy (2024 test period)")
ax4.set_ylabel("Accuracy %")
ax4.legend(fontsize=9)
ax4.grid(alpha=0.3)

plt.savefig("nifty50_dt_v5.png", dpi=150, bbox_inches="tight")
plt.show()
print("\nSaved nifty50_dt_v5.png")

# Tree diagram
fig2, ax2 = plt.subplots(figsize=(26, 12))
plot_tree(model, feature_names=FEATURES,
          class_names=["Down", "Up"],
          filled=True, rounded=True, fontsize=8, ax=ax2)
fig2.suptitle(
    f"Final Decision Tree v5  "
    f"(depth={best['best_depth']}, min_leaf={best['best_leaf']}, "
    f"{best_window})",
    fontsize=12, fontweight="bold")
plt.tight_layout()
plt.savefig("nifty50_dt_v5_tree.png", dpi=150, bbox_inches="tight")
plt.show()
print("Saved nifty50_dt_v5_tree.png")

# ============================================================
# FINE-TUNING -- top 3 features only + class weight nudge
# ============================================================
# Only RSI, Volatility_5, Return_5d had non-zero importance.
# Drop the 8 dead features and try nudging Down class weight
# to improve the 16% Down recall without hurting accuracy.

print("\n" + "=" * 58)
print("FINE-TUNING -- top 3 features + class weight options")
print("=" * 58)

TOP_FEATURES  = ["RSI", "Volatility_5", "Return_5d"]
X_train_top   = X_train_final[TOP_FEATURES]
X_test_top    = X_test[TOP_FEATURES]

weight_options = [
    ("balanced",        "balanced (auto)"),
    ({0: 1.2, 1: 1.0},  "Down x1.2"),
    ({0: 1.5, 1: 1.0},  "Down x1.5"),
    ({0: 2.0, 1: 1.0},  "Down x2.0"),
]

print(f"  {'Weight':<20} {'Accuracy':>10} {'Down recall':>14} {'Up recall':>12}")
print("  " + "-" * 58)

best_fine_acc  = 0
best_fine_pred = pred

for cw, label in weight_options:
    m = DecisionTreeClassifier(
            max_depth=best["best_depth"],
            min_samples_leaf=best["best_leaf"],
            class_weight=cw,
            random_state=42)
    m.fit(X_train_top, y_train_final)
    p   = m.predict(X_test_top)
    acc = accuracy_score(y_test, p)
    if len(set(p)) < 2:
        print(f"  {label:<20} -- one class only, skip")
        continue
    cm_f = confusion_matrix(y_test, p)
    dr   = cm_f[0][0] / cm_f[0].sum()
    ur   = cm_f[1][1] / cm_f[1].sum()
    note = "  <- best" if acc > best_fine_acc else ""
    if acc > best_fine_acc:
        best_fine_acc  = acc
        best_fine_pred = p
    print(f"  {label:<20} {acc*100:>9.1f}%  {dr*100:>11.1f}%  {ur*100:>9.1f}%{note}")

print()
if best_fine_acc >= best["best_acc"]:
    print(f"  Fine-tuned model wins  : {best_fine_acc*100:.1f}%")
    final_acc = best_fine_acc
else:
    print(f"  Original model holds   : {best['best_acc']*100:.1f}%")
    final_acc = best["best_acc"]

print("\n" + "=" * 58)
print("FINAL PROGRESS SUMMARY")
print("=" * 58)
print(f"  v1  Basic DT, raw features        : 45.4%  (all UP)")
print(f"  v2  Better features, depth tune   : 50.0%  (majority trap)")
print(f"  v3  class_weight + grid search    : 44.2%  (honest)")
print(f"  v4  +lags +COVID removed          : 49.6%")
print(f"  v5  Kaggle + 3 windows            : {best['best_acc']*100:.1f}%")
print(f"  v5f fine-tuned top 3 features     : {final_acc*100:.1f}%")
print()
print(f"  Total improvement (honest v3->v5f): +{(final_acc-0.442)*100:.1f}%")