"""Download MNIST and prepare .npy files expected by main.py."""

import argparse
import gzip
import shutil
import urllib.request
from pathlib import Path

import idx2numpy
import numpy as np

BASE_URL = "https://storage.googleapis.com/cvdf-datasets/mnist"
FILES = {
    "train-images-idx3-ubyte.gz": f"{BASE_URL}/train-images-idx3-ubyte.gz",
    "train-labels-idx1-ubyte.gz": f"{BASE_URL}/train-labels-idx1-ubyte.gz",
    "t10k-images-idx3-ubyte.gz": f"{BASE_URL}/t10k-images-idx3-ubyte.gz",
    "t10k-labels-idx1-ubyte.gz": f"{BASE_URL}/t10k-labels-idx1-ubyte.gz",
}


def download(url: str, dest: Path) -> None:
    if dest.exists():
        print(f"[skip] {dest} already exists")
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"[download] {url} -> {dest}", end="", flush=True)
    urllib.request.urlretrieve(url, dest)
    print(" OK")


def gunzip(src: Path, dst: Path) -> None:
    if dst.exists():
        print(f"[skip] {dst} already exists")
        return
    print(f"[gunzip] {src} -> {dst}", end="", flush=True)
    with gzip.open(src, "rb") as f_in, open(dst, "wb") as f_out:
        shutil.copyfileobj(f_in, f_out)
    print(" OK")


def convert_dataset(images_path: Path, labels_path: Path, out_path: Path) -> None:
    print(f"[convert] {images_path.name} + {labels_path.name} -> {out_path}")
    images = idx2numpy.convert_from_file(str(images_path)) / 255.0
    labels = idx2numpy.convert_from_file(str(labels_path))
    new_labels = np.zeros((len(labels), 10))
    new_labels[np.arange(len(labels)), labels] = 1
    dataset = list(zip(images, new_labels, strict=False))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "wb") as f:
        np.save(f, np.array(dataset, dtype=object))
    print(f"  saved {len(dataset)} samples to {out_path}")


def ensure_mnist_weights(path: Path, layer_sizes=(784, 32, 16, 10)) -> None:
    if path.exists():
        print(f"[skip] {path} already exists")
        return
    print(f"[init] creating {path} with layers {layer_sizes}")
    from MultiLayerPerceptron.main import MultiLayerPerceptron

    mlp = MultiLayerPerceptron(list(layer_sizes))
    mlp.save(str(path))
    print("  OK")


def main() -> None:
    parser = argparse.ArgumentParser(description="Download MNIST and prepare npy files")
    parser.add_argument("--raw-dir", default="data/raw", help="directory for raw idx files")
    parser.add_argument("--force", action="store_true", help="re-download even if exists")
    args = parser.parse_args()

    raw_dir = Path(args.raw_dir)
    if args.force:
        for name in FILES:
            p = raw_dir / name
            if p.exists():
                p.unlink()
            ungz = raw_dir / name.replace(".gz", "")
            if ungz.exists():
                ungz.unlink()

    # 1. download
    for name, url in FILES.items():
        download(url, raw_dir / name)

    # 2. gunzip
    for name in FILES:
        gunzip(raw_dir / name, raw_dir / name.replace(".gz", ""))

    # 3. convert to npy expected by main.py
    convert_dataset(
        raw_dir / "train-images-idx3-ubyte",
        raw_dir / "train-labels-idx1-ubyte",
        Path("data/train.npy"),
    )
    convert_dataset(
        raw_dir / "t10k-images-idx3-ubyte",
        raw_dir / "t10k-labels-idx1-ubyte",
        Path("data/test.npy"),
    )

    # 4. ensure initial weights
    ensure_mnist_weights(Path("mnist.npy"))

    print("[done] all files ready")


if __name__ == "__main__":
    main()