Maze Information Prediction

"""
- Input: flatten the 30x30 maze into a one-dimensional vector of length 900 (each cell is mapped to 0~4)
- Network: 900 -> 4096 -> 512 -> 4 (ReLU)
- Training: regression (MSE) on train_data.csv + train_answer.csv, Adam, mini-batches (PyTorch)
- Output: predict for test_data.csv and write result.csv (4 real numbers per row)

Usage:
    python baseline.py train_data.csv train_answer.csv test_data.csv result.csv
"""

import sys
from pathlib import Path
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import zipfile
import os
import random

seed = 42

random.seed(seed)                  # Python built-in random
np.random.seed(seed)               # NumPy
torch.manual_seed(seed)            # PyTorch (CPU)
torch.cuda.manual_seed(seed)       # PyTorch (single GPU)
torch.cuda.manual_seed_all(seed)   # PyTorch (all GPUs)

# Ensures deterministic behavior
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

N = 30
D = N * N

# ===== Data reading and encoding =====

def encode_lines(path: str) -> np.ndarray:
    # Encode each line of 900 characters as a vector of values 0~4:
    # '.'->0, '# '->1, '?'->2, 'S'->3, 'T'->4
    char_id = {".": 0, "#": 1, "?": 2, "S": 3, "T": 4}
    with open(path, "r", encoding="utf-8-sig") as f:
        xs = [[char_id[c] for c in line.strip()] for line in f if line.strip()]
    return np.asarray(xs, dtype=np.float32)


def read_y(path: str) -> np.ndarray:
    # Read the 4 label columns (integers); treated as float during training
    return np.loadtxt(path, delimiter=",", dtype=np.float32, encoding="utf-8-sig")


# ===== Writing the results =====

def write_result(path: Path, pred: np.ndarray) -> None:
    np.savetxt(path, pred, delimiter=",", fmt="%.6f")


# ===== Model training =====

def train_model(train_x_path: str, train_y_path: str, epochs: int) -> nn.Module:
    # Inputs are scaled to about [0,1]; labels are divided by 900 for the regression and multiplied back to the original scale at prediction time.
    x_train = torch.from_numpy(encode_lines(train_x_path) / 4.0)
    y_train = torch.from_numpy(read_y(train_y_path) / 900.0)

    torch.manual_seed(0)
    model = nn.Sequential(
        nn.Linear(D, 4096), nn.ReLU(),
        nn.Linear(4096, 512), nn.ReLU(),
        nn.Linear(512, 4),
    )
    loader = DataLoader(TensorDataset(x_train, y_train), batch_size=64, shuffle=True)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)

    model.train()
    for ep in range(1, epochs + 1):
        total = 0.0
        for xb, yb in loader:
            pred = model(xb)
            loss = nn.functional.mse_loss(pred, yb)
            opt.zero_grad()
            loss.backward()
            opt.step()
            total += float(loss.item()) * xb.shape[0]
        print(f"epoch {ep}/{epochs}  mse={total / len(x_train):.6f}")
    return model


# ===== Prediction =====

def predict(model: nn.Module, test_x_path: str) -> np.ndarray:
    x_test = torch.from_numpy(encode_lines(test_x_path) / 4.0)
    model.eval()
    with torch.no_grad():
        return model(x_test).cpu().numpy() * 900.0


# ===== Main workflow =====


TRAIN_PATH = "/bohr/train-4mzz/v1/"  # training set path


# Training set
train_x_path = TRAIN_PATH + "train_data.csv"
train_y_path = TRAIN_PATH + "train_answer.csv"



out_path = Path("result.csv")
epochs = 3

model = train_model(train_x_path, train_y_path, epochs)




# Save the model weights and the necessary scaling information so that the predictions can be reproduced.
torch.save(
    {
        "state_dict": model.state_dict(),
        "epochs": epochs,
        "architecture": "900-4096-512-4",
        "input_scale": 4.0,
        "output_scale": 900.0,
    },
    out_path.with_name("model.pt"),
)

if os.environ.get("DATA_PATH"):
    DATA_PATH = os.environ.get("DATA_PATH") + "/"  # test set path
else:
    DATA_PATH = "/bohr/mazeval-7zx2/v1/"  # fallback for local testing

# Test set
testA_path = DATA_PATH + "val_data.csv"
testB_path = DATA_PATH + "test_data.csv"

#Predict each set separately
pred_A = predict(model, testA_path)
pred_B = predict(model, testB_path)

#Combine the prediction results

submissionA = pd.DataFrame(pred_A)
submissionA.to_csv("./submission_val.csv", index=False, header=False)

submissionB = pd.DataFrame(pred_B)
submissionB.to_csv("./submission_test.csv", index=False, header=False)

files_to_zip = ['./submission_val.csv', './submission_test.csv']
zip_filename = 'submission.zip'

with zipfile.ZipFile(zip_filename, 'w') as zipf:
    for file in files_to_zip:
        zipf.write(file, os.path.basename(file))

print(f'{zip_filename} is created succefully!')