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
from datetime import datetime
import pytz
import math

from matplotlib.offsetbox import OffsetImage, AnnotationBbox
import matplotlib.image as mpimg
from matplotlib.patches import Rectangle
from mpl_toolkits.axes_grid1 import make_axes_locatable

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|a739574f3ceeac299e7bb47cbe66b524"
    }
    
    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():
    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:
        print(f"Error fetching IMGW data: {response.status_code}")
        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'))

    print(f"Fetched {len(stations)} stations from IMGW data.")
    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):
                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, netatmo_coords, distance_threshold=0.01):
    """
    Remove Netatmo stations that are too close to Traxelektronik stations.
    """
    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 = []

    for netatmo_point in netatmo_points:
        distance, index = trax_tree.query(netatmo_point)
        if distance > distance_threshold:
            netatmo_filtered_coords.append(tuple(netatmo_point))

    print(f"Removed nearby duplicates. Remaining Netatmo stations: {len(netatmo_filtered_coords)}")
    return netatmo_filtered_coords

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 cap_local_residuals(stations_gdf, k=5, z_threshold=2):
    """
    Cap residuals based on local neighbor statistics to mitigate the influence of anomalous stations.
    """
    from sklearn.neighbors import NearestNeighbors

    neighbors = NearestNeighbors(n_neighbors=k+1)  # +1 to include the point itself
    coords = stations_gdf[['X', 'Y']].values
    neighbors.fit(coords)
    distances, indices = neighbors.kneighbors(coords)

    capped_residuals = stations_gdf['Residual'].copy()

    for i in range(len(stations_gdf)):
        neighbor_idxs = indices[i][1:]  # Exclude self
        neighbor_residuals = stations_gdf['Residual'].iloc[neighbor_idxs]

        median = neighbor_residuals.median()
        mad = np.median(np.abs(neighbor_residuals - median))

        if mad == 0:
            mad = 1e-6  # Prevent division by zero

        z = (stations_gdf.iloc[i]['Residual'] - median) / mad

        if z > z_threshold:
            capped_residuals.iloc[i] = median + z_threshold * mad
        elif z < -z_threshold:
            capped_residuals.iloc[i] = median - z_threshold * mad

    stations_gdf['Residual'] = capped_residuals
    print("Applied local residual capping to mitigate anomalous station influences.")

def mask_temperature_grid(temperature_grid, masovian_boundary_geo, transform):
    """
    Mask the temperature grid to ensure data is only within Masovian boundaries.
    """
    # Rasterize the Masovian boundary to create a mask
    mask = rasterize(
        [(geometry, 1) for geometry in masovian_boundary_geo],
        out_shape=temperature_grid.shape,
        transform=transform,
        fill=0,
        all_touched=True,
        dtype='uint8'
    )

    # Apply the mask: set temperature outside the boundary to NaN
    temperature_grid_masked = np.where(mask == 1, temperature_grid, np.nan)
    print("Applied Masovian boundary mask to the temperature grid.")
    return temperature_grid_masked

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_dict):
    # 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 = 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_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_dict and cleaned_temperature is not None:
            lat, lon = imgw_coords_dict[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="EPSG:4326")

    # 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 = read_and_clip_water_bodies(masovian_boundary)

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

    # Ensure geometries are in EPSG:4326
    trax_gdf_dem_crs = trax_gdf.to_crs("EPSG:4326")
    netatmo_gdf_dem_crs = netatmo_gdf.to_crs("EPSG:4326")
    imgw_gdf_dem_crs = imgw_gdf.to_crs("EPSG:4326")

    # Get elevation at station points using dem_data and dem_transform
    trax_elevations = get_elevation_at_points(
        trax_gdf_dem_crs.geometry,
        dem_data=dem_data,
        dem_transform=dem_transform
    )
    netatmo_elevations = get_elevation_at_points(
        netatmo_gdf_dem_crs.geometry,
        dem_data=dem_data,
        dem_transform=dem_transform
    )
    imgw_elevations = get_elevation_at_points(
        imgw_gdf_dem_crs.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_dem_crs.geometry,
        raster_data=forest_data,
        raster_transform=forest_transform
    )
    netatmo_forest = get_raster_value_at_points(
        netatmo_gdf_dem_crs.geometry,
        raster_data=forest_data,
        raster_transform=forest_transform
    )
    imgw_forest = get_raster_value_at_points(
        imgw_gdf_dem_crs.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_dem_crs.geometry,
        water_data=water_data,
        water_transform=water_transform
    )
    netatmo_water = get_water_body_at_points(
        netatmo_gdf_dem_crs.geometry,
        water_data=water_data,
        water_transform=water_transform
    )
    imgw_water = get_water_body_at_points(
        imgw_gdf_dem_crs.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

    # Add 'Priority' column to each GeoDataFrame
    trax_gdf['Priority'] = 1
    netatmo_gdf['Priority'] = 2
    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 = "EPSG:4326"
    stations_gdf = stations_gdf.dropna(subset=['Temperature', 'Elevation', 'Tree_Cover'])

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

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

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

    # Fit Random Forest Regression model
    print("Training Random Forest Regression model...")
    rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
    rf_model.fit(X_scaled, y)

    # 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

    # Instead of removing outliers, cap residuals to prevent anomalous influence
    residual_std = stations_gdf['Residual'].std()
    stations_gdf['Residual'] = stations_gdf['Residual'].clip(
        lower=-z_threshold * residual_std,
        upper=z_threshold * residual_std
    )

    print("Capped residuals to reduce the influence of anomalous stations.")

    # Apply local residual capping based on neighboring stations
    cap_local_residuals(stations_gdf, k=5, z_threshold=2)

    # Handle duplicates by prioritizing higher priority sources
    # Sort by priority (higher priority first)
    data = stations_gdf.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()

    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
    # Adjust grid_spacing as per your requirements
    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)

    # Set up variogram model with adjusted parameters
    variogram_model = 'spherical'
    # Adjust these parameters based on your data; you'll need to experiment with them
    variogram_parameters = {'nugget': 0.2, 'sill': 1.0, 'range': 0.05}

    # Perform Universal Kriging
    print("Performing Universal Kriging interpolation for residuals with adjusted variogram parameters...")
    uk = UniversalKriging(
        x=data['X'].values,
        y=data['Y'].values,
        z=data['Residual'].values,
        variogram_model=variogram_model,
        variogram_parameters=variogram_parameters,
        nlags=20,
        verbose=True
    )

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

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

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

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

    num_nans_in_predicted_temp = grid_df['Predicted_Temperature'].isnull().sum()
    print(f"Number of NaNs in Predicted Temperature: {num_nans_in_predicted_temp}")

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

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

    # Final temperature estimation (before topographic cooling)
    grid_df['Temperature'] = grid_df['Predicted_Temperature'] + grid_df['Residual']

    # Save a copy of the temperature before any adjustments
    grid_df['Temperature_Before_Adjustments'] = grid_df['Temperature'].copy()

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

    # Compute hour angle for cosine function
    hour_angle = ((hour % 24) / 24) * 2 * math.pi  # Convert hour to radians

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

    # 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_df['Elevation'].min()
    elevation_max = grid_df['Elevation'].max()
    grid_df['Normalized_Elevation'] = (grid_df['Elevation'] - elevation_min) / (elevation_max - elevation_min)

    # Compute the topographic cooling effect
    grid_df['Topographic_Cooling'] = grid_df['Normalized_Elevation'] * max_cooling_effect * night_weight

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

    # Apply the topographic cooling effect to the temperature
    grid_df['Temperature'] -= grid_df['Topographic_Cooling']

    # ECOSTRESS Adjustments Begin Here

    # Retrieve ECOSTRESS LST values at grid points (same as before)
    grid_ecostress = get_raster_value_at_points(
        [Point(x, y) for x, y in zip(grid_df['X'], grid_df['Y'])],
        raster_data=ecostress_data,
        raster_transform=ecostress_transform
    )
    grid_df['ECOSTRESS_LST'] = grid_ecostress

    # Handle NaNs in ECOSTRESS_LST by filling with the mean value
    grid_df['ECOSTRESS_LST'].fillna(grid_df['ECOSTRESS_LST'].mean(), inplace=True)

    # Compute ECOSTRESS anomalies (deviation from mean)
    ecostress_mean = grid_df['ECOSTRESS_LST'].mean()
    grid_df['ECOSTRESS_Anomaly'] = grid_df['ECOSTRESS_LST'] - ecostress_mean

    # Normalize anomalies to range [-1, 1]
    ecostress_anomaly_min = grid_df['ECOSTRESS_Anomaly'].min()
    ecostress_anomaly_max = grid_df['ECOSTRESS_Anomaly'].max()
    grid_df['ECOSTRESS_Anomaly_Normalized'] = (
        (grid_df['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

    # Define threshold for coldest areas (e.g., where ECOSTRESS anomaly is significantly low)
    coldest_threshold = -0.5  # You can adjust this value
    coldest_areas_mask = grid_df['ECOSTRESS_Anomaly_Normalized'] <= coldest_threshold

    # Apply anomalies with the new night_weight and only to coldest areas
    adjustment = grid_df['ECOSTRESS_Anomaly_Normalized'] * amplification_factor * night_weight

    # Apply adjustments only in coldest areas
    adjustment[coldest_areas_mask] = 0

    # Prevent ECOSTRESS adjustments on water bodies
    adjustment[grid_df['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_df['Temperature'] += adjustment

    # Apply cooling effect on water bodies
    cooling_effect = -2.0  # degrees Celsius
    grid_df.loc[grid_df['Water_Body'] == 1, 'Temperature'] += cooling_effect

    # ECOSTRESS Adjustments End Here

    # Apply Gaussian smoothing to the temperature grid to normalize large gradients
    # Create a masked 2D temperature array
    temperature_grid = np.full((nrows, ncols), np.nan)
    temperature_grid[grid_df['row_idx'], grid_df['col_idx']] = grid_df['Temperature'].values

    # Apply Gaussian filter while ignoring NaNs
    temperature_smoothed = gaussian_filter_nan(temperature_grid, sigma=1)

    # Replace the grid_df Temperature with the smoothed values
    # Flatten the smoothed grid to match grid_df's order
    smoothed_flat = temperature_smoothed.flatten()
    grid_df['Temperature_Smoothed'] = smoothed_flat[grid_df['row_idx'] * ncols + grid_df['col_idx']]

    grid_df['Temperature'] = grid_df['Temperature_Smoothed']
    grid_df.drop(columns=['Temperature_Smoothed'], inplace=True)

    # Mask the temperature grid to ensure data is only within Masovian boundaries
    masovian_boundary_geo = masovian_boundary.geometry
    temperature_smoothed_masked = mask_temperature_grid(temperature_smoothed, masovian_boundary_geo, transform)

    # Replace the grid_df Temperature with masked values
    grid_df['Temperature'] = temperature_smoothed_masked.flatten()[grid_df['row_idx'] * ncols + grid_df['col_idx']]

    # Read the custom color scale from the CSV file
    color_scale_path = 'input/color_scale.csv'
    if not os.path.exists(color_scale_path):
        print(f"Color scale file '{color_scale_path}' not found.")
        return None, None
    color_scale_df = pd.read_csv(color_scale_path)
    color_scale_df = color_scale_df.sort_values('value')
    
    # Ensure that 'value' column is numeric
    color_scale_df['value'] = pd.to_numeric(color_scale_df['value'], errors='coerce')
    color_scale_df = color_scale_df.dropna(subset=['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 based on color_scale.csv values
    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 based on color_scale.csv
    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 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, temperature_smoothed_masked)
    dem_extent = get_extent(dem_plot_transform, dem_plot_data)

    # Plotting Temperature Map without Contours
    fig1, ax1 = plt.subplots(figsize=(15, 15))

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

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

    # Overlay water bodies
    water_mask = np.where(water_data > 0, 1, 0)  # Adjust condition based on water_data values
    ax1.imshow(water_mask, extent=extent, origin='upper', cmap='Blues', alpha=0.5)  # Increased alpha to 0.5

    # Plot Masovian Voivodeship boundary with a light, dark stroke
    masovian_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')

    # Read the logo image
    logo_path = 'input/logo.png'
    if not os.path.exists(logo_path):
        print(f"Logo file '{logo_path}' not found.")
        return None, None
    logo_img = mpimg.imread(logo_path)

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

    # Position for fig1 logo
    ab_logo1 = AnnotationBbox(
        imagebox1,
        xy=(0.5, 0.02),  # Centered horizontally, slightly above bottom
        xycoords='figure fraction',
        frameon=False,
        box_alignment=(0.5, 0),  # Centered horizontally, aligned at bottom
    )
    fig1.add_artist(ab_logo1)

    # Introduce a subtle background gradient
    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: TempMeso v1.0 (Preview)   |   Resolution: 250 m",
        fontsize=18,  # Increased font size
        ha='center',
        va='top',
        fontname='DejaVu Sans',
        color='#666666'
    )

    # Adjust colorbar position and style
    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')

    # Plot temperature values at the coldest and warmest locations
    min_temp = np.nanmin(grid_df['Temperature'])
    max_temp = np.nanmax(grid_df['Temperature'])

    # Get the indices of the min and max temperatures
    min_indices = np.unravel_index(np.nanargmin(temperature_smoothed_masked), temperature_smoothed_masked.shape)
    max_indices = np.unravel_index(np.nanargmax(temperature_smoothed_masked), temperature_smoothed_masked.shape)

    # Convert grid indices to coordinates
    min_lon, min_lat = rasterio.transform.xy(transform, min_indices[0], min_indices[1], offset='center')
    max_lon, max_lat = rasterio.transform.xy(transform, max_indices[0], max_indices[1], offset='center')

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

    # Format temperature values without '°C'
    ax1.text(min_lon, min_lat, f"{min_temp:.1f}", **text_kwargs)
    ax1.text(max_lon, max_lat, f"{max_temp:.1f}", **text_kwargs)

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

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

    # Now, create the second plot with contours added

    # Plotting Temperature Map with Contours
    fig2, ax2 = plt.subplots(figsize=(15, 15))

    # Plot the interpolated temperature grid with ECOSTRESS adjustments
    img2 = ax2.imshow(temperature_smoothed_masked, extent=extent, origin='upper',
                      cmap=cmap, norm=norm, alpha=1)

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

    # Overlay water bodies
    ax2.imshow(water_mask, extent=extent, origin='upper', cmap='Blues', alpha=0.5)

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

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

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

    # Read the logo image
    logo_img = mpimg.imread(logo_path)

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

    # Position for fig2 logo
    ab_logo2 = AnnotationBbox(
        imagebox2,
        xy=(0.5, 0.02),  # Centered horizontally, slightly above bottom
        xycoords='figure fraction',
        frameon=False,
        box_alignment=(0.5, 0),  # Centered horizontally, aligned at bottom
    )
    fig2.add_artist(ab_logo2)

    # Introduce a subtle background gradient
    ax2.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)
    fig2.text(
        0.5, 0.94,  # Adjusted y-position for better placement
        "High-Resolution Mesoscale Temperature Analysis with Contours",
        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
    fig2.text(
        0.5, 0.91,  # Adjusted y-position to move it closer
        "Model: TempMeso v1.0 (Preview)   |   Resolution: 250 m",
        fontsize=18,  # Increased font size
        ha='center',
        va='top',
        fontname='DejaVu Sans',
        color='#666666'
    )

    # Adjust colorbar position and style
    divider2 = make_axes_locatable(ax2)
    cax2 = divider2.append_axes("right", size="3%", pad=0.1)
    cbar2 = fig2.colorbar(img2, cax=cax2)
    cbar2.set_ticks(ticks)
    cbar2.set_ticklabels(ticks.astype(int))
    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')

    # Plot temperature values at the coldest and warmest locations
    ax2.text(min_lon, min_lat, f"{min_temp:.1f}", **text_kwargs)
    ax2.text(max_lon, max_lat, f"{max_temp:.1f}", **text_kwargs)

    # Add contours to the plot
    min_temp_c = np.floor(np.nanmin(grid_df['Temperature']))
    max_temp_c = np.ceil(np.nanmax(grid_df['Temperature']))
    contour_levels = np.arange(min_temp_c, max_temp_c + 1, 1)

    # Generate contour lines only within the Masovian boundary
    masked_temperature = np.ma.masked_invalid(temperature_smoothed_masked)
    contours = ax2.contour(
        masked_temperature, levels=contour_levels, colors='black', linewidths=0.5,
        extent=extent, origin='upper'
    )
    contour_labels = ax2.clabel(
        contours, inline=False, fontsize=8, fmt='%d', colors='white',  # Removed '°C' from fmt
        inline_spacing=4, manual=False
    )
    for txt in contour_labels:
        txt.set_path_effects([
            patheffects.withStroke(linewidth=3, foreground='black'),
            patheffects.Normal()
        ])

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

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

    return temperature_smoothed_masked, transform

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_interpolated_temperature_at_point(lat, lon, temperature_grid, transform):
    """
    Given a latitude and longitude, return the interpolated temperature value at that point.
    
    Parameters:
    - lat, lon: Latitude and longitude of the point.
    - temperature_grid: The interpolated temperature grid (2D numpy array).
    - transform: Affine transform for the grid.
    
    Returns:
    - Interpolated temperature value at the point, or None if point is outside the grid.
    """
    from rasterio.transform import rowcol
    
    try:
        row, col = rowcol(transform, lon, lat)  # Note that rasterio uses (x, y), so (lon, lat)
        row = int(row)
        col = int(col)
        value = temperature_grid[row, col]
        if np.isnan(value):
            # Value is NaN, meaning outside the interpolated area
            return None
        else:
            return value.item()
    except (IndexError, ValueError):
        return None  # Point is outside the grid

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:
        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_clipped = masovian_boundary.to_crs(src.crs)
            print("Reprojected Masovian boundary to match water body raster CRS.")
        else:
            masovian_boundary_clipped = masovian_boundary
            print("Masovian boundary CRS matches water body raster CRS.")
        
        # Perform the clipping
        clipped_water, water_transform = rasterio.mask.mask(
            src, masovian_boundary_clipped.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_ecostress(masovian_boundary):
    ecostress_path = 'input/raw_ecostress.tif'
    with rasterio.open(ecostress_path) as src:
        print("Reading ECOSTRESS LST data (assumed to be in degrees Celsius).")
        # Read data
        ecostress_data = src.read(1, masked=True)
        ecostress_data = np.where(ecostress_data == src.nodata, np.nan, ecostress_data)
        ecostress_transform = src.transform
    return ecostress_data, ecostress_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:
        coords, was_cleaned = get_coordinates(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()

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

    # Geocode IMGW stations
    imgw_coords_dict = {}
    progress_bar = tqdm(total=len(imgw_stations), desc="Geocoding IMGW Stations", unit="station")
    for station in imgw_stations:
        coords, was_cleaned = get_coordinates(station, geolocator)
        if coords:
            imgw_coords_dict[station] = coords
        progress_bar.update(1)
    progress_bar.close()
    
    # Remove nearby duplicates from Netatmo data
    netatmo_filtered_coords = remove_nearby_duplicates(trax_coords, netatmo_coords, distance_threshold=0.01)
    netatmo_filtered_temperatures_filtered = [
        temp for coord, temp in zip(netatmo_coords, netatmo_temperatures) 
        if tuple(coord) in netatmo_filtered_coords
    ]
    netatmo_filtered_stations = [
        station for station, coord in zip(netatmo_stations, netatmo_coords) 
        if tuple(coord) in netatmo_filtered_coords
    ]

    # Update Netatmo data after removing duplicates
    netatmo_stations = netatmo_filtered_stations
    netatmo_temperatures = netatmo_filtered_temperatures_filtered
    netatmo_coords = netatmo_filtered_coords

    # After plotting, get the interpolated temperature at specific coordinates
    # Coordinates to test
    test_lat = 51.40118
    test_lon = 21.14805

    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_dict
    )

    interpolated_temp = get_interpolated_temperature_at_point(
        test_lat, test_lon, smoothed_temperature, transform
    )

    if interpolated_temp is not None:
        print(f"Interpolated temperature value at {test_lat}, {test_lon} is {interpolated_temp:.2f}")
    else:
        print(f"Coordinates {test_lat}, {test_lon} are outside the interpolated grid or have no data.")

There, it says that linearsegmentcolormap is not defined