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

@register_keras_serializable()
def surface_runoff_amplifier(inputs):
    rain = inputs[:, 0]
    impervious = inputs[:, 1]
    rain_boost = tf.sigmoid((rain - 60) * 0.06)
    impervious_boost = tf.sigmoid((impervious - 0.6) * 10)
    return (1.0 + 0.3 * rain_boost * impervious_boost)[:, None]

@register_keras_serializable()
def drainage_penalty(inputs):
    dd = inputs[:, 2]
    return (1.0 - 0.4 * tf.sigmoid((dd - 3.5) * 2))[:, None]

@register_keras_serializable()
def convergence_suppressor(inputs):
    ci = inputs[:, 4]
    return (1.0 + 0.3 * tf.sigmoid((ci - 0.5) * 8))[:, None]

@register_keras_serializable()
def clip_modulation(x):
    return tf.clip_by_value(x, 0.7, 1.3)

CUSTOM_OBJECTS = {
    'drainage_penalty': drainage_penalty,
    'convergence_suppressor': convergence_suppressor,
    'surface_runoff_amplifier': surface_runoff_amplifier,
    'clip_modulation': clip_modulation
}