import geopandas as gpd
import matplotlib.pyplot as plt
from geopy.geocoders import Nominatim
import requests
from bs4 import BeautifulSoup
import re
from tqdm import tqdm
from shapely.geometry import Point, box
import json
import pandas as pd
import numpy as np
import rasterio
import warnings
from scipy.spatial import cKDTree
from rasterio.transform import from_origin
from rasterio.features import rasterize
from rasterio.warp import reproject, Resampling, calculate_default_transform
import os
import rasterio.mask
from rasterio.io import MemoryFile
import matplotlib.colors as mcolors
import matplotlib.patheffects as patheffects
from pykrige.uk import UniversalKriging
from xgboost import XGBRegressor  # Integrated XGBoost
from sklearn.model_selection import GridSearchCV, cross_validate
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score, median_absolute_error, explained_variance_score
from sklearn.pipeline import Pipeline
from datetime import datetime
from scipy.ndimage import gaussian_filter, distance_transform_edt
import pytz
import glob
import math
import joblib  # For model saving/loading
from matplotlib.font_manager import FontProperties

# Set the target CRS
TARGET_CRS = "EPSG:2180"

# Set environment variable for PROJ_LIB if necessary
# Uncomment and set the correct path if you encounter PROJ_LIB errors
# os.environ['PROJ_LIB'] = r'C:\Users\Lenovo\Documents\conda\envs\temp_env\Library\share\proj'

warnings.filterwarnings("ignore")

# Keywords to remove
keywords_to_remove = [
    r"(centrum)", "rondo", "most", "Obwodnica", "obwodnica I", "obwodnica II", "obwodnica",
    "- Zawiszyn", "OUD Węzeł", "PPO", "Węzeł", "N (wąwóz)", "S (most)", "stacja", "z pomiarów", "min", "max", "średnia"
]

def clean_station_name(station_name):
    pattern = re.compile("|".join(keywords_to_remove), re.IGNORECASE)
    cleaned_name = pattern.sub("", station_name).strip()
    return cleaned_name

def get_masovian_boundary():
    url = "https://raw.githubusercontent.com/ppatrzyk/polska-geojson/master/wojewodztwa/wojewodztwa-medium.geojson"
    # Fetch the content using requests
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
        # Set the CRS to EPSG:4326 when creating the GeoDataFrame
        gdf = gpd.GeoDataFrame.from_features(data["features"], crs="EPSG:4326")
        masovian = gdf[gdf['nazwa'] == 'mazowieckie']
        # Reproject to TARGET_CRS
        masovian = masovian.to_crs(TARGET_CRS)
        return masovian
    else:
        print(f"Error fetching GeoJSON data: {response.status_code}")
        return None

# Global variable for Masovian boundary
masovian_boundary = get_masovian_boundary()

def is_in_masovian(lat, lon):
    point = gpd.GeoSeries([Point(lon, lat)], crs="EPSG:4326").to_crs(TARGET_CRS)
    masovian_geom = masovian_boundary.unary_union
    return masovian_geom.contains(point.buffer(100).iloc[0])

def fetch_traxelektronik_data():
    url = 'https://www.traxelektronik.pl/pogoda/zbiorcza.php?RejID=10'
    response = requests.get(url)
    soup = BeautifulSoup(response.text, 'html.parser')

    stations = []
    temperatures = []

    table_rows = soup.find_all('tr')[2:]

    for row in table_rows:
        columns = row.find_all('td')
        if len(columns) > 1:
            station = columns[0].get_text(strip=True)
            temperature = columns[1].get_text(strip=True)
            if temperature != "-":
                stations.append(station)  # Keep original station names
                temperatures.append(temperature)

    return stations, temperatures

def fetch_netatmo_data():
    url = "https://api.netatmo.com/api/getpublicdata"
    params = {
        "lat_ne": 53,
        "lon_ne": 22,
        "lat_sw": 51,
        "lon_sw": 19,
        "required_data": "temperature",
        "filter": "true",
        "access_token": "66e1e80975134797710422cd|33aaf2b32e9c1848f4588e1fd32d7568"
    }

    response = requests.get(url, params=params)
    if response.status_code != 200:
        print(f"Error fetching Netatmo data: {response.status_code}")
        return [], [], []
    data = json.loads(response.text)

    stations = []
    temperatures = []
    coordinates = []

    # List of cities to exclude
    excluded_cities = ["Płońsk", "Łochów"]

    for station in data.get('body', []):
        place = station.get('place', {})
        location = place.get('location')
        if location:
            lon, lat = location
            if is_in_masovian(lat, lon):
                city = place.get('city', 'Unknown')
                if city in excluded_cities:
                    # Skip this station
                    continue
                for module, module_data in station.get('measures', {}).items():
                    if 'temperature' in module_data.get('type', []):
                        res = module_data.get('res', {})
                        if res:
                            latest_timestamp = max(res.keys())
                            temperature = res[latest_timestamp][0]
                            stations.append(city)
                            temperatures.append(temperature)
                            coordinates.append((lat, lon))
                            break

    return stations, temperatures, coordinates

def fetch_imgw_data():
    url = "https://rafalraczynski.com.pl/imgw/dane-imgw/getJSON.php?type=table&province=14&sort=temp&order=asc"
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
    else:
        return [], [], []

    target_stations = [
        "ANDRZEJEWO", "BORKOWO", "CZARNOWO", "KAZANÓW",
        "RUSZKOWICE", "RYBIENKO", "RZĄŚNIK WŁOŚCIAŃSKI", "WIELGOLAS"
    ]

    stations = []
    temperatures = []
    station_ids = []

    for item in data:
        statName = item.get('statName')
        temp = item.get('temp')
        if statName and temp and statName.upper() in target_stations:
            stations.append(statName)
            temperatures.append(temp)
            station_ids.append(item.get('statId', 'Unknown'))

    return stations, temperatures, station_ids

def get_coordinates(station_name, geolocator):
    try:
        location = geolocator.geocode(f"{station_name}, Masovian Voivodeship, Poland")
        if location:
            if is_in_masovian(location.latitude, location.longitude):
                # Reproject the point to TARGET_CRS
                point = gpd.GeoSeries([Point(location.longitude, location.latitude)], crs="EPSG:4326").to_crs(TARGET_CRS).iloc[0]
                return (point.y, point.x), False
            else:
                return None, False
        return None, False
    except Exception as e:
        return None, False

def clean_temperature(temp):
    cleaned_temp = re.sub(r'[↓↑\s]', '', temp)
    try:
        return float(cleaned_temp)
    except ValueError:
        return None

def remove_nearby_duplicates(trax_coords, netatmo_coords, distance_threshold=50):  # distance_threshold in meters
    trax_points = np.array([(lon, lat) for lat, lon in trax_coords.values()])
    netatmo_points = np.array([(lon, lat) for lat, lon in netatmo_coords])
    trax_tree = cKDTree(trax_points)

    netatmo_filtered_coords = []
    netatmo_filtered_temperatures = []
    netatmo_filtered_stations = []

    for idx, (netatmo_point, station, temperature) in enumerate(zip(netatmo_points, netatmo_coords, netatmo_temperatures)):
        distance, _ = trax_tree.query(netatmo_point)
        if distance > distance_threshold:
            netatmo_filtered_coords.append(tuple(netatmo_point))
            netatmo_filtered_temperatures.append(temperature)
            netatmo_filtered_stations.append(station)

    return netatmo_filtered_coords, netatmo_filtered_temperatures, netatmo_filtered_stations

def hillshade(array, azimuth=315, angle_altitude=45):
    # Convert angles to radians
    azimuth_rad = np.radians(azimuth)
    altitude_rad = np.radians(angle_altitude)

    # Calculate gradients in x and y directions
    x, y = np.gradient(array)

    # Calculate the slope and aspect
    slope = np.pi/2.0 - np.arctan(np.sqrt(x*x + y*y))
    aspect = np.arctan2(-x, y)

    # Calculate shaded relief
    shaded = np.sin(altitude_rad) * np.sin(slope) + np.cos(altitude_rad) * np.cos(slope) * np.cos(azimuth_rad - aspect)

    # Normalize the shaded relief
    shaded = (shaded + 1) / 2
    shaded = (shaded * 255).astype(np.uint8)

    return shaded

def get_elevation_at_points(geometry, dem_data=None, dem_transform=None):
    from rasterio.transform import rowcol

    if dem_data is not None and dem_transform is not None:
        elevations = []
        for point in geometry:
            x, y = point.x, point.y
            try:
                row, col = rowcol(dem_transform, x, y)
                row = int(row)
                col = int(col)
                elevation = dem_data[row, col]
                elevations.append(elevation)
            except (IndexError, ValueError):
                elevations.append(np.nan)
        return elevations
    else:
        raise ValueError("Provide dem_data and dem_transform.")

def get_raster_value_at_points(geometry, raster_data=None, raster_transform=None):
    from rasterio.transform import rowcol

    if raster_data is not None and raster_transform is not None:
        values = []
        for point in geometry:
            x, y = point.x, point.y
            try:
                row, col = rowcol(raster_transform, x, y)
                row = int(row)
                col = int(col)
                if (0 <= row < raster_data.shape[0]) and (0 <= col < raster_data.shape[1]):
                    value = raster_data[row, col]
                    # Check for NaN in raster data
                    if np.isnan(value):
                        values.append(np.nan)
                    else:
                        values.append(value)
                else:
                    values.append(np.nan)
            except (IndexError, ValueError):
                values.append(np.nan)
        return values
    else:
        raise ValueError("Provide raster_data and raster_transform.")

def get_water_body_at_points(geometry, water_data=None, water_transform=None):
    from rasterio.transform import rowcol

    if water_data is not None and water_transform is not None:
        water_bodies = []
        for point in geometry:
            x, y = point.x, point.y
            try:
                row, col = rowcol(water_transform, x, y)
                row = int(row)
                col = int(col)
                if (0 <= row < water_data.shape[0]) and (0 <= col < water_data.shape[1]):
                    value = water_data[row, col]
                    # Adjust the condition based on your raster's water encoding
                    # For binary raster where 1 = water, 0 = land
                    water_bodies.append(1 if value else 0)
                else:
                    water_bodies.append(0)
            except (IndexError, ValueError):
                water_bodies.append(0)
        return water_bodies
    else:
        raise ValueError("Provide water_data and water_transform.")

def read_and_clip_dem(masovian_boundary):
    dem_path = 'input/masovian_dem.tif'
    with rasterio.open(dem_path) as src:
        # Reproject DEM to TARGET_CRS
        transform, width, height = calculate_default_transform(
            src.crs, TARGET_CRS, src.width, src.height, *src.bounds)
        dem_data = np.empty((height, width), dtype=src.meta['dtype'])
        reproject(
            source=rasterio.band(src, 1),
            destination=dem_data,
            src_transform=src.transform,
            src_crs=src.crs,
            dst_transform=transform,
            dst_crs=TARGET_CRS,
            resampling=Resampling.nearest)
        dem_transform = transform

        # Update the profile for the reprojected DEM
        dem_profile = src.profile.copy()
        dem_profile.update({
            'crs': TARGET_CRS,
            'transform': dem_transform,
            'width': width,
            'height': height,
            'nodata': np.nan,
            'count': 1
        })

        # Create an in-memory raster with the reprojected DEM data
        with MemoryFile() as memfile:
            with memfile.open(**dem_profile) as dataset:
                dataset.write(dem_data, 1)
                
                # Attempt to mask the reprojected DEM with the Masovian boundary
                try:
                    dem_clipped, dem_clipped_transform = rasterio.mask.mask(
                        dataset, masovian_boundary.geometry, crop=True, filled=True, nodata=np.nan, all_touched=True)
                    dem_clipped = dem_clipped[0]
                    return dem_clipped, dem_clipped_transform
                except ValueError as ve:
                    print(f"Error during masking DEM: {ve}")
                    return None, None

def read_and_clip_forest(masovian_boundary):
    forest_path = 'input/forests.tif'
    with rasterio.open(forest_path) as src:
        # Reproject Forest Cover to TARGET_CRS
        transform, width, height = calculate_default_transform(
            src.crs, TARGET_CRS, src.width, src.height, *src.bounds)
        forest_data = np.empty((height, width), dtype=src.meta['dtype'])
        reproject(
            source=rasterio.band(src, 1),
            destination=forest_data,
            src_transform=src.transform,
            src_crs=src.crs,
            dst_transform=transform,
            dst_crs=TARGET_CRS,
            resampling=Resampling.nearest)
        forest_transform = transform

        # Update the profile for the reprojected Forest Cover
        forest_profile = src.profile.copy()
        forest_profile.update({
            'crs': TARGET_CRS,
            'transform': forest_transform,
            'width': width,
            'height': height,
            'nodata': 0,
            'count': 1
        })

        # Create an in-memory raster with the reprojected Forest Cover data
        with MemoryFile() as memfile:
            with memfile.open(**forest_profile) as dataset:
                dataset.write(forest_data, 1)
                
                # Attempt to mask the reprojected Forest Cover with the Masovian boundary
                try:
                    forest_clipped, forest_clipped_transform = rasterio.mask.mask(
                        dataset, masovian_boundary.geometry, crop=True, filled=True, nodata=0, all_touched=True)
                    forest_clipped = forest_clipped[0]
                    # Replace nodata values with 0
                    forest_clipped = np.where(forest_clipped == 0, 0, forest_clipped)
                    return forest_clipped, forest_clipped_transform
                except ValueError as ve:
                    print(f"Error during masking Forest Cover: {ve}")
                    return None, None

def read_and_clip_water_bodies(masovian_boundary):
    water_path = 'input/water_bodies.tif'
    with rasterio.open(water_path) as src:
        # Reproject Water Bodies to TARGET_CRS
        transform, width, height = calculate_default_transform(
            src.crs, TARGET_CRS, src.width, src.height, *src.bounds)
        water_data = np.empty((height, width), dtype=src.meta['dtype'])
        reproject(
            source=rasterio.band(src, 1),
            destination=water_data,
            src_transform=src.transform,
            src_crs=src.crs,
            dst_transform=transform,
            dst_crs=TARGET_CRS,
            resampling=Resampling.nearest)
        water_transform = transform

        # Update the profile for the reprojected Water Bodies
        water_profile = src.profile.copy()
        water_profile.update({
            'crs': TARGET_CRS,
            'transform': water_transform,
            'width': width,
            'height': height,
            'nodata': 0,
            'count': 1
        })

        # Create an in-memory raster with the reprojected Water Bodies data
        with MemoryFile() as memfile:
            with memfile.open(**water_profile) as dataset:
                dataset.write(water_data, 1)
                
                # Attempt to mask the reprojected Water Bodies with the Masovian boundary
                try:
                    water_clipped, water_clipped_transform = rasterio.mask.mask(
                        dataset, masovian_boundary.geometry, crop=True, filled=True, nodata=0, all_touched=True)
                    water_clipped = water_clipped[0]
                    # Replace nodata values with 0
                    water_clipped = np.where(water_clipped == 0, 0, water_clipped)
                    return water_clipped, water_clipped_transform, dataset.crs
                except ValueError as ve:
                    print(f"Error during masking Water Bodies: {ve}")
                    return None, None, None

def read_and_clip_ecostress(masovian_boundary):
    ecostress_path = 'input/raw_ecostress.tif'
    with rasterio.open(ecostress_path) as src:
        # Reproject ECOSTRESS to TARGET_CRS
        transform, width, height = calculate_default_transform(
            src.crs, TARGET_CRS, src.width, src.height, *src.bounds)
        ecostress_data = np.empty((height, width), dtype=src.meta['dtype'])
        reproject(
            source=rasterio.band(src, 1),
            destination=ecostress_data,
            src_transform=src.transform,
            src_crs=src.crs,
            dst_transform=transform,
            dst_crs=TARGET_CRS,
            resampling=Resampling.nearest)
        ecostress_transform = transform

        # Update the profile for the reprojected ECOSTRESS
        ecostress_profile = src.profile.copy()
        ecostress_profile.update({
            'crs': TARGET_CRS,
            'transform': ecostress_transform,
            'width': width,
            'height': height,
            'nodata': np.nan,
            'count': 1
        })

        # Create an in-memory raster with the reprojected ECOSTRESS data
        with MemoryFile() as memfile:
            with memfile.open(**ecostress_profile) as dataset:
                dataset.write(ecostress_data, 1)
                
                # Attempt to mask the reprojected ECOSTRESS with the Masovian boundary
                try:
                    ecostress_clipped, ecostress_clipped_transform = rasterio.mask.mask(
                        dataset, masovian_boundary.geometry, crop=True, filled=True, nodata=np.nan, all_touched=True)
                    ecostress_clipped = ecostress_clipped[0]
                    return ecostress_clipped, ecostress_clipped_transform
                except ValueError as ve:
                    print(f"Error during masking ECOSTRESS: {ve}")
                    return None, None

def plot_stations_with_boundary_and_interpolation(
        trax_stations, trax_temperatures, trax_coords,
        netatmo_stations, netatmo_temperatures, netatmo_coords,
        imgw_stations, imgw_temperatures, imgw_coords,
        modis_day_folder, modis_night_folder):
    # Define a function to compute the cooling effect of water bodies based on the hour of the day
    def compute_water_body_cooling_effect(hour):
        """
        Computes the cooling effect of water bodies based on the hour of the day.
        Cooling effect varies between -1°C during the day to -4°C at night.
        """
        # Normalize hour to range 0 to 24
        hour = hour % 24
        # Use a cosine function to model the cooling effect
        # Peak cooling at midnight (hour=0), minimal cooling at noon (hour=12)
        cooling_effect = -3.0 - 2.0 * math.cos((hour / 24) * 2 * math.pi)
        return cooling_effect

    # Determine the current time in CET
    cet = pytz.timezone('CET')
    now_cet = datetime.now(cet)
    hour = now_cet.hour + now_cet.minute / 60.0

    # Define daytime and nighttime hours
    day_start = 6
    day_end = 18
    is_daytime = day_start <= hour <= day_end

    # List of Traxelektronik stations to exclude during daytime
    exclude_trax_stations_daytime = [
        'Ceranów', 'Baranów', 'Drobin', 'Kamion',
        'Mińsk Mazowiecki', 'Mostówka', 'Paplin', 'Pilawa',
        'Stromiec', 'Węzeł Łabiszyńska', 'Węzeł Marki',
        'Węzeł Marynarska', 'Wilchta', 'Wiśniew'
    ]
    # Clean the exclude list for matching
    exclude_trax_stations_daytime_cleaned = [clean_station_name(name) for name in exclude_trax_stations_daytime]

    # Map station names to cleaned station names
    cleaned_station_dict = {station: clean_station_name(station) for station in trax_stations}

    # Process Traxelektronik data
    trax_filtered_stations = []
    trax_filtered_temperatures = []
    trax_geometries = []

    # Filter and clean Traxelektronik data
    for station, temperature in zip(trax_stations, trax_temperatures):
        cleaned_station = cleaned_station_dict[station]
        # Exclude certain stations during daytime
        if is_daytime and cleaned_station in exclude_trax_stations_daytime_cleaned:
            continue  # Skip this station
        cleaned_temperature = clean_temperature(temperature)
        if cleaned_station in trax_coords and cleaned_temperature is not None:
            lat, lon = trax_coords[cleaned_station]
            trax_filtered_stations.append(cleaned_station)
            trax_filtered_temperatures.append(cleaned_temperature)
            trax_geometries.append(Point(lon, lat))

    # Create GeoDataFrame for Traxelektronik data
    trax_gdf = gpd.GeoDataFrame({
        'Station': trax_filtered_stations,
        'Temperature': trax_filtered_temperatures,
        'Source': 'Traxelektronik'
    }, geometry=trax_geometries, crs=TARGET_CRS)

    # Handle Netatmo data
    netatmo_filtered_coords = []
    netatmo_filtered_temperatures = []
    netatmo_filtered_stations = []

    for station, temperature, coord in zip(netatmo_stations, netatmo_temperatures, netatmo_coords):
        if isinstance(temperature, (int, float)) and coord is not None:
            netatmo_filtered_coords.append(coord)
            netatmo_filtered_temperatures.append(temperature)
            netatmo_filtered_stations.append(station)

    if len(netatmo_filtered_coords) > 0:
        netatmo_geometries = [Point(lon, lat) for lat, lon in netatmo_filtered_coords]
        netatmo_gdf = gpd.GeoDataFrame({
            'Station': netatmo_filtered_stations,
            'Temperature': netatmo_filtered_temperatures,
            'Source': 'Netatmo'
        }, geometry=netatmo_geometries, crs=TARGET_CRS)
    else:
        netatmo_gdf = gpd.GeoDataFrame(columns=['Station', 'Temperature', 'Source', 'geometry'], crs=TARGET_CRS)

    # Handle IMGW data
    imgw_filtered_stations = []
    imgw_filtered_temperatures = []
    imgw_geometries = []

    # Clean and process IMGW data
    for station, temperature in zip(imgw_stations, imgw_temperatures):
        cleaned_temperature = clean_temperature(temperature)
        if station in imgw_coords and cleaned_temperature is not None:
            lat, lon = imgw_coords[station]
            imgw_filtered_stations.append(station)
            imgw_filtered_temperatures.append(cleaned_temperature)
            imgw_geometries.append(Point(lon, lat))

    # Create GeoDataFrame for IMGW data
    imgw_gdf = gpd.GeoDataFrame({
        'Station': imgw_filtered_stations,
        'Temperature': imgw_filtered_temperatures,
        'Source': 'IMGW'
    }, geometry=imgw_geometries, crs=TARGET_CRS)
    
    # Add X and Y coordinates to imgw_gdf
    imgw_gdf['X'] = imgw_gdf.geometry.x
    imgw_gdf['Y'] = imgw_gdf.geometry.y

    # Read and clip DEM
    dem_data, dem_transform = read_and_clip_dem(masovian_boundary)

    # Read and clip Forest Cover
    forest_data, forest_transform = read_and_clip_forest(masovian_boundary)

    # Read and clip Water Bodies
    water_data, water_transform, water_crs = read_and_clip_water_bodies(masovian_boundary)

    # Read and clip ECOSTRESS data
    ecostress_data, ecostress_transform = read_and_clip_ecostress(masovian_boundary)

    # Ensure all GeoDataFrames are in TARGET_CRS
    trax_gdf = trax_gdf.to_crs(TARGET_CRS)
    netatmo_gdf = netatmo_gdf.to_crs(TARGET_CRS)
    imgw_gdf = imgw_gdf.to_crs(TARGET_CRS)

    # Get elevation at station points using dem_data and dem_transform
    trax_elevations = get_elevation_at_points(
        trax_gdf.geometry,
        dem_data=dem_data,
        dem_transform=dem_transform
    )
    netatmo_elevations = get_elevation_at_points(
        netatmo_gdf.geometry,
        dem_data=dem_data,
        dem_transform=dem_transform
    )
    imgw_elevations = get_elevation_at_points(
        imgw_gdf.geometry,
        dem_data=dem_data,
        dem_transform=dem_transform
    )

    # Add elevation to the GeoDataFrames
    trax_gdf['Elevation'] = trax_elevations
    netatmo_gdf['Elevation'] = netatmo_elevations
    imgw_gdf['Elevation'] = imgw_elevations

    # Get tree cover density at station points using forest_data and forest_transform
    trax_forest = get_raster_value_at_points(
        trax_gdf.geometry,
        raster_data=forest_data,
        raster_transform=forest_transform
    )
    netatmo_forest = get_raster_value_at_points(
        netatmo_gdf.geometry,
        raster_data=forest_data,
        raster_transform=forest_transform
    )
    imgw_forest = get_raster_value_at_points(
        imgw_gdf.geometry,
        raster_data=forest_data,
        raster_transform=forest_transform
    )

    # Add tree cover density to the GeoDataFrames
    trax_gdf['Tree_Cover'] = trax_forest
    netatmo_gdf['Tree_Cover'] = netatmo_forest
    imgw_gdf['Tree_Cover'] = imgw_forest

    # Get water body data at station points
    trax_water = get_water_body_at_points(
        trax_gdf.geometry,
        water_data=water_data,
        water_transform=water_transform
    )
    netatmo_water = get_water_body_at_points(
        netatmo_gdf.geometry,
        water_data=water_data,
        water_transform=water_transform
    )
    imgw_water = get_water_body_at_points(
        imgw_gdf.geometry,
        water_data=water_data,
        water_transform=water_transform
    )

    # Add water body data to the GeoDataFrames
    trax_gdf['Water_Body'] = trax_water
    netatmo_gdf['Water_Body'] = netatmo_water
    imgw_gdf['Water_Body'] = imgw_water

    # Compute distance to water bodies for station points
    water_mask_array = water_data.astype(bool)
    grid_spacing = water_transform.a  # Assuming square pixels
    distance_to_water_array = distance_transform_edt(~water_mask_array) * grid_spacing  # Distance in meters

    # Get distance to water bodies at station points
    trax_distance_to_water = get_raster_value_at_points(
        trax_gdf.geometry,
        raster_data=distance_to_water_array,
        raster_transform=water_transform
    )
    netatmo_distance_to_water = get_raster_value_at_points(
        netatmo_gdf.geometry,
        raster_data=distance_to_water_array,
        raster_transform=water_transform
    )
    imgw_distance_to_water = get_raster_value_at_points(
        imgw_gdf.geometry,
        raster_data=distance_to_water_array,
        raster_transform=water_transform
    )

    # Add distance to water to the GeoDataFrames
    trax_gdf['Distance_to_Water'] = trax_distance_to_water
    netatmo_gdf['Distance_to_Water'] = netatmo_distance_to_water
    imgw_gdf['Distance_to_Water'] = imgw_distance_to_water

    # Add 'Priority' column to each GeoDataFrame
    trax_gdf['Priority'] = 2
    netatmo_gdf['Priority'] = 1
    imgw_gdf['Priority'] = 3  # Higher priority

    # Combine GeoDataFrames
    stations_gdf = gpd.GeoDataFrame(pd.concat([trax_gdf, netatmo_gdf, imgw_gdf], ignore_index=True))
    stations_gdf.crs = TARGET_CRS
    stations_gdf = stations_gdf.dropna(subset=['Temperature', 'Elevation', 'Tree_Cover', 'Distance_to_Water'])

    # Extract coordinates
    stations_gdf['X'] = stations_gdf.geometry.x
    stations_gdf['Y'] = stations_gdf.geometry.y

    # Read ECOSTRESS data
    ecostress_data, ecostress_transform = read_and_clip_ecostress(masovian_boundary)

    # Define grid spacing
    grid_spacing = 250  # 250 meters

    # Get Masovian boundary bounds in TARGET_CRS
    x_min, y_min, x_max, y_max = masovian_boundary.total_bounds

    # Compute the number of cols and rows with the new grid_spacing
    ncols = int(np.ceil((x_max - x_min) / grid_spacing))
    nrows = int(np.ceil((y_max - y_min) / grid_spacing))

    # Create the affine transform for the grid
    transform = from_origin(x_min, y_max, grid_spacing, grid_spacing)

    # Create arrays of row and column indices
    cols = np.arange(ncols)
    rows = np.arange(nrows)

    # Create a meshgrid of indices
    cols_indices, rows_indices = np.meshgrid(cols, rows)

    # Flatten the indices
    cols_flat = cols_indices.ravel()
    rows_flat = rows_indices.ravel()

    # Get the x and y coordinates for each grid cell center
    xs_flat, ys_flat = rasterio.transform.xy(transform, rows_flat, cols_flat)

    # Create grid_df as a Pandas DataFrame
    grid_df = pd.DataFrame({
        'X': xs_flat,
        'Y': ys_flat,
        'row_idx': rows_flat,
        'col_idx': cols_flat
    })

    # Create geometry column
    grid_df['geometry'] = gpd.points_from_xy(grid_df['X'], grid_df['Y'])

    # Convert to GeoDataFrame
    grid_gdf = gpd.GeoDataFrame(grid_df, geometry='geometry', crs=TARGET_CRS)

    # Assign ECOSTRESS_LST to grid_gdf
    grid_gdf['ECOSTRESS_LST'] = get_raster_value_at_points(
        grid_gdf.geometry,
        raster_data=ecostress_data,
        raster_transform=ecostress_transform
    )

    # Handle missing ECOSTRESS_LST values
    ecostress_lst_mean_grid = np.nanmean(grid_gdf['ECOSTRESS_LST'])
    grid_gdf['ECOSTRESS_LST'].fillna(ecostress_lst_mean_grid, inplace=True)

    # Assign ECOSTRESS_LST to stations_gdf
    stations_gdf['ECOSTRESS_LST'] = get_raster_value_at_points(
        stations_gdf.geometry,
        raster_data=ecostress_data,
        raster_transform=ecostress_transform
    )

    # Handle missing ECOSTRESS_LST values in stations_gdf
    ecostress_lst_mean_stations = np.nanmean(stations_gdf['ECOSTRESS_LST'])
    stations_gdf['ECOSTRESS_LST'].fillna(ecostress_lst_mean_stations, inplace=True)

    # Verify 'ECOSTRESS_LST' in stations_gdf
    if 'ECOSTRESS_LST' not in stations_gdf.columns:
        print("Error: 'ECOSTRESS_LST' column is missing from stations_gdf.")
    else:
        print("'ECOSTRESS_LST' successfully added to stations_gdf.")
        # Optionally, inspect a few rows
        print(stations_gdf[['Station', 'X', 'Y', 'ECOSTRESS_LST']].head())

    # Continue with data preparation...

    # Prepare sample weights: Assign higher weight to IMGW stations
    weights_map = {'Netatmo': 1, 'Traxelektronik': 1, 'IMGW': 5}
    stations_gdf['Sample_Weight'] = stations_gdf['Source'].map(weights_map)
    sample_weights = stations_gdf['Sample_Weight'].values

    # Prepare training data
    X = stations_gdf[['X', 'Y', 'Elevation', 'Tree_Cover', 'Distance_to_Water', 'ECOSTRESS_LST']]

    # Replace infinite values with NaN
    X = X.replace([np.inf, -np.inf], np.nan)

    # Remove rows with NaN values in X
    valid_indices = ~X.isnull().any(axis=1)
    X = X[valid_indices]
    y = stations_gdf.loc[X.index, 'Temperature']
    sample_weights = sample_weights[valid_indices]

    # Proceed with model training...
    # Initialize XGBoost Regressor
    xgb = XGBRegressor(
        objective='reg:squarederror',
        random_state=42,
        n_jobs=-1,
        verbosity=0  # Silent mode
    )

    # Create a pipeline without PolynomialFeatures and StandardScaler
    pipeline = Pipeline([
        ('xgb', xgb)
    ])

    # Define the hyperparameter grid for XGBoost
    param_grid = {
        'xgb__n_estimators': [100, 200, 300],
        'xgb__learning_rate': [0.01, 0.05, 0.1],
        'xgb__max_depth': [4, 6, 8],
        'xgb__subsample': [0.6, 0.8, 1.0],
        'xgb__colsample_bytree': [0.6, 0.8, 1.0],
        'xgb__reg_alpha': [0, 0.1, 0.5],  # L1 regularization
        'xgb__reg_lambda': [1, 1.5, 2]    # L2 regularization
    }

    # Initialize GridSearchCV with 5-fold cross-validation
    grid_search = GridSearchCV(
        estimator=pipeline,
        param_grid=param_grid,
        cv=5,
        scoring='neg_mean_absolute_error',
        n_jobs=-1,
        verbose=1
    )

    # Fit GridSearchCV with sample weights
    grid_search.fit(X, y, **{'xgb__sample_weight': sample_weights})
    print("GridSearchCV completed.")

    # Retrieve the best model
    best_xgb_pipeline = grid_search.best_estimator_
    print(f"Best parameters found: {grid_search.best_params_}")

    # Define the metrics to evaluate
    scoring_metrics = {
        'MAE': 'neg_mean_absolute_error',
        'MSE': 'neg_mean_squared_error',
        'RMSE': 'neg_root_mean_squared_error',
        'R2': 'r2',
        'Median_Err': 'neg_median_absolute_error',
        'Exp_Variance': 'explained_variance'
    }

    # Perform cross-validation with multiple metrics
    cv_results = cross_validate(
        best_xgb_pipeline,
        X,
        y,
        cv=5,
        scoring=scoring_metrics,
        n_jobs=-1,
        return_train_score=False
    )
    print("Cross-validation completed.")

    # Convert negative scores to positive where necessary
    cv_metrics = {}
    for key in scoring_metrics.keys():
        if key in ['MAE', 'MSE', 'RMSE', 'Median_Err']:
            # Convert negative to positive
            cv_metrics[key] = -cv_results[f'test_{key}']
        else:
            # Keep as is (R² and Explained Variance)
            cv_metrics[key] = cv_results[f'test_{key}']

    # Calculate mean and standard deviation for each metric
    metric_summary = {}
    for metric, values in cv_metrics.items():
        mean_val = np.mean(values)
        std_val = np.std(values)
        metric_summary[metric] = (mean_val, std_val)

    # Print detailed performance metrics
    print("\nDetailed Cross-Validation Performance Metrics:")
    print("---------------------------------------------")
    for metric, (mean, std) in metric_summary.items():
        print(f"{metric}: {mean:.3f} ± {std:.3f}")
    print("---------------------------------------------")

    # Get feature importances from the best model
    xgb_model = best_xgb_pipeline.named_steps['xgb']
    feature_importances = xgb_model.feature_importances_
    feature_names = X.columns

    # Create a DataFrame for better visualization
    feature_importances_df = pd.DataFrame({'Feature': feature_names, 'Importance': feature_importances})
    feature_importances_df = feature_importances_df.sort_values('Importance', ascending=False)
    print("\nFeature importances from the XGBoost model:")
    print(feature_importances_df)

    # Create grid dataframe
    # Already converted to grid_gdf earlier

    # Running Model Progress Bar
    with tqdm(total=100, desc='Running Model', unit='%', ncols=80) as pbar:
        # Retrieve elevations from DEM
        grid_gdf['Elevation'] = get_elevation_at_points(
            grid_gdf.geometry,
            dem_data=dem_data,
            dem_transform=dem_transform
        )
        pbar.update(20)

        # Retrieve tree cover density from forest cover raster
        grid_gdf['Tree_Cover'] = get_raster_value_at_points(
            grid_gdf.geometry,
            raster_data=forest_data,
            raster_transform=forest_transform
        )
        pbar.update(20)

        # Calculate distance to water bodies for grid points
        grid_gdf['Distance_to_Water'] = get_raster_value_at_points(
            grid_gdf.geometry,
            raster_data=distance_to_water_array,
            raster_transform=water_transform
        )
        pbar.update(20)

        # Retrieve water body data at grid points
        grid_gdf['Water_Body'] = get_water_body_at_points(
            grid_gdf.geometry,
            water_data=water_data,
            water_transform=water_transform
        )
        pbar.update(10)

        # Include ECOSTRESS_LST in grid_gdf
        grid_gdf['ECOSTRESS_LST'] = get_raster_value_at_points(
            grid_gdf.geometry,
            raster_data=ecostress_data,
            raster_transform=ecostress_transform
        )
        # Handle missing ECOSTRESS_LST values
        grid_gdf['ECOSTRESS_LST'].fillna(ecostress_lst_mean_grid, inplace=True)
        pbar.update(10)

        # Remove grid points with NaN values (i.e., outside data coverage)
        grid_gdf = grid_gdf.dropna(subset=['Elevation', 'Tree_Cover', 'Distance_to_Water', 'Water_Body', 'ECOSTRESS_LST'])
        pbar.update(10)

        # Verify 'ECOSTRESS_LST' in grid_gdf
        if 'ECOSTRESS_LST' not in grid_gdf.columns:
            print("Error: 'ECOSTRESS_LST' column is missing from grid_gdf.")
        else:
            print("'ECOSTRESS_LST' successfully added to grid_gdf.")

        # Assign grid ECOSTRESS_LST as a feature
        # Proceed to model prediction
        # Prepare data for regression prediction
        X_grid = grid_gdf[['X', 'Y', 'Elevation', 'Tree_Cover', 'Distance_to_Water', 'ECOSTRESS_LST']]
        # Handle infinite values in X_grid
        X_grid = X_grid.replace([np.inf, -np.inf], np.nan)
        # Handle NaNs in X_grid
        X_grid = X_grid.fillna(method='ffill').fillna(method='bfill')

        # Ensure that all values are finite
        # Predict using the trained XGBoost model
        predictions = best_xgb_pipeline.predict(X_grid)
        grid_gdf['Predicted_Temperature'] = predictions
        pbar.update(10)

        # Calculate residuals at station locations
        stations_gdf['Predicted_Temperature'] = best_xgb_pipeline.predict(X)
        stations_gdf['Residual'] = stations_gdf['Temperature'] - stations_gdf['Predicted_Temperature']

        # Force residuals to zero at IMGW stations
        stations_gdf.loc[stations_gdf['Source'] == 'IMGW', 'Residual'] = 0

        # Adjust anomalously large residuals to reduce their impact
        lower_percentile = np.percentile(stations_gdf['Residual'], 2.5)
        upper_percentile = np.percentile(stations_gdf['Residual'], 97.5)

        # Cap the residuals
        stations_gdf['Adjusted_Residual'] = stations_gdf['Residual'].clip(lower=lower_percentile, upper=upper_percentile)

        # Remove duplicates by prioritizing IMGW data
        data = stations_gdf.copy()

        # Sort by priority (higher priority first)
        data = data.sort_values('Priority', ascending=False)

        # Group by coordinates and keep the first record in each group (highest priority)
        data = data.groupby(['X', 'Y'], as_index=False).first()

        # Perform Universal Kriging using Adjusted Residual
        uk = UniversalKriging(
            x=data['X'].values,
            y=data['Y'].values,
            z=data['Adjusted_Residual'].values,
            variogram_model='spherical',
            nlags=20,
            verbose=False
        )

        z, ss = uk.execute(
            'points',
            grid_gdf['X'].values,
            grid_gdf['Y'].values
        )

        grid_gdf['Residual'] = z
        grid_gdf['Kriging_Variance'] = ss

        # Handle NaNs in residuals by filling them with zero
        grid_gdf['Residual'].fillna(0, inplace=True)

        # Final temperature estimation
        grid_gdf['Temperature'] = grid_gdf['Predicted_Temperature'] + grid_gdf['Residual']
        pbar.update(10)

        # Save a copy of the temperature before ECOSTRESS adjustments
        grid_gdf['Temperature_No_ECOSTRESS'] = grid_gdf['Temperature'].copy()

        # Topographic Cooling Effect
        max_cooling_effect = 3.0  # Maximum cooling effect in degrees Celsius
        # Normalize elevation to range from 0 to 1
        elevation_min = grid_gdf['Elevation'].min()
        elevation_max = grid_gdf['Elevation'].max()
        grid_gdf['Normalized_Elevation'] = (grid_gdf['Elevation'] - elevation_min) / (elevation_max - elevation_min)

        # Compute the topographic cooling effect
        grid_gdf['Topographic_Cooling'] = grid_gdf['Normalized_Elevation'] * max_cooling_effect * (1 - is_daytime)

        # Ensure topographic cooling is not applied over water bodies
        grid_gdf.loc[grid_gdf['Water_Body'] == 1, 'Topographic_Cooling'] = 0

        # Apply the topographic cooling effect to the temperature
        grid_gdf['Temperature'] -= grid_gdf['Topographic_Cooling']
        pbar.update(2)

        # ECOSTRESS Adjustments

        # Retrieve ECOSTRESS anomalies
        ecostress_anomaly_mean = np.nanmean(grid_gdf['ECOSTRESS_LST'])
        grid_gdf['ECOSTRESS_Anomaly'] = grid_gdf['ECOSTRESS_LST'] - ecostress_anomaly_mean

        # Enhanced Anomaly Adjustment for Nighttime Cold Areas
        if not is_daytime:
            # Define the threshold for the coldest ECOSTRESS anomalies (5th percentile)
            cold_threshold = np.nanpercentile(grid_gdf['ECOSTRESS_Anomaly'], 5)
            
            # Identify cells with anomalies below the threshold
            cold_mask = grid_gdf['ECOSTRESS_Anomaly'] <= cold_threshold
            
            # Apply an additional 2°C cooling effect to these cells
            grid_gdf.loc[cold_mask, 'Temperature'] -= 2.0
            print(f"Applied additional 2°C cooling to {cold_mask.sum()} cold areas based on ECOSTRESS data.")
        
        # Normalize anomalies to range [-1, 1]
        ecostress_anomaly_min = grid_gdf['ECOSTRESS_Anomaly'].min()
        ecostress_anomaly_max = grid_gdf['ECOSTRESS_Anomaly'].max()
        grid_gdf['ECOSTRESS_Anomaly_Normalized'] = (
            (grid_gdf['ECOSTRESS_Anomaly'] - ecostress_anomaly_min) / (ecostress_anomaly_max - ecostress_anomaly_min) * 2 - 1
        )

        # Increase the impact of ECOSTRESS LST data
        amplification_factor = 3.0  # Adjust as needed

        # Compute night_weight: maximum at midnight, zero at noon
        night_weight = (math.cos(((hour % 24) / 24) * 2 * math.pi) + 1) / 2  # Ranges from 1 at midnight to 0 at noon

        # Apply anomalies with the night_weight
        adjustment = grid_gdf['ECOSTRESS_Anomaly_Normalized'] * amplification_factor * night_weight

        # Prevent ECOSTRESS adjustments on water bodies
        adjustment[grid_gdf['Water_Body'] == 1] = 0

        # Limit maximum adjustment to avoid over-adjusting
        max_adjustment = 5.0  # degrees Celsius
        adjustment = adjustment.clip(-max_adjustment, max_adjustment)

        # Apply adjustments to the interpolated temperatures
        grid_gdf['Temperature'] += adjustment

        # Apply the variable cooling effect to water bodies based on time of day
        cooling_effect = compute_water_body_cooling_effect(hour)

        # Apply the cooling effect to grid cells identified as water bodies
        grid_gdf.loc[grid_gdf['Water_Body'] == 1, 'Temperature'] += cooling_effect

        # Create 2D arrays of temperatures
        grid_temperature_with_ecostress = np.full((nrows, ncols), np.nan)
        grid_temperature_with_ecostress[grid_gdf['row_idx'], grid_gdf['col_idx']] = grid_gdf['Temperature'].values

        # Mask the grid_temperature with the Masovian boundary
        masovian_shape = [masovian_boundary.geometry.unary_union]
        mask = rasterize(
            masovian_shape,
            out_shape=(nrows, ncols),
            transform=transform,
            fill=0,
            all_touched=True,
            default_value=1,
            dtype='uint8'
        )

        temperature_with_ecostress = grid_temperature_with_ecostress

        temperature_with_ecostress_masked = np.ma.masked_where(mask == 0, temperature_with_ecostress)

        min_temp = np.nanmin(temperature_with_ecostress)
        max_temp = np.nanmax(temperature_with_ecostress)

        # Also create grid_variance array (we'll skip smoothing variance for simplicity)
        grid_variance = np.full((nrows, ncols), np.nan)
        grid_variance[grid_gdf['row_idx'], grid_gdf['col_idx']] = grid_gdf['Kriging_Variance']
        grid_variance = np.ma.array(grid_variance, mask=mask == 0)

        # Compute Standard Deviation from Variance
        grid_std_dev = np.sqrt(grid_variance)
        grid_std_dev = np.ma.array(grid_std_dev, mask=mask == 0)

        # Extract temperatures over water bodies
        water_temperatures = grid_gdf.loc[grid_gdf['Water_Body'] == 1, 'Temperature']
        land_temperatures = grid_gdf.loc[grid_gdf['Water_Body'] == 0, 'Temperature']

        # Compare average temperatures
        avg_water_temp = water_temperatures.mean()
        avg_land_temp = land_temperatures.mean()
        print(f"Average Water Temperature: {avg_water_temp:.2f}°C")
        print(f"Average Land Temperature: {avg_land_temp:.2f}°C")
        print(f"Temperature Difference (Land - Water): {(avg_land_temp - avg_water_temp):.2f}°C")

        # Read the custom color scale from the CSV file
        color_scale_df = pd.read_csv('input/color_scale.csv')
        color_scale_df = color_scale_df.sort_values('value')
        values = color_scale_df['value'].values
        hex_colors = color_scale_df['color'].values

        # Ensure colors start with '#'
        hex_colors = [color if color.startswith('#') else f"#{color}" for color in hex_colors]

        # Create a custom colormap
        # Normalize the values to 0..1
        vmin = values.min()
        vmax = values.max()
        normalized_values = (values - vmin) / (vmax - vmin)

        # Create a list of tuples (position, color)
        color_tuples = list(zip(normalized_values, hex_colors))

        # Create a LinearSegmentedColormap
        cmap = mcolors.LinearSegmentedColormap.from_list('custom_cmap', color_tuples)

        # Set the normalization
        norm = mcolors.Normalize(vmin=vmin, vmax=vmax)

    with rasterio.open('input/masovian_dem.tif') as dem_src:
        # Reproject Masovian boundary to DEM CRS if necessary
        if masovian_boundary.crs != dem_src.crs:
            masovian_boundary_dem_crs = masovian_boundary.to_crs(dem_src.crs)
        else:
            masovian_boundary_dem_crs = masovian_boundary

        # Clip the DEM to the Masovian boundary
        dem_plot_data, dem_plot_transform = rasterio.mask.mask(
            dem_src, masovian_boundary_dem_crs.geometry, crop=True, filled=True, nodata=np.nan
        )

        # Read the first band (assuming single-band DEM)
        dem_plot_data = dem_plot_data[0]  # Since mask returns data as (bands, rows, cols)

        # Handle NaNs in the DEM data
        dem_plot_data = np.nan_to_num(dem_plot_data, nan=np.nanmin(dem_plot_data))

        # Calculate the extent of the clipped DEM for plotting
        dem_left, dem_top = dem_plot_transform * (0, 0)
        dem_right, dem_bottom = dem_plot_transform * (dem_plot_data.shape[1], dem_plot_data.shape[0])
        dem_extent = [dem_left, dem_right, dem_bottom, dem_top]

        # Create hillshade using DEM
        hs = hillshade(dem_plot_data, azimuth=315, angle_altitude=45)

        # Define extent for plotting
        def get_extent(transform, array):
            x_min = transform.c
            y_max = transform.f
            x_max = x_min + transform.a * array.shape[1]
            y_min = y_max + transform.e * array.shape[0]
            return [x_min, x_max, y_min, y_max]

        extent = get_extent(transform, temperature_with_ecostress)
        dem_extent = get_extent(dem_transform, dem_plot_data)

        # Read water data and get CRS
        water_data, water_transform, water_crs = read_and_clip_water_bodies(masovian_boundary)

        # Create an empty array to hold the resampled water data
        water_resampled = np.empty((nrows, ncols), dtype=water_data.dtype)

        # Define bounding box coordinates
        xmin, ymin, xmax, ymax = 20.188477, 51.017203, 22.061646, 51.821759

        water_mask = np.where(water_resampled > 0, 1, 0)

        # Compute the indices for the bounding box
        from rasterio.transform import rowcol

        # Convert coordinates to row and column indices in the temperature grid
        row_ul, col_ul = rowcol(transform, xmin, ymax)
        row_lr, col_lr = rowcol(transform, xmax, ymin)

        # Ensure indices are within valid ranges
        row_ul = max(0, min(nrows - 1, row_ul))
        row_lr = max(0, min(nrows - 1, row_lr))
        col_ul = max(0, min(ncols - 1, col_ul))
        col_lr = max(0, min(ncols - 1, col_lr))

        # Ensure row_ul <= row_lr and col_ul <= col_lr
        if row_ul > row_lr:
            row_ul, row_lr = row_lr, row_ul
        if col_ul > col_lr:
            col_ul, col_lr = col_lr, col_ul

        # Extract the subset arrays
        temperature_subset = temperature_with_ecostress[row_ul:row_lr+1, col_ul:col_lr+1]
        hs_subset = hs[row_ul:row_lr+1, col_ul:col_lr+1]
        water_mask_subset = water_mask[row_ul:row_lr+1, col_ul:col_lr+1]
        mask_subset = mask[row_ul:row_lr+1, col_ul:col_lr+1]

        # Apply the Masovian mask to the subsets
        temperature_subset = np.ma.array(temperature_subset, mask=mask_subset == 0)
        hs_subset = np.ma.array(hs_subset, mask=mask_subset == 0)
        water_mask_subset = np.ma.array(water_mask_subset, mask=mask_subset == 0)

        # Compute new transform for the subset
        grid_spacing_x = transform.a  # Positive grid spacing in x direction
        grid_spacing_y = transform.e  # Negative grid spacing in y direction

        new_x_min = transform.c + col_ul * grid_spacing_x
        new_y_max = transform.f + row_ul * grid_spacing_y  # grid_spacing_y is negative

        new_transform = from_origin(new_x_min, new_y_max, grid_spacing_x, -grid_spacing_y)

        # Compute the new extent
        extent_subset = get_extent(new_transform, temperature_subset)

        # Filter stations within the bounding box
        bbox_polygon = box(xmin, ymin, xmax, ymax)
        stations_in_bbox = stations_gdf[stations_gdf.geometry.within(bbox_polygon)]

        # Clip the Masovian boundary to the bounding box
        masovian_boundary_clipped = masovian_boundary.clip(bbox_polygon)

        # Plotting Temperature Map with Contours
        # Calculate the aspect ratio of the bounding box
        width = xmax - xmin
        height = ymax - ymin
        aspect_ratio = width / height

        # Set the figure size based on the aspect ratio
        desired_height_inches = 10  # You can adjust this value as desired
        desired_width_inches = desired_height_inches * aspect_ratio

        # Plotting Temperature Map with ECOSTRESS Adjustments
        fig1, ax1 = plt.subplots(figsize=(15, 15))

        # Plot the interpolated temperature grid with ECOSTRESS adjustments
        img1 = ax1.imshow(temperature_with_ecostress_masked, extent=extent, origin='upper',
                          cmap=cmap, norm=norm, alpha=1, interpolation='nearest')

        # Plot Masovian Voivodeship boundary with a light, dark stroke
        masovian_boundary.boundary.plot(ax=ax1, edgecolor='#333333', linewidth=1.5)

        # Plot stations
        stations_gdf.plot(ax=ax1, color='black', marker='s', markersize=5)

        # Remove axes for a minimalistic look
        ax1.axis('off')

        # Draw a light, dark stroke around the plot extent
        from matplotlib.patches import Rectangle
        rect = Rectangle(
            (extent[0], extent[2]),  # Lower-left corner (x0, y0)
            extent[1] - extent[0],   # Width
            extent[3] - extent[2],   # Height
            linewidth=1.5,
            edgecolor='#333333',
            facecolor='none',
            transform=ax1.transData
        )
        ax1.add_patch(rect)

        # Introduce a subtle background gradient
        from matplotlib.colors import LinearSegmentedColormap
        background_cmap = LinearSegmentedColormap.from_list(
            'background_gradient', ['#ffffff', '#f0f0f0'], N=256)
        ax1.imshow(
            np.linspace(0, 1, 256).reshape(1, -1),
            cmap=background_cmap,
            extent=(extent[0], extent[1], extent[2], extent[3]),
            aspect='auto',
            alpha=0.2,
            zorder=-1  # Place it below other elements
        )

        # Adjust font properties for a modern look
        plt.rcParams.update({
            'font.size': 14,  # Base font size
            'font.family': 'sans-serif',
            'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'],
            'text.color': '#333333'
        })

        # Enlarge and reposition the product name (Title)
        fig1.text(
            0.5, 0.94,  # Adjusted y-position for better placement
            "High-Resolution Mesoscale Temperature Analysis",
            fontsize=28,  # Increased font size
            fontweight='bold',
            ha='center',
            va='top',
            fontname='DejaVu Sans',
            color='#333333'
        )

        # Move the model name/resolution text closer to the product name text
        fig1.text(
            0.5, 0.91,  # Adjusted y-position to move it closer
            "Model: HRMTA v1.1 (Preview)   |   Resolution: 250 m",
            fontsize=18,  # Increased font size
            ha='center',
            va='top',
            fontname='DejaVu Sans',
            color='#666666'
        )

        # Adjust colorbar position and style
        from mpl_toolkits.axes_grid1 import make_axes_locatable

        divider = make_axes_locatable(ax1)
        cax = divider.append_axes("right", size="3%", pad=0.1)
        cbar1 = fig1.colorbar(img1, cax=cax)
        ticks = np.arange(np.floor(vmin / 5) * 5, np.ceil(vmax / 5) * 5 + 1, 5)
        cbar1.set_ticks(ticks)
        cbar1.set_ticklabels(ticks.astype(int))
        cbar1.ax.tick_params(labelsize=14)

        # Enhance colorbar aesthetics
        cbar1.outline.set_visible(False)
        cbar1.ax.yaxis.set_tick_params(width=0.5)
        cbar1.ax.tick_params(labelsize=14, colors='#333333')

        # Move the logo to the bottom right corner, directly below the colorbar
        from matplotlib.offsetbox import OffsetImage, AnnotationBbox
        import matplotlib.image as mpimg

        # Read the logo image
        logo_path = 'input/logo.png'
        logo_img = mpimg.imread(logo_path)

        # Adjust the size of the logo
        logo_zoom = 0.25
        imagebox = OffsetImage(logo_img, zoom=logo_zoom)

        # Get the position of the colorbar axes
        cbar_bbox = cax.get_position()
        cbar_x0 = cbar_bbox.x0  # Left
        cbar_x1 = cbar_bbox.x1  # Right

        # Calculate the x position for the logo (centered below the colorbar)
        logo_x = (cbar_x0 + cbar_x1) / 2
        # Set the y position as low as possible, but ensure the logo is fully visible
        logo_y = 0.02  # Adjust this value if necessary

        ab = AnnotationBbox(
            imagebox,
            xy=(logo_x, logo_y),
            xycoords='figure fraction',
            frameon=False,
            box_alignment=(0.5, 0),  # Centered horizontally, aligned at bottom
        )

        fig1.add_artist(ab)

        # Adjust the bottom margin to ensure the logo is not cut off
        plt.subplots_adjust(left=0.05, right=0.95, top=0.88, bottom=0.05)

        # Ensure the temperature grid and water mask are aligned
        if temperature_with_ecostress.shape != water_mask.shape:
            raise ValueError("Temperature grid and water mask dimensions do not match.")

        # Create a masked array for land temperatures by excluding water bodies
        temperature_land = temperature_with_ecostress.copy()
        # If temperature_with_ecostress is not already a masked array, convert it
        if not isinstance(temperature_land, np.ma.MaskedArray):
            temperature_land = np.ma.array(temperature_land)
        # Update the mask to include water bodies
        temperature_land.mask = temperature_land.mask | (water_mask == 1)

        # Plot the temperature values on the map (without markers)
        # Use white text with a thick black stroke for better visibility
        text_kwargs = dict(
            fontsize=18,  # Increased font size for clarity
            fontname='Arial',
            ha='center',
            va='center',
            color='white',
            transform=ax1.transData,
            path_effects=[patheffects.withStroke(linewidth=3, foreground='black')]
        )

        # Save and show the plot
        plt.savefig('output/temperature_analysis_modern.png', dpi=300)
        plt.show()

        print("Updated temperature analysis plot created successfully.")

        # Calculate the aspect ratio of the bounding box
        width = xmax - xmin
        height = ymax - ymin
        aspect_ratio = width / height

        # Set the maximum desired width
        max_width_inches = 15  # Adjust this value as desired

        # Calculate the figure height based on the aspect ratio and maximum width
        desired_width_inches = min(max_width_inches, aspect_ratio * desired_height_inches)
        desired_height_inches = desired_width_inches / aspect_ratio

        # Ensure the height doesn't exceed a certain value
        desired_height_inches = min(desired_height_inches, 10)  # Adjust as needed

        fig2, ax2 = plt.subplots(figsize=(desired_width_inches, desired_height_inches))

        # Proceed with previous plotting code, using adjusted fig2 and ax2

        # Use the subset data and extent with applied mask
        img2 = ax2.imshow(temperature_subset, extent=extent_subset, origin='upper',
                          cmap=cmap, norm=norm, alpha=1, interpolation='nearest')

        # Plot Masovian boundary
        masovian_boundary_clipped.boundary.plot(ax=ax2, edgecolor='#333333', linewidth=1.5)

        # Plot stations within the bounding box
        stations_in_bbox.plot(ax=ax2, color='black', marker='s', markersize=5)

        # Remove axes for a minimalistic look
        ax2.axis('off')

        # Introduce a subtle background gradient using the subset extent
        ax2.imshow(
            np.linspace(0, 1, 256).reshape(1, -1),
            cmap=background_cmap,
            extent=extent_subset,
            aspect='auto',
            alpha=0.2,
            zorder=-1  # Place it below other elements
        )

        # Adjust font properties for a modern look
        plt.rcParams.update({
            'font.size': 14,
            'font.family': 'sans-serif',
            'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'],
            'text.color': '#333333'
        })

        # Enlarge and reposition the product name (Title)
        fig2.text(
            0.5, 0.995,  # Moved up towards the top edge
            "High-Resolution Mesoscale Temperature Analysis",
            fontsize=12,
            fontweight='bold',
            ha='center',
            va='top',
            fontname='DejaVu Sans',
            color='#333333'
        )

        # Move the model name/resolution text closer to the product name text
        fig2.text(
            0.5, 0.970,  # Moved up slightly below the title
            "Model: HRMTA v1.1 (Preview)   |   Resolution: 250 m",
            fontsize=9,
            ha='center',
            va='top',
            fontname='DejaVu Sans',
            color='#666666'
        )

        # Add data source attribution at the bottom left corner
        fig2.text(
            0.01, 0.02,
            "Data source: IMGW-PIB (processed)",
            fontsize=12,
            ha='left',
            va='bottom',
            fontname='DejaVu Sans',
            color='#333333'
        )

        # Adjust colorbar position and style
        from mpl_toolkits.axes_grid1 import make_axes_locatable

        divider = make_axes_locatable(ax2)
        cax2 = divider.append_axes("right", size="3%", pad=0.1)
        cbar2 = fig2.colorbar(img2, cax=cax2)
        cbar2.set_ticks(ticks)
        cbar2.set_ticklabels(ticks.astype(int))

        # Import FontProperties and load the font
        from matplotlib.font_manager import FontProperties

        # Replace 'input/RobotoCondensed-Regular.ttf' with the actual path to your font file
        roboto_font_path = 'input/RobotoCondensed-Regular.ttf'
        roboto_condensed = FontProperties(fname=roboto_font_path)

        # Apply font properties to the colorbar tick labels
        for label in cbar2.ax.yaxis.get_ticklabels():
            label.set_fontproperties(roboto_condensed)

        cbar2.ax.tick_params(labelsize=14)

        # Enhance colorbar aesthetics
        cbar2.outline.set_visible(False)
        cbar2.ax.yaxis.set_tick_params(width=0.5)
        cbar2.ax.tick_params(labelsize=14, colors='#333333')

        # Move the logo to the bottom right corner, directly below the colorbar
        from matplotlib.offsetbox import OffsetImage, AnnotationBbox
        import matplotlib.image as mpimg

        # Read the logo image again for the second plot
        logo_img2 = mpimg.imread(logo_path)

        # Adjust the size of the logo
        imagebox2 = OffsetImage(logo_img2, zoom=logo_zoom)

        # Get the position of the colorbar axes
        cbar_bbox2 = cax2.get_position()
        logo_x2 = (cbar_bbox2.x0 + cbar_bbox2.x1) / 2
        logo_y2 = 0.02  # Adjust this value if necessary

        ab2 = AnnotationBbox(
            imagebox2,
            xy=(logo_x2, logo_y2),
            xycoords='figure fraction',
            frameon=False,
            box_alignment=(0.5, 0),
        )
        fig2.add_artist(ab2)

        # Adjust the bottom margin to ensure the logo and data attribution text are not cut off
        plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.06)

        # Ensure the temperature grid and water mask are aligned
        if temperature_with_ecostress.shape != water_mask.shape:
            raise ValueError("Temperature grid and water mask dimensions do not match.")

        # Create a masked array for land temperatures by excluding water bodies
        temperature_land = temperature_with_ecostress.copy()
        # If temperature_with_ecostress is not already a masked array, convert it
        if not isinstance(temperature_land, np.ma.MaskedArray):
            temperature_land = np.ma.array(temperature_land)
        # Update the mask to include water bodies
        temperature_land.mask = temperature_land.mask | (water_mask == 1)

        contour_levels = np.arange(np.floor(min_temp), np.ceil(max_temp) + 1, 1)

    if min_temp < max_temp:
        # Create contour levels with 1°C intervals
        contour_levels = np.arange(np.floor(min_temp), np.ceil(max_temp) + 1, 1)
    else:
        # If min and max temperatures are equal, define a small range around the temperature
        # to allow contouring
        contour_levels = np.arange(np.floor(min_temp) - 1, np.ceil(max_temp) + 2, 1)

    # Ensure contour_levels has at least two distinct levels
    if len(contour_levels) < 2:
        print("Insufficient temperature range for contour plotting. Skipping contour lines.")
        contours = None
    else:
        # Continue with plotting contours
        contours = ax2.contour(
            temperature_subset, levels=contour_levels, colors='black', linewidths=0.5,
            extent=extent_subset, origin='upper'
        )
        contour_labels = ax2.clabel(
            contours, inline=False, fontsize=8, fmt='%d', colors='white',
            inline_spacing=4, manual=False
        )
        for txt in contour_labels:
            txt.set_path_effects([
                patheffects.Stroke(linewidth=3, foreground='black'),
                patheffects.Normal()
            ])

        # Define font properties
        from matplotlib.font_manager import FontProperties
        roboto_font_path = 'input/RobotoCondensed-Regular.ttf'
        roboto_condensed = FontProperties(fname=roboto_font_path)

    # Save and show the plot
    plt.savefig('output/temperature_analysis_with_contours2.png', dpi=300)
    plt.show()

    print("Temperature analysis plot with contours created successfully.")

    # At the end, print feature importances to confirm ECOSTRESS influence
    print("\nFeature importances from the XGBoost model:")
    print(feature_importances_df)

    # Check if 'Distance_to_Water' has significant importance
    if 'Distance_to_Water' in feature_importances_df['Feature'].values:
        distance_water_importance = feature_importances_df.loc[feature_importances_df['Feature'] == 'Distance_to_Water', 'Importance'].values[0]
        if distance_water_importance > 0.05:  # Threshold can be adjusted
            print(f"\n'Distance_to_Water' feature has significant influence on the model (importance: {distance_water_importance:.4f}).")
        else:
            print(f"\n'Distance_to_Water' feature has low influence on the model (importance: {distance_water_importance:.4f}).")
    else:
        print("\n'Distance_to_Water' feature not found in feature importances.")

    # Save the trained model for future use
    joblib.dump(best_xgb_pipeline, 'output/best_xgb_pipeline_enhanced.joblib')
    print("\nTrained XGBoost model saved successfully.")

    return temperature_with_ecostress, transform

def get_interpolated_temperature_at_two_points(x1, y1, x2, y2, temperature_grid, transform):
    """
    Given two coordinate pairs, return the interpolated temperature values at those points.

    Parameters:
    - x1, y1: X and Y coordinates of the first point.
    - x2, y2: X and Y coordinates of the second point.
    - temperature_grid: The interpolated temperature grid (2D numpy array).
    - transform: Affine transform for the grid.

    Returns:
    - A tuple containing interpolated temperature values at the two points,
      or None for a point if it's outside the grid or has invalid data.
    """
    import numpy as np
    from rasterio.transform import rowcol

    temps = []
    for x, y in [(x1, y1), (x2, y2)]:
        try:
            # Convert geographic coordinates to raster indices
            row, col = rowcol(transform, x, y)  # rasterio uses (x, y)
            row = int(row)
            col = int(col)

            # Retrieve the value from the temperature grid
            value = temperature_grid[row, col]

            # Check if the value is valid
            if np.ma.is_masked(value) or np.isnan(value):
                # Value is masked or NaN, meaning outside the interpolated area
                temps.append(None)
            else:
                temps.append(value.item())
        except (IndexError, ValueError):
            # Point is outside the grid
            temps.append(None)
    return tuple(temps)

# Ensure to adjust your main execution block accordingly
if __name__ == "__main__":
    # Fetching Data Progress Bar
    with tqdm(total=100, desc='Fetching Data', unit='%', ncols=80) as pbar:
        # Fetch Traxelektronik data
        trax_stations, trax_temperatures = fetch_traxelektronik_data()
        pbar.update(20)

        # Geocode Traxelektronik stations
        geolocator = Nominatim(user_agent="geoapi")
        trax_coords = {}

        for station in tqdm(trax_stations, desc="Geocoding Traxelektronik Stations", unit="station"):
            cleaned_station = clean_station_name(station)
            coords, was_cleaned = get_coordinates(cleaned_station, geolocator)
            if coords:
                trax_coords[cleaned_station] = coords  # Use cleaned station names as keys
        pbar.update(30)

        # Fetch Netatmo data
        netatmo_stations, netatmo_temperatures, netatmo_coords = fetch_netatmo_data()
        pbar.update(20)

        # Fetch IMGW data
        imgw_stations, imgw_temperatures, imgw_station_ids = fetch_imgw_data()
        pbar.update(10)

        # Geocode IMGW stations
        imgw_coords = {}
        for station in tqdm(imgw_stations, desc="Geocoding IMGW Stations", unit="station"):
            coords, was_cleaned = get_coordinates(station, geolocator)
            if coords:
                imgw_coords[station] = coords
        pbar.update(20)

    # Paths to MODIS data folders (no longer used, but kept for compatibility)
    modis_day_folder = 'input/MODIS_day'
    modis_night_folder = 'input/MODIS_night'

    # Call the function with all required arguments
    smoothed_temperature, transform = plot_stations_with_boundary_and_interpolation(
        trax_stations, trax_temperatures, trax_coords,
        netatmo_stations, netatmo_temperatures, netatmo_coords,
        imgw_stations, imgw_temperatures, imgw_coords,
        modis_day_folder, modis_night_folder
    )

    # Coordinates to test (ensure these are in EPSG:2180)
    test_x1, test_y1 = 691000, 5130000  # Example coordinates in EPSG:2180
    test_x2, test_y2 = 692000, 5131000  # Example coordinates in EPSG:2180

    # Call the modified function
    interpolated_temp1, interpolated_temp2 = get_interpolated_temperature_at_two_points(
        test_x1, test_y1, test_x2, test_y2, smoothed_temperature, transform
    )

    if interpolated_temp1 is not None and interpolated_temp2 is not None:
        print(f"\nLocation 1 [{test_x1}, {test_y1}]: {interpolated_temp1:.2f}°C")
        print(f"Location 2 [{test_x2}, {test_y2}]: {interpolated_temp2:.2f}°C")
    else:
        print(f"\nOne or both coordinates are outside the interpolated grid or have no data.")

Traceback (most recent call last):
  File "c:\Users\Lenovo\Downloads\pythonStationsData\xgboost_interpolation_enhanced.py", line 1719, in <module>
    smoothed_temperature, transform = plot_stations_with_boundary_and_interpolation(
  File "c:\Users\Lenovo\Downloads\pythonStationsData\xgboost_interpolation_enhanced.py", line 1595, in plot_stations_with_boundary_and_interpolation
    contours = ax2.contour(
  File "C:\Users\Lenovo\Documents\conda\envs\temp_env\lib\site-packages\matplotlib\__init__.py", line 1473, in inner
    return func(
  File "C:\Users\Lenovo\Documents\conda\envs\temp_env\lib\site-packages\matplotlib\axes\_axes.py", line 6659, in contour
    contours = mcontour.QuadContourSet(self, *args, **kwargs)
  File "C:\Users\Lenovo\Documents\conda\envs\temp_env\lib\site-packages\matplotlib\contour.py", line 813, in __init__
    kwargs = self._process_args(*args, **kwargs)
  File "C:\Users\Lenovo\Documents\conda\envs\temp_env\lib\site-packages\matplotlib\contour.py", line 1470, in _process_args
    x, y, z = self._contour_args(args, kwargs)
  File "C:\Users\Lenovo\Documents\conda\envs\temp_env\lib\site-packages\matplotlib\contour.py", line 1507, in _contour_args
    x, y = self._initialize_x_y(z)
  File "C:\Users\Lenovo\Documents\conda\envs\temp_env\lib\site-packages\matplotlib\contour.py", line 1583, in _initialize_x_y
    raise TypeError(f"Input z must be at least a (2, 2) shaped array, "
TypeError: Input z must be at least a (2, 2) shaped array, but has shape (1, 1)

Fix this error here please