import nbformat as nbf


nb = nbf.v4.new_notebook()


text1 = """\
# Part 1: Decision Tree Classifier

This notebook adds the Decision Tree part of the assignment using only the biologically selected dataset `slice_1_significant_7.csv`.

The task is a binary classification problem:

- `0`: Healthy control
- `1`: Exfoliative Glaucoma (XFG)

I use the iris tutorial as a template for the overall workflow:

- define `X` and `y`
- split into training and test sets
- train a `DecisionTreeClassifier`
- evaluate the model
- visualize the final tree

Because this dataset is much smaller than iris, I also use cross-validation and a small hyperparameter search to make the model selection more defensible.
"""


code1 = """\
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from IPython.display import display
from sklearn.inspection import permutation_importance
from sklearn.metrics import (
    ConfusionMatrixDisplay,
    accuracy_score,
    classification_report,
)
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split
from sklearn.tree import DecisionTreeClassifier, plot_tree

sns.set_theme(style="whitegrid")

df = pd.read_csv("slice_1_significant_7.csv")
identifier_cols = ["Internal LIMS ID", "Original ID", "group"]
feature_names = [c for c in df.columns if c not in identifier_cols]

print(f"Samples: {df.shape[0]}")
print(f"Features: {len(feature_names)}")
print("Feature names:")
print(feature_names)
"""


text2 = """\
## Step 1: Prepare Features and Targets

Following the same style as the iris tutorial, I separate the predictors into `X` and the class labels into `y`.
"""


code2 = """\
X = df.drop(columns=identifier_cols)
y = df["group"]

display(X.head())
print("Target distribution:")
print(y.value_counts().sort_index())
"""


text3 = """\
## Step 2: Train-Test Split

I use an 80/20 stratified split so both classes remain balanced in the training and test sets. The training set is used for tuning, while the test set is held back for the final evaluation.
"""


code3 = """\
X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.2,
    stratify=y,
    random_state=42,
)

print("Training set shape:", X_train.shape)
print("Test set shape:", X_test.shape)
"""


text4 = """\
## Step 3: Iteratively Tune the Decision Tree

The assignment asks for iterative development, so I compare multiple settings for:

- `max_depth`
- `min_samples_leaf`
- split criterion: `gini` or `entropy`

For each configuration I record:

- mean 5-fold cross-validation accuracy on the training set
- training accuracy
- test accuracy
- actual tree depth and number of leaves

This helps distinguish a tree that genuinely generalizes from one that merely memorizes the training set.
"""


code4 = """\
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
rows = []

for max_depth in [1, 2, 3, 4, 5, None]:
    for criterion in ["gini", "entropy"]:
        for min_samples_leaf in [1, 2, 4]:
            model = DecisionTreeClassifier(
                max_depth=max_depth,
                criterion=criterion,
                min_samples_leaf=min_samples_leaf,
                random_state=42,
            )

            cv_scores = cross_val_score(
                model,
                X_train,
                y_train,
                cv=cv,
                scoring="accuracy",
            )

            model.fit(X_train, y_train)

            rows.append(
                {
                    "max_depth": "None" if max_depth is None else str(max_depth),
                    "criterion": criterion,
                    "min_samples_leaf": min_samples_leaf,
                    "cv_accuracy_mean": cv_scores.mean(),
                    "cv_accuracy_std": cv_scores.std(),
                    "train_accuracy": accuracy_score(y_train, model.predict(X_train)),
                    "test_accuracy": accuracy_score(y_test, model.predict(X_test)),
                    "actual_tree_depth": model.get_depth(),
                    "n_leaves": model.get_n_leaves(),
                }
            )

results_df = pd.DataFrame(rows).sort_values(
    ["cv_accuracy_mean", "test_accuracy", "n_leaves"],
    ascending=[False, False, True],
)

display(results_df.head(12))
"""


text5 = """\
The top row by cross-validation score is a deeper tree, but it does not perform as well on the held-out test set. That makes overfitting a real concern.

A simpler model with:

- `criterion="entropy"`
- `max_depth=2`
- `min_samples_leaf=2`

has nearly the same cross-validation performance, gives better test accuracy in this split, and is much easier to interpret. I therefore use that as the final tree.
"""


code5 = """\
depth_summary = (
    results_df[results_df["criterion"] == "entropy"]
    .groupby("max_depth", as_index=False)["cv_accuracy_mean"]
    .max()
)

depth_order = ["1", "2", "3", "4", "5", "None"]

plt.figure(figsize=(8, 5))
sns.lineplot(
    data=depth_summary,
    x="max_depth",
    y="cv_accuracy_mean",
    marker="o",
    sort=False,
)
plt.title("Cross-Validation Accuracy by Max Depth")
plt.xlabel("Max depth")
plt.ylabel("Cross-validation accuracy")
plt.show()
"""


text6 = """\
## Step 4: Fit the Final Decision Tree

The final model is deliberately shallow. This trades a small amount of flexibility for better transparency and lower overfitting risk.
"""


code6 = """\
final_params = {
    "criterion": "entropy",
    "max_depth": 2,
    "min_samples_leaf": 2,
    "random_state": 42,
}

final_model = DecisionTreeClassifier(**final_params)
final_model.fit(X_train, y_train)

y_train_pred = final_model.predict(X_train)
y_test_pred = final_model.predict(X_test)

print("Final model parameters:", final_params)
print("Training accuracy:", round(accuracy_score(y_train, y_train_pred), 3))
print("Test accuracy:", round(accuracy_score(y_test, y_test_pred), 3))
print()
print(classification_report(y_test, y_test_pred, digits=3))
"""


text7 = """\
## Step 5: Evaluate the Final Model

The confusion matrix shows how the model handles the two classes rather than only reporting one summary number.
"""


code7 = """\
ConfusionMatrixDisplay.from_predictions(
    y_test,
    y_test_pred,
    display_labels=["Healthy", "XFG"],
    cmap="Blues",
)
plt.title("Confusion Matrix for the Final Decision Tree")
plt.grid(False)
plt.show()
"""


text8 = """\
## Step 6: Interpret Which Proteins the Tree Uses

Decision trees only assign impurity-based importance to features that are actually used in splits. That means it is completely possible for a shallow tree to put all of its importance on only one or two proteins while the others receive zero.

That is not automatically a biological conclusion. In this setting it mostly means:

- the tree found that two proteins were enough for its chosen split rules
- some proteins may contain overlapping information, so the tree only needs one of them
- with only 60 samples, the exact split choices are unstable and can change with a different train/test split

To make that limitation clearer, I look at both the tree's built-in importance values and permutation importance on the test set.
"""


code8 = """\
gini_importance = (
    pd.Series(final_model.feature_importances_, index=feature_names)
    .sort_values(ascending=False)
)

display(gini_importance.to_frame("tree_importance"))

plt.figure(figsize=(10, 5))
sns.barplot(
    x=gini_importance.values,
    y=gini_importance.index,
    hue=gini_importance.index,
    dodge=False,
    palette="crest",
    legend=False,
)
plt.title("Impurity-Based Feature Importance")
plt.xlabel("Importance")
plt.ylabel("Protein fragment")
plt.show()
"""


code9 = """\
perm = permutation_importance(
    final_model,
    X_test,
    y_test,
    n_repeats=50,
    random_state=42,
)

permutation_df = (
    pd.DataFrame(
        {
            "feature": feature_names,
            "permutation_importance_mean": perm.importances_mean,
            "permutation_importance_std": perm.importances_std,
        }
    )
    .sort_values("permutation_importance_mean", ascending=False)
)

display(permutation_df)
"""


text9 = """\
If only two proteins receive non-zero tree importance, that is therefore not necessarily weird. It is a normal consequence of a shallow greedy tree.

What would be weird is claiming from this alone that the other five proteins are unimportant biologically. A single tree cannot support that conclusion, especially with a small dataset. The safer interpretation is that this particular classifier can achieve its decisions using a small subset of the available predictors.
"""


code10 = """\
plt.figure(figsize=(18, 8))
plot_tree(
    final_model,
    feature_names=feature_names,
    class_names=["Healthy", "XFG"],
    filled=True,
    rounded=True,
    fontsize=10,
)
plt.title("Final Decision Tree")
plt.show()
"""


text10 = """\
## Conclusion

Using only the significant seven proteins, the best final decision tree is a shallow entropy-based model with `max_depth=2` and `min_samples_leaf=2`.

This is a reasonable final model because:

- it is built from biologically motivated features
- it was selected through iterative tuning rather than one-shot fitting
- it remains interpretable
- it avoids the strongest overfitting seen in deeper trees

The fact that the tree mainly uses two proteins is not by itself a problem. It reflects how decision trees work: they only reward features that reduce impurity through actual splits. That should be interpreted as a property of this classifier, not as a definitive ranking of biological importance.
"""


nb["cells"] = [
    nbf.v4.new_markdown_cell(text1),
    nbf.v4.new_code_cell(code1),
    nbf.v4.new_markdown_cell(text2),
    nbf.v4.new_code_cell(code2),
    nbf.v4.new_markdown_cell(text3),
    nbf.v4.new_code_cell(code3),
    nbf.v4.new_markdown_cell(text4),
    nbf.v4.new_code_cell(code4),
    nbf.v4.new_markdown_cell(text5),
    nbf.v4.new_code_cell(code5),
    nbf.v4.new_markdown_cell(text6),
    nbf.v4.new_code_cell(code6),
    nbf.v4.new_markdown_cell(text7),
    nbf.v4.new_code_cell(code7),
    nbf.v4.new_markdown_cell(text8),
    nbf.v4.new_code_cell(code8),
    nbf.v4.new_code_cell(code9),
    nbf.v4.new_markdown_cell(text9),
    nbf.v4.new_code_cell(code10),
    nbf.v4.new_markdown_cell(text10),
]


with open("Decision_Tree.ipynb", "w") as f:
    nbf.write(nb, f)

print("Decision_Tree.ipynb created successfully.")