import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from tensorflow.keras.optimizers import Adam
import os
import shutil
import tensorflow as tf

# Disable GPU to avoid cuDNN/cuFFT/cuBLAS conflicts
os.environ["CUDA_VISIBLE_DEVICES"] = ""
print("CUDA_VISIBLE_DEVICES set to empty to disable GPU usage.")

# Check TensorFlow version
print(f"TensorFlow version: {tf.__version__}")

# Check disk space
available_space = shutil.disk_usage('.').free / (1024 * 1024)  # Free space in MB
if available_space < 10:
    raise RuntimeError(f"Insufficient disk space: {available_space:.2f} MB available. Need at least 10 MB.")
print(f"Available disk space: {available_space:.2f} MB")

# Ensure the data file exists
data_path = 'data/energy_data.csv'
if not os.path.exists(data_path):
    raise FileNotFoundError(f"Energy data file not found at {data_path}. Please ensure the file exists.")
print(f"Data file found at {data_path}")

# Load energy generation dataset
try:
    data = pd.read_csv(data_path)
except Exception as e:
    raise RuntimeError(f"Failed to load {data_path}: {e}")
print(f"Data loaded successfully. Number of rows: {len(data)}")

# Validate the dataset
if 'energy_generation' not in data.columns:
    raise ValueError("CSV file must contain an 'energy_generation' column.")
if len(data) < 60:
    raise ValueError(f"Dataset must have at least 60 rows for LSTM training. Found {len(data)} rows.")
energy_data = data['energy_generation'].values.reshape(-1, 1)
print("Dataset validated: 'energy_generation' column present.")

# Normalize the data
scaler = MinMaxScaler(feature_range=(0, 1))
energy_data = scaler.fit_transform(energy_data)
print("Data normalized successfully.")

# Prepare the dataset for LSTM
X, y = [], []
for i in range(60, len(energy_data)):
    X.append(energy_data[i-60:i, 0])
    y.append(energy_data[i, 0])
X, y = np.array(X), np.array(y)
print(f"Prepared data for LSTM: X shape = {X.shape}, y shape = {y.shape}")

# Reshape X for LSTM input
X = X.reshape((X.shape[0], X.shape[1], 1))
print(f"Reshaped X for LSTM input: X shape = {X.shape}")

# Define the LSTM model
model = Sequential()
model.add(LSTM(units=50, return_sequences=True, input_shape=(X.shape[1], 1)))
model.add(LSTM(units=50, return_sequences=False))
model.add(Dense(units=1))
print("LSTM model defined successfully.")

# Compile and train the model
model.compile(optimizer=Adam(), loss='mean_squared_error')
model.fit(X, y, epochs=10, batch_size=32, verbose=1)
print("Model training completed.")

# Ensure the models directory exists
model_path = 'models/lstm_energy_model'
os.makedirs(model_path, exist_ok=True)
if not os.access(model_path, os.W_OK):
    raise PermissionError("No write permissions for the 'models/lstm_energy_model' directory.")
print("Models directory created and writable.")

# Remove any existing model directory
if os.path.exists(model_path) and os.path.isdir(model_path):
    shutil.rmtree(model_path)
    print(f"Removed existing directory at {model_path}")
os.makedirs(model_path, exist_ok=True)

# Save the trained model in SavedModel format
model.save(model_path, save_format='tf')
print(f"✅ LSTM model trained and saved as SavedModel at {model_path}")