# It work in my personal laptop but no in the Valencia server, IDK why :'(
import torch
import benchmark
import opensr_test
import matplotlib.pyplot as plt
from sr4rs.utils import load_cesbio_sr, run_sr4rs

# Load the model
model = load_cesbio_sr()

# Load the dataset
dataset = opensr_test.load("naip")
lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"]

# Predict a image
results = run_sr4rs(
    model=model,
    lr=lr_dataset[2],
    hr=hr_dataset[2],
)

# Display the results
fig, ax = plt.subplots(1, 3, figsize=(10, 5))
ax[0].imshow(results["lr"].transpose(1, 2, 0)/3000)
ax[0].set_title("LR")
ax[0].axis("off")
ax[1].imshow(results["sr"].transpose(1, 2, 0)/3000)
ax[1].set_title("SR")
ax[1].axis("off")
ax[2].imshow(results["hr"].transpose(1, 2, 0) / 3000)
ax[2].set_title("HR")
plt.show()


# Run the experiment
#benchmark.create_geotiff(model, run_sr4rs, "all", "sr4rs/", force=True)