Should I also normalize day sin and day cos from here, in the context of a machine learning approach? The other topographic predictors are normalized from 0 to 1. import pandas as pd
import xarray as xr
import numpy as np
import os
import calendar
from sklearn.preprocessing import MinMaxScaler
import shutil

# Import settings from our config file
import config

import pandas as pd
import xarray as xr
import numpy as np
import os
import calendar
from sklearn.preprocessing import MinMaxScaler
import shutil

# Import settings from our config file
import config

def prepare_data(df, dem, predictor_names, output_dir):
    """
    Loads raw DEM and weather data, processes it for a multi-task (temp & precip) model,
    normalizes it, and saves it into the daily directory structure required for training.
    """
    print("Starting data preparation...")
    os.makedirs(output_dir, exist_ok=True)

    # --- 1. Prepare Predictor Attributes from DEM ---
    print("Pre-computing topographic predictors for all locations...")
    
    spatial_predictor_names = [p for p in predictor_names if p not in ['day_sin', 'day_cos', 'year_scaled']]
    dem_predictors_xr = dem[spatial_predictor_names]

    # --- Pre-compute for ALL unique stations in the dataframe once ---
    # This is more efficient than querying the DEM inside the loop.
    unique_stations_df = df[['cod', 'X', 'Y']].drop_duplicates().reset_index(drop=True)
    station_locations = unique_stations_df[['X', 'Y']].values
    
    station_spatial_attrs = dem_predictors_xr.sel(
        easting=xr.DataArray(station_locations[:, 0], dims='points'),
        northing=xr.DataArray(station_locations[:, 1], dims='points'),
        method='nearest'
    ).to_array().values.T

    # Create a new DataFrame with station codes and their spatial attributes
    station_attrs_df = pd.DataFrame(station_spatial_attrs, columns=spatial_predictor_names)
    station_attrs_df['cod'] = unique_stations_df['cod']
    
    # --- Pre-compute for the full target grid ---
    target_easting = dem['easting'].values
    target_northing = dem['northing'].values
    easting_grid, northing_grid = np.meshgrid(target_easting, target_northing)
    target_locations_s_star = np.stack([easting_grid.flatten(), northing_grid.flatten()], axis=1)
    
    target_spatial_attrs = dem_predictors_xr.sel(
        easting=xr.DataArray(target_locations_s_star[:, 0], dims='points'),
        northing=xr.DataArray(target_locations_s_star[:, 1], dims='points'),
        method='nearest'
    ).to_array().values.T

    # --- 2. Prepare Normalization Scalers for Target Variables ---
    print("Fitting scalers for temperature and precipitation...")
    scaler_temp = MinMaxScaler()
    df['prec_log'] = np.log1p(df['prec'])
    scaler_precip = MinMaxScaler()
    
    # Fit on a copy to avoid SettingWithCopyWarning
    scaler_temp.fit(df[['tavg']])
    scaler_precip.fit(df[['prec_log']])

    np.save(os.path.join(output_dir, "target_temp_min.npy"), scaler_temp.min_)
    np.save(os.path.join(output_dir, "target_temp_scale.npy"), scaler_temp.scale_)
    np.save(os.path.join(output_dir, "target_precip_min.npy"), scaler_precip.min_)
    np.save(os.path.join(output_dir, "target_precip_scale.npy"), scaler_precip.scale_)
    print("Saved normalization parameters.")

    # --- 3. Process and Save Data for Each Day ---
    print("Processing and saving data for each day...")
    dates = df['Date'].unique()
    
    # Group by date for efficiency
    grouped = df.groupby('Date')
    
    for date, df_day in grouped:
        date_str = pd.to_datetime(date).strftime("%Y-%m-%d")
        day_output_dir = os.path.join(output_dir, f"day_{date_str}")
        os.makedirs(day_output_dir, exist_ok=True)

        # --- A. Calculate Temporal Predictors ---
        day_of_year = date.timetuple().tm_yday
        n_days = 366 if calendar.isleap(date.year) else 365
        day_sin = np.sin(2 * np.pi * day_of_year / n_days)
        day_cos = np.cos(2 * np.pi * day_of_year / n_days)
        year_scaled = (date.year - 1962.5) / 61.5 # Adjust as needed
        temporal_values = np.array([day_sin, day_cos, year_scaled])

        # --- B. Prepare Context Data (Observed Stations) ---
        
        # *** ROBUSTNESS FIX: Use a merge to safely combine attributes and values ***
        df_day_merged = pd.merge(df_day, station_attrs_df, on='cod', how='left')
        
        # Sort by station code to ensure consistent order, although not strictly necessary after merge
        df_day_merged = df_day_merged.sort_values('cod').reset_index(drop=True)

        # Extract spatial attributes from the merged dataframe
        context_spatial_attrs_day = df_day_merged[spatial_predictor_names].values

        # Combine with temporal attributes
        num_context_stations = len(df_day_merged)
        temporal_attrs_context = np.tile(temporal_values, (num_context_stations, 1))
        context_attributes_X = np.concatenate([context_spatial_attrs_day, temporal_attrs_context], axis=-1)

        # Prepare the 2-column normalized target values from the merged dataframe
        context_temp_norm = scaler_temp.transform(df_day_merged[['tavg']])
        context_precip_log_norm = scaler_precip.transform(df_day_merged[['prec_log']])
        context_values_y_normalized = np.hstack([context_temp_norm, context_precip_log_norm])

        # --- C. Prepare Target Data (Full Grid) ---
        num_grid_points = len(target_locations_s_star)
        temporal_attrs_target = np.tile(temporal_values, (num_grid_points, 1))
        target_attributes_X_star = np.concatenate([target_spatial_attrs, temporal_attrs_target], axis=-1)

        # --- D. Save Files ---
        np.save(os.path.join(day_output_dir, "context_attributes_X.npy"), context_attributes_X.astype(np.float32))
        np.save(os.path.join(day_output_dir, "context_values_y_normalized.npy"), context_values_y_normalized.astype(np.float32))
        np.save(os.path.join(day_output_dir, "target_attributes_X_star.npy"), target_attributes_X_star.astype(np.float32))
        np.save(os.path.join(day_output_dir, "target_locations_s_star.npy"), target_locations_s_star.astype(np.float32))

    print(f"Data preparation complete. Processed {len(dates)} days.")


if __name__ == "__main__":
    # --- Configuration ---
    # These predictors must match the variable names in your NetCDF DEM file
    predictor_names = [
        'smooth_elev', 'topo_position', 'dist_to_coast', 
        'latwgs', 'lonwgs', 
        'day_sin', 'day_cos', 'year_scaled'
    ]
    output_dir = config.DATA_DIR_ALL

    # --- Execution ---
    # Clean the output directory to ensure a fresh start
    if os.path.exists(output_dir):
        print(f"Removing existing data directory: {output_dir}")
        shutil.rmtree(output_dir)

    print("Loading source data...")
    df = pd.read_parquet(config.SOURCE_WEATHER_PATH_ALL)
    dem = xr.open_dataset(config.SOURCE_DEM_PATH)

    # Filter data to the desired time range and type
    # The model will be trained on 'train' stations from 2016 onwards.
    df['Date'] = pd.to_datetime(df['Date'])
    df = df[df['Date'] >= pd.to_datetime("2016-01-01")]
    df = df[df['type'] == 'train']
    
    if df.empty:
        raise ValueError("The DataFrame is empty after filtering. Check date range and 'type' column.")

    prepare_data(df, dem, predictor_names, output_dir)