import torch
import joblib

checkpoint = torch.load("models/firetrace_model.pt") # Directories may be different depending on specific environment.
model = FiretraceMLP(width=checkpoint['model_size'][0], depth=checkpoint['model_size'][1])

# OPTIONAL:
# compiled_model = torch.compile(model, fullgraph=True, mode="max-autotune")
# and use compiled_model instead of model. Note that this is not compatible with some environments.

model.load_state_dict(checkpoint['model_state_dict'])
model.eval() # Prepare model for inference

# Get your input data. You may choose to do this via a numpy array. Here's how we do it.
# It's quite important that the data is presented in the same order we have here.
inputs = [soi, max_t_bne, max_t_mel, max_t_cns, max_t_pth, max_t_syd, sin_month, cos_month, year]

# Prepare scalers
x_scaler = joblib.load("models/x_scaler.save")
y_scaler = joblib.load("models/y_scaler.save")

# Scale inputs (while converting to numpy array)
scaled_inputs = x_scaler.transform(np.array(inputs).reshape(1, -1))

model_output = model(torch.tensor(scaled_inputs).float())

# Interpret the following tensor as you wish.
scaled_output = y_scaler.inverse_transform(model_output.detach().numpy().reshape(1, -1))