# Load data
# Note: For the documentation interactive mode, we only query 20000 records from the dataset.  Feel free to remove for a more realistic analysis.
data = (
    mapper
    .loader
    .from_huggingface("oscur/pluto", number_of_rows=20000, streaming=True)
    .with_columns("longitude", "latitude")
#    .with_columns(geometry_column=<geometry_column_name>") # Replace <geometry_column_name> with the actual name of your geometry column instead of latitude and longitude columns.
    .load()
    # From the loader module, from the following file within the OSCUR HuggingFace datasets hub and with the `longitude` and `latitude` or only with `geometry`
)    

# Create an urban layer (needed for the imputer)
# See further in the urban_layer example at examples/Basics/urban_layer.ipynb
layer = (
    mapper.urban_layer.with_type("streets_intersections") # From the urban layer module and with the type streets_intersections
    .from_place("Downtown Brooklyn, New York City, USA") # From place
    .build()
)

print(f"[Before Impute] Number of missing values in the longitude column: {data['longitude'].isnull().sum()}")
print(f"[Before Impute] Number of missing values in the latitude column: {data['latitude'].isnull().sum()}")

# Apply the imputer
imputed_data = (
    mapper
    .imputer # From the imputer module
    .with_type("SimpleGeoImputer") # With the type SimpleGeoImputer
    .on_columns(longitude_column="longitude", latitude_column="latitude") # On the columns longitude and latitude
#    .on_columns(geometry_column=<geometry_column_name>") # Replace <geometry_column_name> with the actual name of your geometry column instead of latitude and longitude columns.      
    .transform(data, layer)  # All imputers require access to the urban layer in case they need to extract information from it.
)

print(f"[After Impute] Number of missing values in the longitude column: {imputed_data['longitude'].isnull().sum()}")
print(f"[After Impute] Number of missing values in the latitude column: {imputed_data['latitude'].isnull().sum()}")

imputed_data

print(mapper.imputer.preview())

# Load CSV data
data1 = (
    mapper
    .loader
    .from_huggingface("oscur/pluto", number_of_rows=1000, streaming=True)
    .with_columns("longitude", "latitude")
#    .with_columns(geometry_column=<geometry_column_name>") # Replace <geometry_column_name> with the actual name of your geometry column instead of latitude and longitude columns.
    .load()
    # From the loader module, from the following file and with the `longitude` and `latitude` or only `geometry`
)

# Load Parquet data
data2 = (
    mapper
    .loader
    .from_huggingface("oscur/taxisvis1M", number_of_rows=1000, streaming=True) # To update with your own path
    .with_columns("pickup_longitude", "pickup_latitude").load() # Inform your long and lat columns
#    .with_columns(geometry_column=<geometry_column_name>") # Replace <geometry_column_name> with the actual name of your geometry column instead of latitude and longitude columns.      
)

data = {
  "pluto_data": data1,
  "taxi_data": data2,
}

# Apply the imputer.
# If the same imputer is applied to all datasets, and longitude_column/latitude_column have different names in each dataset, you can use loader.with_map 
# to map columns, standardizing the column names
imputed_data = (
    mapper
    .imputer # From the imputer module
    .with_type("SimpleGeoImputer") # With the type SimpleGeoImputer
    .on_columns(longitude_column="longitude", latitude_column="latitude") # On the columns longitude and latitude
#    .on_columns(geometry_column=<geometry_column_name>") # Replace <geometry_column_name> with the actual name of your geometry column instead of latitude and longitude columns.      
    .with_data(data_id="pluto_data") # On a specific data from the dictionary
    .transform(data, layer)  # All imputers require access to the urban layer in case they need to extract information from it.
)