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

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', 'liq','qc_liq']]
ds_cbh = xr.open_dataset('cbh_ceil.nc')

ds_mfrs = xr.open_mfdataset('sgpmfrsr7nchC1.b1/*')
ds_mfrs = ds_mfrs[['computed_cosine_correction_filter6', 'direct_normal_narrowband_filter6',
                'qc_direct_normal_narrowband_filter6']]
ds_mfrs

ds_mfrs = ds_mfrs.compute()

# 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.direct_normal_narrowband_filter6 > 0.15, drop=True)
ds_radio = ds_radio.where(ds_radio.qc_vap == 0, drop=True)
ds_cbh = ds_cbh.where(ds_cbh.qc_first_cbh == 0, drop=True)
ds_cbh = ds_cbh.where(ds_cbh.first_cbh >= 5000, drop=True)

# 1. Round or floor the time coordinates to the 1-minute mark ('min')
ds_radio_rounded = ds_radio.assign_coords(time=ds_radio.time.dt.floor("min"))
ds_mfrs_rounded = ds_mfrs.assign_coords(time=ds_mfrs.time.dt.floor("min"))
ds_cbh_rounded = ds_cbh.assign_coords(time=ds_cbh.time.dt.floor("min"))

# 2. If flooring introduces duplicate timestamps within the same dataset, clean them up:
ds_radio_rounded = ds_radio_rounded.drop_duplicates("time")
ds_mfrs_rounded = ds_mfrs_rounded.drop_duplicates("time")
ds_cbh_rounded = ds_cbh_rounded.drop_duplicates("time")

# 3. Now your inner merge works smoothly
ds_merged = xr.merge(
    [ds_radio_rounded, ds_mfrs_rounded, ds_cbh_rounded], join="inner"
)
ds_merged

plt.scatter(ds_merged.direct_normal_narrowband_filter6.values*2, ds_merged.vap.values/ds_merged.computed_cosine_correction_filter6.values*0.067, alpha=0.4)
plt.title("Direct Normal vs VAP")
plt.xlabel("Filter 6 Values")
plt.ylabel("VAP")
plt.grid(True)

np.max(ds_merged.vap.values/ds_merged.computed_cosine_correction_filter6.values)

# split the datasets into training, validation, and testing
x_data = ds_merged[['direct_normal_narrowband_filter6', 'computed_cosine_correction_filter6']].to_array().values.T
x_data = x_data
y_data = ds_merged.vap.values # / ds_merged.computed_cosine_correction_filter6.values *0.067

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)

import torch.nn as nn
import torch.optim as optim

# =====================================================================
# SIMPLE 3-HIDDEN-LAYER NEURAL NETWORK
# =====================================================================
class SimpleRegressor(nn.Module):
    def __init__(self, input_dim):
        super(SimpleRegressor, self).__init__()
        
        # Pure linear connections without any extra regularization or normalization
        self.layer1 = nn.Linear(input_dim, 64)  # Input -> Hidden 1
        self.layer2 = nn.Linear(64, 32)         # Hidden 1 -> Hidden 2
        self.layer3 = nn.Linear(32, 16)         # Hidden 2 -> Hidden 3
        self.output = nn.Linear(16, 1)          # Hidden 3 -> 1 Final Prediction
        
        self.relu = nn.ReLU() # Standard activation function
        
    def forward(self, x):
        # Linearly combine, activate, repeat
        x = self.relu(self.layer1(x))
        x = self.relu(self.layer2(x))
        x = self.relu(self.layer3(x))
        return self.output(x)

# 1. Initialize the simple network
input_features = X_train_tensor.shape[1] # Automatically reads your 3 input variables
model = SimpleRegressor(input_dim=input_features)

# 2. Standard Loss Function 
criterion = nn.MSELoss() 

# 3. Clean Adam Optimizer (no extra weight decay)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# --- EARLY STOPPING CONFIGURATION ---
NUM_EPOCHS = 100          # Raised this so early stopping actually has room to work!
patience = 10             # Number of epochs to wait for improvement before quitting
patience_counter = 0      # Tracks how many bad epochs in a row we've had
best_val_loss = float('inf') # Initialize at infinity so any real loss beats it
best_model_wts = copy.deepcopy(model.state_dict()) # To store our best weights

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)           
        loss = criterion(outputs, batch_y)  
        loss.backward()
        optimizer.step()
        
        running_train_loss += loss.item() * batch_X.size(0)
        
    epoch_train_loss = running_train_loss / len(train_loader.dataset)
    
    # --- VALIDATION PHASE ---
    model.eval()  
    running_val_loss = 0.0
    
    with torch.no_grad():  
        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:02d}/{NUM_EPOCHS} | Train Loss: {epoch_train_loss:.4f} | Val Loss: {epoch_val_loss:.4f}")
    
    # --- EARLY STOPPING CHECK ---
    if epoch_val_loss < best_val_loss:
        # We found a better model! Reset counter and save the weights
        best_val_loss = epoch_val_loss
        best_model_wts = copy.deepcopy(model.state_dict())
        patience_counter = 0
        print(f"--> New best model found! Saving weights...")
    else:
        # No improvement, increase the penalty counter
        patience_counter += 1
        print(f"--> No improvement. Patience: {patience_counter}/{patience}")
        
        if patience_counter >= patience:
            print(f"\n🛑 Early stopping triggered! Stopping training at epoch {epoch+1}.")
            break

# =====================================================================
# RESTORE BEST MODEL WEIGHTS
# =====================================================================
print("Loading best model weights for final evaluation...")
model.load_state_dict(best_model_wts)

# =====================================================================
# 4. FINAL TESTING PHASE
# =====================================================================
model.eval()

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"========================================================")

plt.scatter(all_targets, all_preds)