pipelines/model_preparation_pipeline.py

from prefect import flow, task
from prefect.artifacts import create_table_artifact
from prefect.task_runners import ThreadPoolTaskRunner
import numpy as np
import pandas as pd
from pipelines.preprocessing_pipeline import save_to_file
from prefect.logging import get_run_logger
import yaml
import pickle

from typing import Literal

with open('config/params.yml') as file:
    config = yaml.safe_load(file)

with open('config/fieldsensor_irrigator_mapping_anonym.yaml') as file:
    fieldsensor_irrigator_mapping = yaml.safe_load(file)


data_availability = config['data_availability']

datetime_col = config['datetime_col']
value_col = config['value_col']
datastream_id_col = config['datastream_id_col']
datastream_name_col = config['datastream_name_col']
sensor_type_col = config['sensor_type_col']
ground_offset_col = config['ground_offset_col']

resampling_window = config['resampling_window']
days_weather_forecast = config['days_weather_forecast']

meteo_method = config['meteo_method']


### Pipeline nodes

@task(task_run_name="get_consortia")
def get_consortia():
    """
    Retrieves the list of consortia from the configuration.

    Returns:
        List[str]: List of consortia names.
    """
    logger = get_run_logger()
    consortia = config['consortia']
    logger.info(f"We work with the following consortia: {consortia}")
    return consortia


@task(task_run_name="read_input_data_{consortium_name}")
def read_input_data(consortium_name, has_weather_data=True, has_crop_data=True, has_soil_data=True, has_remote_sensing_data=True):
    field_sensor_data = pd.read_parquet(f'data//03_primary//field_sensor_data_{consortium_name}.parquet')
    irrigation_data = pd.read_parquet(f'data//03_primary//irrigation_data_{consortium_name}.parquet')
    if has_weather_data:
        weather_data = pd.read_parquet(f'data//03_primary//weather_data_{consortium_name}.parquet')
    else:
        weather_data = None
    historical_weather_data = pd.read_parquet(f'data//03_primary//historical_weather_data_{consortium_name}.parquet')
    forecasted_weather_data = pd.read_parquet(f'data//03_primary//forecasted_weather_data_{consortium_name}.parquet')
    if has_crop_data:
        with open(f'data//03_primary//crop_type_data_{consortium_name}.pickle', 'rb') as handle:
            crop_type_data = pickle.load(handle)
    if has_soil_data:
        soil_type_data = pd.read_parquet(f'data//03_primary//soil_type_data_{consortium_name}.parquet')
    else:
        soil_type_data = None
    locations_ids = pd.read_parquet(f'data//03_primary//locations_ids_{consortium_name}.parquet')

    if has_remote_sensing_data:
        #remote_sensing_data = xr.open_dataset(f'data//03_primary//remote_sensing_data_resampled_{consortium_name}.nc',
        #                                      engine='netcdf4')
        remote_sensing_data = pd.read_parquet(f'data//03_primary//remote_sensing_data_final_{consortium_name}.parquet')
    else:
        remote_sensing_data = None


    info_artifact = [
        {'variable': 'Length of field sensors DataFrame', 'amount': len(field_sensor_data) }, 
        {'variable': 'Length of irrigation data DataFrame', 'amount': len(irrigation_data) }, 
        {'variable': 'Length of weather data DataFrame', 'amount': len(weather_data) if weather_data is not None else None},
        {'variable': 'Length of historical weather data DataFrame', 'amount': len(historical_weather_data) }, 
        {'variable': 'Length of non-NaN field sensors DataFrame', 'amount': len(field_sensor_data.dropna()) }, 
        {'variable': 'Length of non-NaN irrigation data DataFrame', 'amount': len(irrigation_data.dropna()) }, 
        {'variable': 'Length of non-NaN weather data DataFrame', 'amount': len(weather_data.dropna()) if weather_data is not None else None},
        {'variable': 'Length of non-NaN historical weather data DataFrame', 'amount': len(historical_weather_data.dropna()) }, 
    ]
    create_table_artifact(
        key=f"data-{consortium_name}-info",
        table=info_artifact,
        description= "# Data info!"
    )

    return field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data, crop_type_data, soil_type_data, locations_ids, remote_sensing_data


def get_type(name):
    if 'ELMED' in name:
        return 'ELMED'
    elif 'TN' in name:
        return 'TN'
    elif '_tens_' in name:
        return 'TN'
    

def get_sector(sensor_name):
    candidates = [y for y in sensor_name.split(' ')[0].split('_') if 'sector' in y]
    if len(candidates) == 1:
        return candidates[0]
    else:
        return None
    

def get_management(sensor_name):
    if '_management2' in sensor_name.lower():
        return 'management2'
    elif '_management1' in sensor_name.lower():
        return 'management1'
    elif 'management0' in sensor_name.lower():
        return 'management1'
    else:
        return ''
    

def get_associated_sensors(sensors):
    sensors['type'] = sensors[datastream_name_col].transform(get_type)
    sensors['sector'] = sensors[datastream_name_col].transform(lambda x: get_sector(x))
    sensors['management'] = sensors[datastream_name_col].transform(lambda x: get_management(x))
    sensors['associated_sensors'] = sensors.apply(lambda x: sensors[(sensors['type'] == x['type'])&(sensors['sector'] == x['sector'])&(sensors['management'] == x['management'])&(sensors[datastream_name_col] != x[datastream_name_col])][datastream_name_col].values, axis=1)

    # if all(sensors['associated_sensors'].transform(lambda x: len(x) == 0)):
    #     set_sensors = set(sensors[datastream_name_col])
    #     dict_association = {}
    #     for sensorA in set_sensors:
    #         for sensorB in set_sensors:
    #             candidatesA = set(sensorA.split(' ')[0].split('_'))
    #             candidatesB = set(sensorB.split(' ')[0].split('_'))

    #             if len(candidatesA - candidatesB) == 1 and len(candidatesB - candidatesA) == 1:
    #                 el_diff = list(candidatesA - candidatesB)[0]
    #                 if el_diff.isnumeric() and float(el_diff) < 2000:
    #                     dict_association[sensorA] = sensorB
    #     sensors['associated_sensors'] = sensors[datastream_name_col].transform(lambda x: [dict_association[x]])
        
    return sensors


@task(task_run_name = 'get_sensors_meteo_{consortium_name}')
def get_sensors_meteo(
        consortium_name:str,
        meteo_df: pd.DataFrame,
        location_ids: pd.DataFrame,
        sensor_df: pd.DataFrame,
        method: Literal['closest', 'interpolated'] = 'closest',
        idw_power: float = 2.0
) -> pd.DataFrame:
    """
    Enrich field sensor locations with meteo data from weather sensors.

    Parameters:
    -----------
    consortium_name: str
        Name of the processing consortium
    meteo_df : pd.DataFrame
        Time-series meteo data with datetime index and columns like
        'Air temperature F0057_mean', 'Precipitation F5108', etc.
    location_ids : pd.DataFrame
        Location data with columns: 'datastream_name', 'x' (lon), 'y' (lat)
    sensor_df : pd.DataFrame
        Field sensor data containing 'datastream_name' column. Unique values
        identify the target locations to enrich.
    method : str
        'closest' for nearest sensor, 'interpolated' for IDW
    idw_power : float
        Power parameter for inverse distance weighting (default: 2.0)

    Returns:
    --------
    pd.DataFrame
        Enriched dataframe with columns: 'result_time', 'location_id',
        and sensor-agnostic meteo features (e.g., 'temperature_mean_local_meteo')
    """

    # make sure - datetime index
    if datetime_col in meteo_df.columns:
        meteo_df = meteo_df.set_index(datetime_col)

    # Step 1 --> Identify meteo sensor columns and parse them
    meteo_columns = {}  # {base_metric: [(full_col_name, sensor_id, suffix), ...]}

    for col in meteo_df.columns:
        # matching datastream in location_ids
        matching = location_ids[
            location_ids[datastream_name_col].apply(lambda x: x in col if pd.notna(x) else False)
        ]

        if len(matching) > 0:
            # get base metric and suffix
            datastream_name = matching.iloc[0][datastream_name_col]

            # remove datastream_name from column to get suffix
            suffix = col.replace(datastream_name, '').strip('_')

            # norm base metric name (lowercase, remove extra spaces)
            base_metric = datastream_name.rsplit('_', 1)[0].lower().replace(' ', '_')

            # create full metric name with suffix
            if suffix:
                metric_key = f"{base_metric}_{suffix}"
            else:
                metric_key = base_metric

            if metric_key not in meteo_columns:
                meteo_columns[metric_key] = []

            meteo_columns[metric_key].append({
                'col_name': col,
                datastream_name_col: datastream_name,
                'sensor_id': datastream_name.rsplit('_', 1)[1] if ' ' in datastream_name else datastream_name
            })

    # Step 2 --> get meteo sensor locations
    meteo_sensors = {}  # {datastream_name: (x, y)}
    for metric_data in meteo_columns.values():
        for sensor_info in metric_data:
            ds_name = sensor_info[datastream_name_col]
            if ds_name not in meteo_sensors:
                loc = location_ids[location_ids[datastream_name_col] == ds_name]
                if len(loc) > 0:
                    meteo_sensors[ds_name] = (loc.iloc[0]['x'], loc.iloc[0]['y'])

    # Step 3 --> get target locations (field sensors)
    target_datastream_names = sensor_df[datastream_name_col].unique()

    # get locations for these target sensors
    target_locations = location_ids[
        location_ids[datastream_name_col].isin(target_datastream_names)
    ].copy()

    # Step 4 --> Build enriched dataframe
    results = []

    for _, target_loc in target_locations.iterrows():
        target_name = target_loc[datastream_name_col]
        target_x, target_y = target_loc['x'], target_loc['y']

        # each timestamp
        for timestamp in meteo_df.index:
            row_data = {
                datetime_col: timestamp,
                datastream_name_col: target_name
            }

            # each metric
            for metric_key, sensor_list in meteo_columns.items():
                # obtain available sensors with data for this timestamp
                available_sensors = []

                for sensor_info in sensor_list:
                    col_name = sensor_info['col_name']
                    ds_name = sensor_info[datastream_name_col]

                    if ds_name in meteo_sensors and pd.notna(meteo_df.loc[timestamp, col_name]):
                        sensor_x, sensor_y = meteo_sensors[ds_name]
                        value = meteo_df.loc[timestamp, col_name]

                        # calculate distance
                        distance = np.sqrt((target_x - sensor_x) ** 2 + (target_y - sensor_y) ** 2)

                        available_sensors.append({
                            'value': value,
                            'distance': distance
                        })

                # calculate local meteo value
                if len(available_sensors) == 0:
                    local_value = np.nan
                elif method == 'closest' or len(available_sensors) == 1:
                    # use closest sensor
                    closest = min(available_sensors, key=lambda s: s['distance'])
                    local_value = closest['value']
                else:  # interpolated with IDW
                    # Inverse Distance Weighting
                    total_weight = 0
                    weighted_sum = 0

                    for sensor in available_sensors:
                        if sensor['distance'] == 0:
                            # target is at sensor location
                            local_value = sensor['value']
                            break
                        else:
                            weight = 1 / (sensor['distance'] ** idw_power)
                            weighted_sum += weight * sensor['value']
                            total_weight += weight
                    else:
                        # normal IDW calculation
                        local_value = weighted_sum / total_weight if total_weight > 0 else np.nan

                row_data[f"{metric_key}_local_meteo"] = local_value

            results.append(row_data)

    # create final dataframe
    result_df = pd.DataFrame(results)

    # columns: result_time, location_id, then all meteo features
    meteo_cols = [col for col in result_df.columns if col.endswith('_local_meteo')]
    result_df = result_df[[datetime_col, datastream_name_col] + meteo_cols]

    return result_df

@task(task_run_name='merge_all_data_{consortium_name}')
def merge_all_data(consortium_name, field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data, remote_sensing_data, location_ids):
    associated_irrigators = []
    for year in fieldsensor_irrigator_mapping[consortium_name]:
        tmp = pd.DataFrame(fieldsensor_irrigator_mapping[consortium_name][year].items(), columns=[datastream_name_col, 'irrigator'])
        tmp['year'] = year
        associated_irrigators.append(tmp)
    associated_irrigators = pd.concat(associated_irrigators)

    field_sensor_data['year'] = field_sensor_data.index.year

    full_table = field_sensor_data.drop(columns=[datastream_id_col]).reset_index().merge(
        associated_irrigators,
        on = [datastream_name_col, 'year'],
        how = 'left'
    ).drop(columns=['year']).merge(
        irrigation_data.drop(columns=[datastream_id_col]).reset_index().rename(columns={datastream_name_col: 'irrigator', value_col: 'irrigation'}),
        on = ['irrigator', datetime_col],
        how = 'left'
    ).merge(
        forecasted_weather_data,
        on = [datetime_col, datastream_name_col],
        how= 'left'
    ).merge(
        historical_weather_data.reset_index(),
        on = [datetime_col, datastream_name_col],
        how= 'left'
    )

    if weather_data is not None:   # TODO: ADD OPTION OF CLOSER/INTERPOLATED
        full_table = full_table.merge(
            weather_data,
            on=[datetime_col, datastream_name_col],
            how='left'
        )
    
    if remote_sensing_data is not None: # TODO: remove when we have data for Tres
        # add band values (short term solution)
        indices = [
            "ndvi", "grvi", "rvi", "rgi", "aci", "maci", "gndvi", "ngrdi", "ngbdi", "bgvi", "brvi",
            "wi", "varig", "gli", "g_perc", "ndmi", "ndwi", "reci", "ndre_lower_end",
            "ndre_upper_end", "msavi", "arvi", "sipi", "gci"
        ]

        full_table = full_table.merge(
            remote_sensing_data[[datastream_name_col ,datetime_col] + indices],
            on=[datastream_name_col, datetime_col],
            how='left'
        )

    full_table['irrigation'] = full_table['irrigation'].fillna(0)

    sensors = pd.DataFrame(full_table[datastream_name_col].drop_duplicates().reset_index(drop=True))
    sensors = get_associated_sensors(sensors)

    full_table = full_table.merge(
        sensors[[datastream_name_col, 'associated_sensors']],
        on = [datastream_name_col],
        how = 'left'
    )
    
    return full_table


@task(task_run_name='add_irrigation_forecast')
def add_irrigation_forecast(full_table):
    for n in range(days_weather_forecast * 24 // int(resampling_window.split('h')[0])):
        full_table[f'forecasted_irrigation_next_{n+1}period'] = full_table['irrigation'].shift(-n)
        full_table[f'forecasted_irrigation_next_{n+1}period'] = full_table[f'forecasted_irrigation_next_{n+1}period'].fillna(0)

    return full_table


@flow(name='model_preparation_pipeline', retries=1, task_runner=ThreadPoolTaskRunner())
def model_preparation_pipeline(override_consortia=None, override_data_availability=None) -> list[str]:
    logger = get_run_logger()
    logger.info(f'Starting model preparation pipeline!')

    if not override_consortia:
        consortia = get_consortia()
    else:
        consortia = override_consortia

    for consortium_name in consortia:
        if override_data_availability:
            data_availability_consortium = override_data_availability
        else:
            data_availability_consortium = data_availability[consortium_name]

        # Read data
        field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data, crop_type_data, soil_type_data, location_ids, remote_sensing_data = read_input_data(consortium_name, **data_availability_consortium)


        # transform local meteo in single var per sensor
        if data_availability_consortium['has_weather_data']:
            weather_data = get_sensors_meteo(
                consortium_name=consortium_name,
                meteo_df = weather_data,
                location_ids = location_ids,
                sensor_df = field_sensor_data,
                method = meteo_method)

        # Nodes here
        full_table = merge_all_data.submit(consortium_name, field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data, remote_sensing_data, location_ids).result()
        full_table = add_irrigation_forecast(full_table)

        save_to_file(
            df=full_table,
            output_file=f'data//04_model_input//full_table_{consortium_name}.parquet',
        )

    return


if __name__ == "__main__":
    model_preparation_pipeline()