import pandas as pd
from loading_climate_data import load_and_process_climate_data
from loading_crop_data import engineer_crop_features
from loading_ndvi_data import process_ndvi_data
from loading_soil_data import get_soil_data


def combine_datasets(crop_data_file,
                     soil_data_file,
                     climate_data_file,
                     ndvi_data_file,
                     ndvi=False):
    dataset_loaders = {
        'soil': lambda: get_soil_data(soil_data_file),
        'climate': lambda: load_and_process_climate_data(climate_data_file),
        'ndvi': lambda: process_ndvi_data(ndvi_data_file),
        'crop': lambda: engineer_crop_features(crop_data_file)
    }

    merge_keys = {
        'soil': ['PostalCode'],
        'climate': ['PostalCode', 'Year'],
        'ndvi': ['PostalCode', 'Year'],
        'crop': ['PostalCode', 'Year']
    }

    combined_data = dataset_loaders['crop']()

    # Merge climate data
    climate_data = dataset_loaders['climate']()
    combined_data = pd.merge(combined_data, climate_data, on=merge_keys['climate'], how='left')

    # Merge NDVI data if required
    if ndvi:
        ndvi_data = dataset_loaders['ndvi']()
        combined_data = combined_data[combined_data['Year'] >= 2000]
        combined_data = pd.merge(combined_data, ndvi_data, on=merge_keys['ndvi'], how='left')

    # Merge soil data
    soil_data = dataset_loaders['soil']()
    combined_data = pd.merge(combined_data, soil_data, on=merge_keys['soil'], how='left')
    postal_codes = combined_data['PostalCode']
    years = combined_data['Year']

    # Drop irrelevant or redundant columns
    features_to_exclude = ['PostalCode', 'Year', 'SoilID']

    combined_data = combined_data.drop(columns=features_to_exclude, errors='ignore')

    return combined_data


def combine_dataset_pc(crop_data_file,
                       soil_data_file,
                       climate_data_file,
                       ndvi_data_file,
                       ndvi=False):
    dataset_loaders = {
        'soil': lambda: get_soil_data(soil_data_file),
        'climate': lambda: load_and_process_climate_data(climate_data_file),
        'ndvi': lambda: process_ndvi_data(ndvi_data_file),
        'crop': lambda: engineer_crop_features(crop_data_file)
    }

    merge_keys = {
        'soil': ['PostalCode'],
        'climate': ['PostalCode', 'Year'],
        'ndvi': ['PostalCode', 'Year'],
        'crop': ['PostalCode', 'Year']
    }

    combined_data = dataset_loaders['crop']()

    # Merge climate data
    climate_data = dataset_loaders['climate']()
    combined_data = pd.merge(combined_data, climate_data, on=merge_keys['climate'], how='left')

    # Merge NDVI data if required
    if ndvi:
        ndvi_data = dataset_loaders['ndvi']()
        combined_data = combined_data[combined_data['Year'] >= 2000]
        combined_data = pd.merge(combined_data, ndvi_data, on=merge_keys['ndvi'], how='left')

    # Merge soil data
    soil_data = dataset_loaders['soil']()
    combined_data = pd.merge(combined_data, soil_data, on=merge_keys['soil'], how='left')

    # Drop irrelevant or redundant columns
    features_to_exclude = ['Year', 'SoilID']
    combined_data = combined_data.drop(columns=features_to_exclude, errors='ignore')

    return combined_data
