"""
===================================================
Decision Boundaries & Margin Scores
===================================================

Visualizes the discrete decision boundaries generated by ternary weights
:math:`w_j \\in \\{-1, 0, 1\\}` on a 2D synthetic dataset.
"""

import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler
from tlmnet import TlmMilpClassifier

X, y = make_classification(
    n_samples=100,
    n_features=2,
    n_redundant=0,
    n_informative=2,
    random_state=42,
    class_sep=1.2,
)
X = StandardScaler().fit_transform(X)

clf = TlmMilpClassifier(C=1.0)
clf.fit(X, y)

print(f"Fitted Coefficients: {clf.coef_}")

# Plot Decision Boundary
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200), np.linspace(y_min, y_max, 200))

Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)

plt.figure(figsize=(7, 6))
plt.contourf(xx, yy, Z, alpha=0.3, cmap=plt.cm.coolwarm)
plt.scatter(X[:, 0], X[:, 1], c=y, cmap=plt.cm.coolwarm, edgecolors="k", linewidths=0.5)
plt.xlabel("Standardized Feature 1")
plt.ylabel("Standardized Feature 2")
plt.title(f"TLM Decision Surface (coef={list(clf.coef_)})")
plt.tight_layout()
plt.show()