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
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
import os
import rasterio.mask
import matplotlib.colors as mcolors
import matplotlib.patheffects as patheffects
from pykrige.uk import UniversalKriging
from sklearn.linear_model import HuberRegressor
from scipy.stats import zscore
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from scipy.ndimage import gaussian_filter

os.environ['PROJ_LIB'] = r'C:\Users\Lenovo\Documents\conda\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']
        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")
    masovian_geom = masovian_boundary.to_crs(epsg=2180).reset_index(drop=True)
    point_projected = point.to_crs(epsg=2180)
    return masovian_geom.contains(point_projected.buffer(100)).values[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 != "-":
                cleaned_station = clean_station_name(station)
                stations.append(cleaned_station)
                temperatures.append(temperature)

    print(f"Fetched {len(stations)} stations with valid temperature data from Traxelektronik.")
    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|5f5006a604c524ff5e8753b2290a594a"
    }
    
    response = requests.get(url, params=params)
    if response.status_code == 200:
        print("Successfully fetched data from Netatmo")
    else:
        print(f"Error fetching data: {response.status_code}")
    data = json.loads(response.text)
    
    stations = []
    temperatures = []
    coordinates = []
    
    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):
                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]
                            city = place.get('city', 'Unknown')
                            stations.append(city)
                            temperatures.append(temperature)
                            coordinates.append((lat, lon))
                            break
    
    print(f"Fetched {len(stations)} stations with valid temperature data from Netatmo in Masovian Voivodeship.")
    return stations, temperatures, coordinates

def fetch_imgw_data():
    """
    Fetches additional weather station data from the IMGW JSON endpoint.
    Targets specific stations and retrieves their current temperatures.
    """
    url = "https://rafalraczynski.com.pl/imgw/dane-imgw/getJSON.php?type=table&province=14&sort=temp&order=asc&_=1726673311725"
    target_stations = [
        "ANDRZEJEWO", "BORKOWO", "CZARNOWO", "KAZANÓW",
        "RUSZKOWICE", "RYBIENKO", "RZĄŚNIK WŁOŚCIAŃSKI", "WIELGOLAS"
    ]
    
    response = requests.get(url)
    if response.status_code != 200:
        raise ConnectionError(f"Error fetching IMGW data: {response.status_code}")
    
    try:
        data = response.json()
    except json.JSONDecodeError:
        raise ValueError("Failed to parse JSON response from IMGW data endpoint.")
    
    stations = []
    temperatures = []
    
    for record in data:
        stat_name = record.get("statName", "").upper()
        temp_str = record.get("temp", "")
        if stat_name in target_stations and temp_str:
            try:
                temp = float(re.sub(r'[↓↑\s]', '', temp_str))
                stations.append(stat_name.title())  # Capitalize properly
                temperatures.append(temp)
            except ValueError:
                print(f"Invalid temperature value for station {stat_name}: {temp_str}")
    
    print(f"Fetched {len(stations)} additional IMGW stations with valid temperature data.")
    return stations, temperatures

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):
                print(f"Geocoded '{station_name}' to {location.latitude}, {location.longitude} (in Masovian).")
                return (location.latitude, location.longitude), False
            else:
                print(f"'{station_name}' is not in Masovian Voivodeship.")
                return None, False
        print(f"Failed to geocode '{station_name}'.")
        return None, False
    except Exception as e:
        print(f"Error geocoding {station_name}: {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, trax_temperatures, netatmo_coords, netatmo_temperatures, imgw_coords, imgw_temperatures, distance_threshold=0.01):
    trax_points = np.array([(lon, lat) for station, (lat, lon) in trax_coords.items()])
    netatmo_points = np.array([(lon, lat) for lat, lon in netatmo_coords])
    imgw_points = np.array([(lon, lat) for lat, lon in imgw_coords])
    
    trax_tree = cKDTree(trax_points)
    netatmo_tree = cKDTree(netatmo_points)
    
    netatmo_filtered_coords = []
    netatmo_filtered_temperatures = []
    
    for netatmo_point, netatmo_temp in zip(netatmo_points, netatmo_temperatures):
        distance, index = trax_tree.query(netatmo_point)
        if distance > distance_threshold:
            netatmo_filtered_coords.append(tuple(netatmo_point))
            netatmo_filtered_temperatures.append(netatmo_temp)
    
    imgw_filtered_coords = []
    imgw_filtered_temperatures = []
    
    for imgw_point, imgw_temp in zip(imgw_points, imgw_temperatures):
        # Check against Traxelektronik data
        distance_trax, _ = trax_tree.query(imgw_point)
        # Check against Netatmo data
        distance_netatmo, _ = netatmo_tree.query(imgw_point)
        if distance_trax > distance_threshold and distance_netatmo > distance_threshold:
            imgw_filtered_coords.append(tuple(imgw_point))
            imgw_filtered_temperatures.append(imgw_temp)
    
    return netatmo_filtered_coords, netatmo_filtered_temperatures, imgw_filtered_coords, imgw_filtered_temperatures

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 gaussian_filter_nan(image, sigma):
    nan_mask = np.isnan(image)
    image_filled = np.where(nan_mask, 0, image)
    weights = np.where(nan_mask, 0, 1)
    image_filtered = gaussian_filter(image_filled, sigma=sigma, mode='nearest')
    weights_filtered = gaussian_filter(weights, sigma=sigma, mode='nearest')
    with np.errstate(divide='ignore', invalid='ignore'):
        image_final = image_filtered / weights_filtered
    image_final[nan_mask] = np.nan
    return image_final

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):
    """
    Plots the temperature interpolation and error maps with all station data.
    Prioritizes IMGW data over Netatmo and Traxelektronik data.
    """
    # Initialize lists for GeoDataFrames
    trax_filtered_stations = []
    trax_filtered_temperatures = []
    trax_geometries = []

    # Filter and clean Traxelektronik data
    for station, temperature in zip(trax_stations, trax_temperatures):
        cleaned_station = clean_station_name(station)
        cleaned_temperature = clean_temperature(temperature)
        if station in trax_coords and cleaned_temperature is not None:
            lat, lon = trax_coords[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="EPSG:4326")

    # 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="EPSG:4326")
    else:
        netatmo_gdf = gpd.GeoDataFrame(columns=['Station', 'Temperature', 'Source', 'geometry'], crs="EPSG:4326")

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

    for station, temperature in zip(imgw_stations, imgw_temperatures):
        if isinstance(temperature, (int, float)):
            # Geocode the station to get coordinates
            coords, _ = get_coordinates(station, geolocator=geolocator)
            if coords:
                imgw_filtered_coords.append(coords)
                imgw_filtered_temperatures.append(temperature)
                imgw_filtered_stations.append(station)

    if len(imgw_filtered_coords) > 0:
        imgw_geometries = [Point(lon, lat) for lat, lon in imgw_filtered_coords]
        imgw_gdf = gpd.GeoDataFrame({
            'Station': imgw_filtered_stations,
            'Temperature': imgw_filtered_temperatures,
            'Source': 'IMGW'
        }, geometry=imgw_geometries, crs="EPSG:4326")
    else:
        imgw_gdf = gpd.GeoDataFrame(columns=['Station', 'Temperature', 'Source', 'geometry'], crs="EPSG:4326")

    # Read and clip DEM, Forest Cover, and Water Bodies
    dem_data, dem_transform = read_and_clip_dem(masovian_boundary)
    forest_data, forest_transform = read_and_clip_forest(masovian_boundary)
    water_data, water_transform = read_and_clip_water_bodies(masovian_boundary)

    # Create GeoDataFrames for each source with additional attributes
    for gdf in [trax_gdf, netatmo_gdf, imgw_gdf]:
        if not gdf.empty:
            # Add elevation
            gdf['Elevation'] = get_elevation_at_points(gdf.geometry, dem_data=dem_data, dem_transform=dem_transform)
            # Add tree cover
            gdf['Tree_Cover'] = get_raster_value_at_points(gdf.geometry, raster_data=forest_data, raster_transform=forest_transform)
            # Add water body presence
            gdf['Water_Body'] = get_water_body_at_points(gdf.geometry, water_data=water_data, water_transform=water_transform)

    # Combine GeoDataFrames with priority: IMGW > Traxelektronik > Netatmo
    stations_gdf = pd.concat([imgw_gdf, trax_gdf, netatmo_gdf], ignore_index=True)
    stations_gdf = gpd.GeoDataFrame(stations_gdf, crs="EPSG:4326")

    # Drop rows with missing essential data
    essential_columns = ['Temperature', 'Elevation', 'Tree_Cover']
    stations_gdf = stations_gdf.dropna(subset=essential_columns)

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

    # Assign weights based on source priority (IMGW > Traxelektronik > Netatmo)
    weight_mapping = {'IMGW': 3, 'Traxelektronik': 2, 'Netatmo': 1}
    stations_gdf['Weight'] = stations_gdf['Source'].map(weight_mapping)

    # Prepare data for regression
    X = stations_gdf[['X', 'Y', 'Elevation', 'Tree_Cover', 'Water_Body']]
    y = stations_gdf['Temperature']
    sample_weight = stations_gdf['Weight']

    # Normalize covariates
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # Fit Random Forest Regression model with sample weights
    print("Training Random Forest Regression model with weighted samples...")
    rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
    rf_model.fit(X_scaled, y, sample_weight=sample_weight)

    # Predict temperatures and compute residuals
    stations_gdf['Predicted_Temperature'] = rf_model.predict(X_scaled)
    stations_gdf['Residual'] = stations_gdf['Temperature'] - stations_gdf['Predicted_Temperature']

    # Outlier detection using Z-score
    stations_gdf['Residual_Z'] = zscore(stations_gdf['Residual'])
    z_threshold = 2  # Set your threshold
    stations_gdf_inliers = stations_gdf[np.abs(stations_gdf['Residual_Z']) <= z_threshold].copy()

    # Remove duplicates by averaging residuals at duplicate locations
    data = pd.DataFrame({
        'X': stations_gdf_inliers['X'],
        'Y': stations_gdf_inliers['Y'],
        'Residual': stations_gdf_inliers['Residual'],
        'Elevation': stations_gdf_inliers['Elevation'],
        'Tree_Cover': stations_gdf_inliers['Tree_Cover'],
        'Water_Body': stations_gdf_inliers['Water_Body']
    })
    data = data.groupby(['X', 'Y'], as_index=False).mean()
    unique_points = data[['X', 'Y']].values
    unique_residuals = data['Residual'].values
    unique_elevation = data['Elevation'].values
    unique_tree_cover = data['Tree_Cover'].values
    unique_water_body = data['Water_Body'].values

    # Create grid in EPSG:4326 using rasterio's transform functions
    grid_spacing = 0.0025  # Approximately ~5 km

    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 dataframe
    grid_df = pd.DataFrame({
        'X': xs_flat,
        'Y': ys_flat,
        'row_idx': rows_flat,
        'col_idx': cols_flat
    })

    # Retrieve elevations from DEM
    grid_elevations = get_elevation_at_points(
        [Point(x, y) for x, y in zip(grid_df['X'], grid_df['Y'])],
        dem_data=dem_data,
        dem_transform=dem_transform
    )
    grid_df['Elevation'] = grid_elevations

    # Retrieve tree cover density from forest cover raster
    grid_forest = get_raster_value_at_points(
        [Point(x, y) for x, y in zip(grid_df['X'], grid_df['Y'])],
        raster_data=forest_data,
        raster_transform=forest_transform
    )
    grid_df['Tree_Cover'] = grid_forest

    # Retrieve water body data at grid points
    grid_water = get_water_body_at_points(
        [Point(x, y) for x, y in zip(grid_df['X'], grid_df['Y'])],
        water_data=water_data,
        water_transform=water_transform
    )
    grid_df['Water_Body'] = grid_water

    # Remove grid points with NaN values (i.e., outside data coverage)
    grid_df = grid_df.dropna(subset=['Elevation', 'Tree_Cover'])

    # Prepare data for regression prediction
    X_grid = grid_df[['X', 'Y', 'Elevation', 'Tree_Cover', 'Water_Body']]
    X_grid_scaled = scaler.transform(X_grid)
    grid_df['Predicted_Temperature'] = rf_model.predict(X_grid_scaled)

    # Universal Kriging for residual interpolation
    variogram_model = 'spherical'

    print("Performing Universal Kriging interpolation for residuals...")
    uk = UniversalKriging(
        x=data['X'].values,
        y=data['Y'].values,
        z=data['Residual'].values,
        variogram_model=variogram_model,
        nlags=20,
        verbose=False
    )

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

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

    # Final temperature estimation
    grid_df['Temperature'] = grid_df['Predicted_Temperature'] + grid_df['Residual']

    # Apply Gaussian smoothing to the interpolated temperature grid
    print("Applying Gaussian smoothing to the interpolated temperature grid...")
    # First, create a 2D array of temperatures
    grid_temperature = np.full((nrows, ncols), np.nan)
    grid_temperature[grid_df['row_idx'], grid_df['col_idx']] = grid_df['Temperature']

    # Mask the grid_temperature
    masovian_shape = [masovian_boundary.geometry.unary_union]

    # Generate mask with all_touched=True to include partial pixels
    mask = rasterize(
        masovian_shape,
        out_shape=(nrows, ncols),
        transform=transform,
        fill=0,
        all_touched=True,
        default_value=1,
        dtype='uint8'
    )

    sigma = 1
    smoothed_temperature = gaussian_filter_nan(grid_temperature, sigma=sigma)

    # Apply the mask to the temperature data to limit it to Masovian region
    smoothed_temperature = np.ma.array(smoothed_temperature, mask=mask == 0)

    # Also create grid_variance array (we'll skip smoothing variance for simplicity)
    grid_variance = np.full((nrows, ncols), np.nan)
    grid_variance[grid_df['row_idx'], grid_df['col_idx']] = grid_df['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)

    # 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()
    if vmax - vmin == 0:
        vmin = 0
        vmax = 1
    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.where(np.isnan(dem_plot_data), np.nanmin(dem_plot_data), 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]

    # Compute the 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]
        
    hs = hillshade(dem_plot_data, azimuth=315, angle_altitude=45)
        
    extent = get_extent(transform, smoothed_temperature)
    dem_extent = get_extent(dem_plot_transform, dem_plot_data)

    # Plotting Temperature Map
    fig, ax = plt.subplots(figsize=(15, 15))

    # Set modern and minimalistic font properties with increased size
    plt.rcParams.update({
        'font.size': 10,
        'font.family': 'sans-serif',
        'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans']
    })

    # Plot the interpolated temperature grid with custom colormap
    img = ax.imshow(smoothed_temperature, extent=extent, origin='upper', cmap=cmap, norm=norm, alpha=1)

    # Plot hillshade on top
    dem_img = ax.imshow(hs, extent=dem_extent, cmap='gray', origin='upper', alpha=0.3)

    # Overlay water bodies with blue color
    water_mask = np.where(water_data > 0, 1, 0)  # Adjust condition based on water_data values
    ax.imshow(
        water_mask, 
        extent=extent, 
        origin='upper', 
        cmap='Blues', 
        alpha=0.3
    )

    # Plot Masovian Voivodeship boundary
    masovian_boundary.boundary.plot(ax=ax, color='black', linewidth=2)

    # Plot stations as smaller rectangles
    stations_gdf.plot(
        ax=ax, 
        color='black', 
        marker='s',             # 's' for square markers (rectangles)
        markersize=5,           # Smaller size than before
        label='Stations'
    )

    # Add contours for every 0.5°C value
    min_temp = np.floor(smoothed_temperature.min())
    max_temp = np.ceil(smoothed_temperature.max())
    contour_levels = np.arange(min_temp, max_temp + 0.5, 0.5)

    # Plot contour lines without interrupting them
    contours = ax.contour(
        smoothed_temperature, levels=contour_levels, colors='black', linewidths=0.5,
        extent=extent, origin='upper'
    )

    # Add labels without cutting the lines
    contour_labels = ax.clabel(
        contours, inline=False, fontsize=8, fmt='%0.1f', colors='white',
        inline_spacing=4, manual=False
    )

    # Add black outline to contour labels for better visibility
    for txt in contour_labels:
        txt.set_path_effects([
            patheffects.Stroke(linewidth=3, foreground='black'),
            patheffects.Normal()
        ])

    # Add colorbar and labels
    cbar = fig.colorbar(img, ax=ax, label="Temperature (°C)")

    # Generate ticks every 5 units within the temperature range
    ticks = np.arange(np.floor(vmin / 5) * 5, np.ceil(vmax / 5) * 5 + 1, 5)
    
    # Adjust colorbar ticks and labels to show every 5 units
    cbar.set_ticks(ticks)
    cbar.set_ticklabels(ticks.astype(int))

    ax.set_title("Temperature Interpolation in Masovian Voivodeship", fontsize=18)
    ax.legend()
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")

    plt.savefig('output/output_testing.png', dpi=300, bbox_inches='tight')
    plt.show()

    print(f"Temperature interpolation plot created successfully with smoothed data.")

    # Plotting Error Map with Standard Deviation
    fig_error, ax_error = plt.subplots(figsize=(15, 15))

    # Plot the kriging standard deviation grid (temperature error interpolation)
    cax = ax_error.imshow(
        grid_std_dev, 
        extent=extent, 
        origin='upper', 
        cmap='RdBu', 
        alpha=1  # Full opacity for the error data
    )

    # Overlay the hillshade data on top with 25% transparency
    dem_img = ax_error.imshow(
        hs, 
        extent=dem_extent, 
        cmap='gray', 
        origin='upper', 
        alpha=0.5  # 25% transparency for the hillshade
    )

    # Overlay water bodies with blue color
    ax_error.imshow(
        water_mask, 
        extent=extent, 
        origin='upper', 
        cmap='Blues', 
        alpha=0.3
    )

    # Overlay the Masovian Voivodeship boundary
    masovian_boundary.boundary.plot(
        ax=ax_error, 
        color='black', 
        linewidth=2
    )

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

    # Add colorbar for the error data
    cbar_error = fig_error.colorbar(
        cax, 
        ax=ax_error, 
        label="Kriging Standard Deviation (°C)"
    )

    # Set plot titles and labels
    ax_error.set_title(
        "Estimated Mean Temperature Error (Standard Deviation)", 
        fontsize=18
    )
    ax_error.legend()
    ax_error.set_xlabel("Longitude")
    ax_error.set_ylabel("Latitude")

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

    print(f"Estimated mean temperature error (standard deviation) plot created successfully.")

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)
                value = raster_data[row, col]
                values.append(value)
            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:
        print("Assuming DEM is in EPSG:4326; skipping reprojection.")
        dem_data = src.read(1, masked=True)
        dem_transform = src.transform
    return dem_data, dem_transform

def read_and_clip_forest(masovian_boundary):
    forest_path = 'input/forests.tif'
    with rasterio.open(forest_path) as src:
        print("Assuming forest cover raster is in EPSG:4326; skipping reprojection.")
        forest_data = src.read(1, masked=True)
        forest_data = np.where(forest_data == src.nodata, 0, forest_data)
        forest_transform = src.transform
    return forest_data, forest_transform

def read_and_clip_water_bodies(masovian_boundary):
    water_path = 'input/water_bodies.tif'
    with rasterio.open(water_path) as src:
        print("Reading and clipping water bodies raster to Masovian boundaries.")
        
        # Ensure the Masovian boundary is in the same CRS as the water raster
        if masovian_boundary.crs != src.crs:
            masovian_boundary = masovian_boundary.to_crs(src.crs)
            print("Reprojected Masovian boundary to match water body raster CRS.")
        else:
            print("Masovian boundary CRS matches water body raster CRS.")
        
        # Perform the clipping
        clipped_water, water_transform = rasterio.mask.mask(
            src, masovian_boundary.geometry, crop=True, filled=True, nodata=0
        )
        
        # Read the first band (assuming single-band raster)
        water_data = clipped_water[0]
        
        # Handle NoData values by setting them to 0 (assuming 0 represents no water)
        water_data = np.where(water_data == src.nodata, 0, water_data)
        
    return water_data, water_transform

def read_and_clip_dem(masovian_boundary):
    dem_path = 'input/masovian_dem.tif'
    with rasterio.open(dem_path) as src:
        print("Assuming DEM is in EPSG:4326; skipping reprojection.")
        dem_data = src.read(1, masked=True)
        dem_transform = src.transform
    return dem_data, dem_transform

# Main execution flow
if __name__ == "__main__":
    # Fetch Traxelektronik data
    trax_stations, trax_temperatures = fetch_traxelektronik_data()

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

    progress_bar = tqdm(total=len(trax_stations), desc="Geocoding Traxelektronik Stations", unit="station")
    
    for station in trax_stations:
        cleaned_station = clean_station_name(station)
        coords, was_cleaned = get_coordinates(cleaned_station, geolocator)
        if coords:
            trax_coords[station] = coords
        progress_bar.update(1)
    
    progress_bar.close()

    # Fetch Netatmo data
    netatmo_stations, netatmo_temperatures, netatmo_coords = fetch_netatmo_data()

    # Geocode Netatmo stations
    netatmo_geocoded = {}
    print("Geocoding Netatmo Stations:")
    for station, temp, coord in zip(netatmo_stations, netatmo_temperatures, netatmo_coords):
        if coord:
            lat, lon = coord
            netatmo_geocoded[station] = coord

    # Fetch additional IMGW data
    imgw_stations, imgw_temperatures = fetch_imgw_data()

    # Geocode IMGW stations
    imgw_coords = {}
    print("\nGeocoding IMGW Stations:")
    for station, temp in zip(imgw_stations, imgw_temperatures):
        cleaned_station = clean_station_name(station)
        coords, _ = get_coordinates(cleaned_station, geolocator)
        if coords:
            imgw_coords[station] = coords
            print(f'[IMGW] "{station}": "{temp}"')

    # Plot all stations with interpolation and error estimation
    plot_stations_with_boundary_and_interpolation(
        trax_stations, trax_temperatures, trax_coords, 
        netatmo_stations, netatmo_temperatures, netatmo_coords,
        imgw_stations, imgw_temperatures, imgw_coords
    )

In this code, it now takes really a lot time to generate the output, especially at "Performing Universal Kriging interpolation for residuals..." phase. Without changing any primary code functions and interpolation methodology, speed up the process as quick as it would be possible.