"""
Digit Recognizer: Chunked SGD Baseline

A small CPU-only baseline for Kaggle's Digit Recognizer competition.
It intentionally uses chunked CSV reads and SGDClassifier so it runs on low-memory machines
without GPU quota. The goal is a clean, reproducible starter model — not leaderboard magic.
"""

from pathlib import Path
import gc
import hashlib
import json

import numpy as np
import pandas as pd
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score

OUTPUT = Path("/kaggle/working")
CLASSES = np.arange(10)


def find_digit_input() -> Path:
    """Find the mounted Digit Recognizer directory across Kaggle image variants."""
    candidates = [Path("/kaggle/input/digit-recognizer")]
    candidates.extend(p.parent for p in Path("/kaggle/input").glob("**/train.csv"))
    for candidate in candidates:
        if (candidate / "train.csv").exists() and (candidate / "test.csv").exists() and (candidate / "sample_submission.csv").exists():
            return candidate
    seen = sorted(str(p) for p in Path("/kaggle/input").glob("**/*") if p.is_file())[:80]
    raise FileNotFoundError("Digit Recognizer CSVs not mounted. Visible input files: " + json.dumps(seen, indent=2))


INPUT = find_digit_input()
print(f"Using input directory: {INPUT}")


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


clf = SGDClassifier(loss="log_loss", alpha=1e-4, max_iter=1, tol=None, random_state=42)
first = True
val_x, val_y = [], []

for epoch in range(8):
    for chunk_id, chunk in enumerate(pd.read_csv(INPUT / "train.csv", chunksize=3500)):
        y = chunk.pop("label").astype("int8").to_numpy()
        x = chunk.astype("float32").to_numpy() / 255.0

        # Tiny deterministic holdout from the first epoch only. This is just a sanity check.
        if epoch == 0 and chunk_id < 2:
            val_x.append(x[:250].copy())
            val_y.append(y[:250].copy())
            x = x[250:]
            y = y[250:]

        clf.partial_fit(x, y, classes=CLASSES if first else None)
        first = False
        del chunk, x, y
    gc.collect()

vx = np.vstack(val_x)
vy = np.concatenate(val_y)
holdout_accuracy = float(accuracy_score(vy, clf.predict(vx)))
print(f"small_holdout_accuracy={holdout_accuracy:.4f}")

predictions = []
for chunk in pd.read_csv(INPUT / "test.csv", chunksize=5000):
    x = chunk.astype("float32").to_numpy() / 255.0
    predictions.extend(clf.predict(x).astype(int).tolist())
    del chunk, x

sample = pd.read_csv(INPUT / "sample_submission.csv")
submission = pd.DataFrame({"ImageId": sample["ImageId"], "Label": predictions})
assert len(submission) == len(sample)
assert submission["ImageId"].equals(sample["ImageId"])
assert submission["Label"].between(0, 9).all()

sub_path = OUTPUT / "submission.csv"
submission.to_csv(sub_path, index=False)
manifest = {
    "competition": "digit-recognizer",
    "model": "chunked SGDClassifier(log_loss), 8 CSV epochs, random_state=42",
    "small_holdout_accuracy": holdout_accuracy,
    "rows": int(len(submission)),
    "sha256": sha256(sub_path),
}
(OUTPUT / "manifest.json").write_text(json.dumps(manifest, indent=2))
print(json.dumps(manifest, indent=2))