import logging
import numpy as np


def apply_land_sea_mask(predictions, features, class_to_index, mask_index=14, water_class_value=210):
    """
    Apply the land-sea mask to predictions and substitute water regions with the water class value.

    :param predictions: np.array, Predictions from the model
    :param features: np.array, Feature array that contains the sea-land mask as one of the features
    :param class_to_index: dict, The mapping of ground truth classes to their index
    :param mask_index: int, Index of the feature corresponding to the sea-land mask (default is 14)
    :param water_class_value: int, Value representing the water class in the ground truth (default is 210)
    :return: np.array, Adjusted predictions
    """
    # Check if 210 (water class) exists in the class_to_index mapping
    if water_class_value not in class_to_index:
        logging.info(f"Class {water_class_value} not found in class_to_index. Skipping water masking.")
        return predictions

    # Get the index of water class from the class_to_index mapping
    water_class_index = class_to_index[water_class_value]

    # Extract the sea-land mask from the features (1 for land, 0 for water)
    land_sea_mask = features[:, mask_index]
    
    # Apply the mask to predictions
    pred_masked = predictions * land_sea_mask
    
    # Replace water regions (where mask is 0) with the water class index
    pred_adjusted = np.where(land_sea_mask == 0, water_class_index, pred_masked)
    
    return pred_adjusted


def apply_land_sea_mask_inference(predictions, land_sea_mask, class_to_index, water_class_value=210):
    """
    Apply the land-sea mask to predictions and substitute water regions with the water class value.

    :param predictions: np.array, Predictions from the model (2D array of shape (lat_dim, lon_dim))
    :param land_sea_mask: np.array, 2D array representing the sea-land mask (1 for land, 0 for water)
    :param class_to_index: dict, The mapping of ground truth classes to their index
    :param water_class_value: int, Value representing the water class in the ground truth (default is 210)
    :return: np.array, Adjusted predictions
    """
    if water_class_value not in class_to_index:
        logging.info(f"Class {water_class_value} not found in class_to_index. Skipping water masking.")
        return predictions

    water_class_index = class_to_index[water_class_value]
    logging.info(f"Water class value (210) is mapped to index: {water_class_index}")

    # Apply the mask to predictions (1 for land, 0 for water)
    pred_masked = predictions * land_sea_mask

    # Replace water regions (where mask is 0) with the **class value 210**, not the index
    pred_adjusted = np.where(land_sea_mask == 0, water_class_value, pred_masked)

    logging.info(f"Number of water regions replaced: {(land_sea_mask == 0).sum()}")
    
    return pred_adjusted