import os

os.environ["KERAS_BACKEND"] = "tensorflow"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

from pathlib import Path

import keras
import numpy as np
from keras import layers

LOOKBACK = 28
TRAIN_END = 190
VALIDATION_END = 230
MODEL_PATH = Path("demand_forecast.keras")

keras.utils.set_random_seed(7)

rng = np.random.default_rng(42)
days = np.arange(260, dtype="float32")
weekly = 12.0 * np.sin(2.0 * np.pi * days / 7.0)
annual = 8.0 * np.sin(2.0 * np.pi * days / 365.0)
trend = 0.08 * days
promotion = ((days % 31) < 4).astype("float32")
noise = rng.normal(0.0, 2.0, size=days.shape[0]).astype("float32")
demand = 120.0 + trend + weekly + annual + (18.0 * promotion) + noise

train_mean = demand[:TRAIN_END].mean()
train_std = demand[:TRAIN_END].std()
demand_scaled = (demand - train_mean) / train_std

features = np.column_stack(
    [
        demand_scaled,
        promotion,
        np.sin(2.0 * np.pi * days / 7.0),
        np.cos(2.0 * np.pi * days / 7.0),
    ]
).astype("float32")


def make_windows(feature_rows, target_values, lookback):
    windows = []
    targets = []
    target_days = []
    for start in range(len(feature_rows) - lookback):
        target_day = start + lookback
        windows.append(feature_rows[start:target_day])
        targets.append(target_values[target_day])
        target_days.append(target_day)
    return (
        np.asarray(windows, dtype="float32"),
        np.asarray(targets, dtype="float32"),
        np.asarray(target_days, dtype="int32"),
    )


x_all, y_all, target_days = make_windows(features, demand_scaled, LOOKBACK)
train_mask = target_days < TRAIN_END
validation_mask = (target_days >= TRAIN_END) & (target_days < VALIDATION_END)
holdout_mask = target_days >= VALIDATION_END

x_train, y_train = x_all[train_mask], y_all[train_mask]
x_val, y_val = x_all[validation_mask], y_all[validation_mask]
x_holdout, y_holdout = x_all[holdout_mask], y_all[holdout_mask]
holdout_days = target_days[holdout_mask]

model = keras.Sequential(
    [
        layers.Input(shape=(LOOKBACK, features.shape[1])),
        layers.LSTM(24),
        layers.Dense(12, activation="relu"),
        layers.Dense(1),
    ]
)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.01),
    loss="mse",
    metrics=[keras.metrics.MeanAbsoluteError(name="mae")],
)

history = model.fit(
    x_train,
    y_train,
    validation_data=(x_val, y_val),
    epochs=8,
    batch_size=32,
    shuffle=False,
    verbose=0,
)
metrics = model.evaluate(x_val, y_val, verbose=0, return_dict=True)
forecast_scaled = model.predict(x_holdout[:7], verbose=0).reshape(-1)
forecast = (forecast_scaled * train_std) + train_mean
actual = (y_holdout[:7] * train_std) + train_mean

model.save(MODEL_PATH)
reloaded = keras.saving.load_model(MODEL_PATH)
reloaded_forecast = reloaded.predict(x_holdout[:1], verbose=0).reshape(-1)[0]
reload_delta = abs(reloaded_forecast - forecast_scaled[0])

print(f"backend: {keras.backend.backend()}")
print(f"train windows: {x_train.shape[0]}")
print(f"validation windows: {x_val.shape[0]}")
print(f"history keys: {', '.join(sorted(history.history.keys()))}")
print(f"validation mae: {metrics['mae'] * train_std:.2f} units")
print("holdout forecast:")
print("day  predicted  actual")
for day, predicted, observed in zip(holdout_days[:7], forecast, actual):
    print(f"{int(day):3d}  {predicted:9.1f}  {observed:6.1f}")
print(f"saved model: {MODEL_PATH}")
print(f"reload delta: {reload_delta:.6f}")