# load imports
import xarray as xr
import matplotlib.pyplot as plt
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

# load data with xarray
ds_radio = xr.open_dataset('radiometer.nc')
ds_radio = ds_radio[['vap', 'qc_vap']]

ds_mfrs = xr.open_dataset('mfrs.nc')
ds_mfrs

# remove data where the value of the QC bit is not 0 (remove the bad data)
ds_mfrs = ds_mfrs.where(ds_mfrs.qc_direct_normal_narrowband_filter6 == 0, drop=True)
ds_mfrs = ds_mfrs.where(ds_mfrs.qc_hemisp_narrowband_filter6 == 0, drop=True)
ds_mfrs = ds_mfrs.where(ds_mfrs.qc_direct_horizontal_narrowband_filter6 == 0, drop=True)

ds_radio = ds_radio.where(ds_radio.qc_vap == 0, drop=True)

# create time mask where times match between the datasets
radio_time = ds_radio.time.values
mfrs_time = ds_mfrs.time.values

mask1 = np.isin(radio_time, mfrs_time)
mask2 = np.isin(mfrs_time, radio_time)

ds_combined = xr.merge([ds_radio, ds_mfrs], join='inner')
ds_combined

# Create a 1-row, 3-column plot layout
fig, axes = plt.subplots(2, 2, figsize=(18, 20), sharey=True)

# 1. Direct Normal
axes[0,0].scatter(ds_combined.direct_normal_narrowband_filter6.values, ds_combined.vap.values, alpha=0.4)
axes[0,0].set_title("Direct Normal vs VAP")
axes[0,0].set_xlabel("Filter 6 Values")
axes[0,0].set_ylabel("VAP")
axes[0,0].grid(True)

# 2. Hemispheric
axes[1,0].scatter(ds_combined.hemisp_narrowband_filter6.values, ds_combined.vap.values, alpha=0.4, color='orange')
axes[1,0].set_title("Hemispheric vs VAP")
axes[1,0].set_xlabel("Filter 6 Values")
axes[1,0].grid(True)

# 3. Direct Horizontal
axes[0,1].scatter(ds_combined.direct_horizontal_narrowband_filter6.values, ds_combined.vap.values, alpha=0.4, color='green')
axes[0,1].set_title("Direct Horizontal vs VAP")
axes[0,1].set_xlabel("Filter 6 Values")
axes[0,1].grid(True)

plt.tight_layout()
plt.show()

# split the datasets into training, validation, and testing
x_data = ds_combined[['direct_normal_narrowband_filter6', 'hemisp_narrowband_filter6',
                      'direct_horizontal_narrowband_filter6']].to_array().values.T
y_data = ds_combined.vap.values / 10

x_train = x_data[:int(x_data.shape[0]*0.8),:]
x_val = x_data[int(x_data.shape[0]*0.8):int(x_data.shape[0]*0.9),:]
x_test = x_data[int(x_data.shape[0]*0.9):,:]

y_train = y_data[:int(x_data.shape[0]*0.8)]
y_val = y_data[int(x_data.shape[0]*0.8):int(x_data.shape[0]*0.9)]
y_test = y_data[int(x_data.shape[0]*0.9):]

print(x_train.shape, y_train.shape)
print(x_test.shape, y_test.shape)

# Target variables (y) must still be reshaped to (-1, 1) to match the output dimension
X_train_tensor = torch.tensor(x_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)

X_val_tensor = torch.tensor(x_val, dtype=torch.float32)
y_val_tensor = torch.tensor(y_val, dtype=torch.float32).reshape(-1, 1)

X_test_tensor = torch.tensor(x_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32).reshape(-1, 1)

print("New Feature Shape:", X_train_tensor.shape)  # Should read: torch.Size([time, 3])
print("New Target Shape:", y_train_tensor.shape)    # Should read: torch.Size([time, 1])

# Rebuild your DataLoaders
BATCH_SIZE = 64
train_loader = DataLoader(TensorDataset(X_train_tensor, y_train_tensor), batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val_tensor, y_val_tensor), batch_size=BATCH_SIZE, shuffle=False)
test_loader = DataLoader(TensorDataset(X_test_tensor, y_test_tensor), batch_size=BATCH_SIZE, shuffle=False)

# Initialize Model, Loss, and Optimizer
model = nn.Sequential(
    nn.Linear(3, 32),
    nn.ReLU(),
    nn.Linear(32, 1)
)

criterion = nn.MSELoss()  # Built-in softmax + negative log-likelihood
optimizer = optim.Adam(model.parameters(), lr=0.001)

NUM_EPOCHS = 20

for epoch in range(NUM_EPOCHS):
    model.train()
    running_train_loss = 0.0
    
    for batch_X, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_X)           # outputs shape will naturally be (64, 1)
        loss = criterion(outputs, batch_y)  # batch_y shape is (64, 1)
        loss.backward()
        optimizer.step()
        
        running_train_loss += loss.item() * batch_X.size(0)
        
    epoch_train_loss = running_train_loss / len(train_loader.dataset)
    print(f"Epoch {epoch+1:02d}/{NUM_EPOCHS} | Train MSE: {epoch_train_loss:.4f}")
    
    # --- VALIDATION PHASE ---
    model.eval()  # Sets the model to evaluation mode (turns off dropout/batchnorm)
    running_val_loss = 0.0
    
    with torch.no_grad():  # Disables gradient calculation to save memory and speed up
        for batch_X, batch_y in val_loader:
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            running_val_loss += loss.item() * batch_X.size(0)
            
    epoch_val_loss = running_val_loss / len(val_loader.dataset)
    
    print(f"Epoch {epoch+1}/{NUM_EPOCHS} | Train Loss: {epoch_train_loss:.4f} | Val Loss: {epoch_val_loss:.4f}")

# 4. Final Testing Phase
model.eval()

# CRITICAL FIX: All 3 tracking containers must be initialized before the loop!
total_absolute_error = 0.0  
all_preds = []
all_targets = []

with torch.no_grad():
    for batch_X, batch_y in test_loader:
        predictions = model(batch_X)
        
        # 1. Accumulate absolute errors for MAE
        absolute_errors = torch.abs(predictions - batch_y)
        total_absolute_error += absolute_errors.sum().item()
        
        # 2. Append the batch values to our lists for R²
        all_preds.extend(predictions.cpu().numpy())
        all_targets.extend(batch_y.cpu().numpy())

average_mae = total_absolute_error / len(test_loader.dataset)
test_r2 = r2_score(all_targets, all_preds)

# Print your complete performance report
print(f"\n================ FINAL TEST PERFORMANCE ================")
print(f"Mean Absolute Error (MAE): {average_mae:.4f}")
print(f"R² Score (Variance Explained): {test_r2:.4f}")
print(f"========================================================")