# -*- coding: utf-8 -*-
"""
Created on Thu Jul 18 14:24:44 2024

@author: ehsan
"""

import netCDF4
import numpy as np
from keras.models import Sequential
from keras.layers import Dense, LSTM
from sklearn.preprocessing import MinMaxScaler
from keras.callbacks import EarlyStopping

# Load the NetCDF file
ncfile = netCDF4.Dataset('tmax_weekly_2024.nc')    #combined_tmax.nc or tmax_weekly_2024.nc

# Extract temperature data and relevant information
sst  = ncfile.variables['tmax'][:]
# Extract temperature data and relevant information
#sst = ncfile.variables['tmax'][-102:]
lat = ncfile.variables['lat'][:]
lon = ncfile.variables['lon'][:]

# Convert longitude values from [0, 360] to [-180, 180]
lon[lon > 180] -= 360


# Crop the data to cover pacific
lon_min, lon_max = 130, 157 # Adjusted longitude range for carrebean  -85, -60  
lat_min, lat_max = -15, 1 # Adjusted latitude range for carrebean   10, 30




# Combine the two longitude ranges
lon_indices = np.where((lon >= lon_min) & (lon <= lon_max))[0]
lat_indices = np.where((lat >= lat_min) & (lat <= lat_max))[0]

sst_cropped = sst[:, lat_indices][:, :, lon_indices]
lon_values_cropped = lon[lon_indices]
lat_values_cropped = lat[lat_indices]

# Example array XXX (replace this with your own array):
# Convert sst_cropped values to NaN where XXX has NaN values
XXX = sst_cropped[-1, :, :]    # last step
XXX[XXX < 0] = np.nan
sst_cropped[:, np.isnan(XXX)] = np.nan

# Build the LSTM model
model = Sequential()
model.add(LSTM(64, input_shape=(1, 1)))
model.add(Dense(1, activation='relu'))
model.compile(loss='mean_squared_error', optimizer='adam')

# Initialize an empty matrix to store the predicted values
predictions = np.zeros((len(lat_values_cropped), len(lon_values_cropped)))

# Define a scaling factor to amplify the predicted values
scaling_factor = 1.0

# Iterate over the cropped latitude and longitude values
for i, lat_val in enumerate(lat_values_cropped):
    for j, lon_val in enumerate(lon_values_cropped):
        # Find the indices for the given latitude and longitude
        lat_index = np.where(lat == lat_val)[0][0]
        lon_index = np.where(lon == lon_val)[0][0]

        # Check if the data contains any "no value" (9.97E+36) values
        if np.isnan(sst_cropped[:, i, j]).any() or (sst_cropped[:, i, j] == 9.97E+36).any():
            print(f"Skipping training for Latitude: {lat_val}, Longitude: {lon_val} due to NaN or 9.96921e+36 values.")
            continue

        # Reshape the cropped SST data to (time, 1)
        sst_reshaped = sst_cropped[:, i, j].reshape(-1, 1)

        # Normalize the SST data
        scaler = MinMaxScaler(feature_range=(-1, 1))
        sst_scaled = scaler.fit_transform(sst_reshaped)

        # Prepare the input data
        X = sst_scaled[:-1]  # Use all but the last time step as input
        y = sst_scaled[1:]   # Predict the next time step

        # Reshape the input data to match the LSTM input shape
        X = X.reshape(X.shape[0], 1, X.shape[1])

        # Early stopping callback to prevent overfitting
        early_stopping = EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=True)

        # Train the model with batch processing and early stopping
        model.fit(X, y, epochs=2, batch_size=32, validation_split=0.1, callbacks=[early_stopping])

        # Make predictions for the given pixel
        input_data = X[-1].reshape(1, 1, X.shape[2])
        predicted = model.predict(input_data)

        # Inverse transform the predicted data
        predicted_inverse = scaler.inverse_transform(predicted)

        # Scale the predicted value using the scaling factor
        predicted_scaled = predicted_inverse * scaling_factor

        # Store the predicted value in the matrix
        predictions[i, j] = predicted_scaled[0, 0]

        # Print the predicted SST for the current latitude and longitude
        print("Predicted SST for Latitude:", lat_val, "Longitude:", lon_val)
        print(predicted_scaled[0, 0])
        print()

# Save the predicted values as a NetCDF file
output_file = netCDF4.Dataset('predicted_values_cropped_CB1.nc', 'w', format='NETCDF4')
output_file.createDimension('lat', len(lat_values_cropped))
output_file.createDimension('lon', len(lon_values_cropped))
lat_var = output_file.createVariable('lat', 'f4', ('lat',))
lon_var = output_file.createVariable('lon', 'f4', ('lon',))
sst_var = output_file.createVariable('sst_predicted', 'f4', ('lat', 'lon'))
lat_var[:] = lat_values_cropped
lon_var[:] = lon_values_cropped
sst_var[:, :] = predictions
output_file.close()

# Save the trained model
model.save('lstm_model_CB.h5')


###############################

#########################mapping
import netCDF4
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.ticker import MultipleLocator

# Open netCDF file
ncfile = netCDF4.Dataset('predicted_values_cropped_CB1.nc')

# Extract SST data
sst = ncfile.variables['sst_predicted'][:]
lat = ncfile.variables['lat'][:]
lon = ncfile.variables['lon'][:]

# Convert XXX values less than 0 to NaN
sst[sst > 60] = np.nan  # no value 9.97E+36
# Convert XXX values less than 0 to NaN
sst[sst == 0] = np.nan  # no value 9.97E+36
sst[sst < 0] = np.nan  # no value 9.97E+36
# Clip the data to the desired range




scaling_factor=1.0
sst_clipped = sst* scaling_factor

# Create a map projection
projection = ccrs.PlateCarree()

fig, ax = plt.subplots(figsize=(10, 6), subplot_kw={'projection': projection})

# Add map features
ax.add_feature(cfeature.COASTLINE)
ax.add_feature(cfeature.BORDERS, linestyle=':')
ax.add_feature(cfeature.LAND, edgecolor='black')

# Define custom contour levels
levels = [0, 5, 10, 15, 20, 25, 30, 35, 40]

# Draw the data with custom contour levels
cf = ax.contourf(lon, lat, sst_clipped, levels=levels, cmap='coolwarm', transform=projection)

# Add colorbar close to the map
cbar = plt.colorbar(cf, ax=ax, orientation='vertical', fraction=0.024, pad=0.05)

# Add gridlines
gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, linewidth=0.5, color='gray', alpha=0.5, linestyle='--')
gl.xlabels_top = False
gl.ylabels_right = False

# Set equal intervals for gridlines
gl.xlocator = MultipleLocator(base=10.0)
gl.ylocator = MultipleLocator(base=10.0)

# Add annotations
plt.title('Predicted Maximum Temperature (°C)')
plt.xlabel('Longitude')
plt.ylabel('Latitude')

plt.show()







# ###################33observation

# import netCDF4
# import numpy as np
# import matplotlib.pyplot as plt
# import cartopy.crs as ccrs
# import cartopy.feature as cfeature
# from matplotlib.ticker import MultipleLocator

# # Load the NetCDF file
# ncfile = netCDF4.Dataset('tmax_weekly_2024.nc')  # Replace 'tmax_weekly_2024.nc' with your file name tmax.2024.nc'

# # Extract temperature data and relevant information
# sst  = ncfile.variables['tmax'][:]
# #sst = ncfile.variables['tmax'][-102:]
# lat = ncfile.variables['lat'][:]
# lon = ncfile.variables['lon'][:]

# # Convert longitude values from [0, 360] to [-180, 180]
# lon[lon > 180] -= 360

# # Crop the data to cover pacific
# lon_min, lon_max = 130, 157 # Adjusted longitude range for carrebean  -85, -60  
# lat_min, lat_max = -15, 1 # Adjusted latitude range for carrebean   10, 30


# lon_indices = np.where((lon >= lon_min) & (lon <= lon_max))[0]
# lat_indices = np.where((lat >= lat_min) & (lat <= lat_max))[0]

# sst_cropped = sst[:, lat_indices][:, :, lon_indices]
# lon_values_cropped = lon[lon_indices]
# lat_values_cropped = lat[lat_indices]

# # Example array XXX (replace this with your own array):
# XXX = sst_cropped[-1, :, :]    # last step

# # Define custom contour levels
# levels = [0, 5, 10, 15, 20, 25, 30, 35, 40]

# # Create a map projection
# projection = ccrs.PlateCarree()

# fig, ax = plt.subplots(figsize=(10, 6), subplot_kw={'projection': projection})

# # Add map features
# ax.add_feature(cfeature.COASTLINE)
# ax.add_feature(cfeature.BORDERS, linestyle=':')
# ax.add_feature(cfeature.LAND, edgecolor='black')

# # Draw the data with custom contour levels
# cf = ax.contourf(lon_values_cropped, lat_values_cropped, XXX, levels=levels, cmap='coolwarm', transform=projection)

# # Add colorbar close to the map
# cbar = plt.colorbar(cf, ax=ax, orientation='vertical', fraction=0.024, pad=0.05)
# cbar.set_label('Temperature (°C)')

# # Add gridlines
# gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, linewidth=0.5, color='gray', alpha=0.5, linestyle='--')
# gl.xlabels_top = False
# gl.ylabels_right = False

# # Set equal intervals for gridlines
# gl.xlocator = MultipleLocator(base=10.0)
# gl.ylocator = MultipleLocator(base=10.0)

# # Add annotations
# plt.title('Onserved Maximum Temperature(deg C)')
# plt.xlabel('Longitude')
# plt.ylabel('Latitude')

# plt.show()



