import tensorflow as tf
from tensorflow.keras import layers, models, callbacks
from tensorflow.keras.saving import register_keras_serializable
import numpy as np

@register_keras_serializable()
def rainfall_proximity_penalty(inputs):
    rainfall = inputs[:, 0]
    distance = inputs[:, 4]
    proximity_score = tf.sigmoid((150 - distance) * 0.04)
    rainfall_score = tf.sigmoid((rainfall - 90) * 0.3)
    return (rainfall_score * proximity_score)[:, None]

@register_keras_serializable()
def flood_risk_booster(inputs):
    slope = inputs[:, 3]
    rainfall = inputs[:, 0]
    slope_boost = tf.sigmoid((slope - 2.0) * 1.5)
    rain_boost = tf.sigmoid((rainfall - 60) * 0.25)
    return (1.0 + 0.25 * slope_boost * rain_boost)[:, None]

@register_keras_serializable()
def flood_suppression_mask(inputs):
    elevation = inputs[:, 2]
    rainfall = inputs[:, 0]
    flatness = tf.sigmoid((elevation - 9.0) * 0.6)
    dryness = tf.sigmoid((20.0 - rainfall) * 0.2)
    return (1.0 - 0.3 * flatness * dryness)[:, None]

CUSTOM_OBJECTS = {
    "rainfall_proximity_penalty": rainfall_proximity_penalty,
    "flood_risk_booster": flood_risk_booster,
    "flood_suppression_mask": flood_suppression_mask
}