# !!!!pred from the question = capacity_kw!!!!
weather_states = pd.read_sql("SELECT stations.id, stations.capacity_kw, start, wind_speed_10m, wind_direction_10m, wind_speed_80m, wind_direction_80m, wind_speed_180m, wind_direction_180m FROM stations INNER JOIN weather_states ON stations.id = weather_states.station WHERE weather_states.source = 'openmeteo_forecast/history/best' AND stations.source = 'wind'", db_client)

grid_states = pd.read_sql("SELECT start, wind FROM grid_states", db_client)

def create_x_y(df: tuple[Any, pd.DataFrame]):
    start = df[1]["start"].iloc[0]
    res = df[1].sort_values("id").drop(["id", "start"], axis=1)
    temp_wind = grid_states.loc[grid_states["start"] == start]["wind"].to_list()
    wind_kw = temp_wind if len(temp_wind) >= 1 else None
    res_flat_df = pd.DataFrame(res.to_numpy().reshape((1, -1)))
    res_flat_df["wind_kw"] = wind_kw
    return res_flat_df

data = pd.concat(map(create_x_y, weather_states.groupby("start"))).dropna()
from sklearn.model_selection import train_test_split


data = data.astype("float32")
train, test = train, test = train_test_split(data.dropna(), test_size=0.2)

train_y = train.pop("wind_kw")
train_x = train


test_y = test.pop("wind_kw")
test_x = test

norm = tf.keras.layers.Normalization()
norm.adapt(train_x)

model = tf.keras.Sequential([
    norm,
    tf.keras.layers.Dense(16, activation="linear"),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(1, activation="linear"),
])


model.compile(
    optimizer=tf.keras.optimizers.legacy.Adam(0.001),
    metrics=[tf.keras.metrics.R2Score(dtype=tf.float32)],
    loss=tf.keras.losses.MeanSquaredError(),
)

model.fit(train_x, train_y, epochs=7, batch_size=2)

tf.keras.models.save_model(model, 'wind.keras')