# Install required libraries via %pip as recommended by I-GUIDE
import sys
# Only run pip install if the libraries aren't already found
if 'pytorch_tabnet' not in sys.modules:
    # %pip install pytorch-tabnet==4.1.0 geopandas matplotlib seaborn scikit-learn
    pass

# Standard Data Science & Geospatial Libraries
import os
import glob
import itertools
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
from pathlib import Path
from scipy.ndimage import label
from matplotlib import gridspec

# Machine Learning & TabNet Dependencies
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from pytorch_tabnet.tab_model import TabNetClassifier
from pytorch_tabnet.tab_network import TabNet
from pytorch_tabnet.tab_network import TabNetNoEmbeddings
# Evaluation Metrics
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score, 
    roc_auc_score, classification_report, roc_curve, 
    confusion_matrix, ConfusionMatrixDisplay
)

# Hardware Verification
device = torch.device("cpu")
print(f"Environment ready. I-GUIDE Platform Evaluation Mode: {device}")

# Seed Initialization
import random

def set_seed(seed=42):
    """
    Sets the seed for all relevant libraries to ensure consistent and 
    reproducible results across the I-GUIDE environment.
    """
    # PyTorch internal reproducibility
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    
    # Standard Python and Numpy seeds
    np.random.seed(seed)
    random.seed(seed)
    
    # Force Deterministic behavior in CuDNN (relevant if GPU is available)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# Initialize the global seed (Default: 42)
set_seed(42)
print("Global seed set to 42. Reproducibility initialized.")

# 2.1 Define data folder using relative paths for platform compatibility
data_folder = Path('./Data')

# Check if the data directory exists to prevent execution errors on the platform
if not data_folder.exists():
    print(f"Error: {data_folder} directory not found. Please ensure it is uploaded to your GitHub repository.")
else:
    # Load the consolidated AI-ready dataset
    file_path = data_folder / 'Master_ML_Ready_Data.csv'
    df = pd.read_csv(file_path)

    # Display dataset dimensions and initial rows for verification
    print(f"Dataset loaded successfully. Shape: {df.shape}")
    display(df.head())

# 2.2 Temporal Preprocessing & Feature Engineering
# To ensure "AI-ready Data," we transform raw timestamps into cyclic features 

# Convert time column to datetime for extraction
df['timestamp'] = pd.to_datetime(df['timestamp'])

# Cyclic Time Transformation (Seasonality)
df['doy'] = df['timestamp'].dt.dayofyear
df['sin_doy'] = np.sin(2 * np.pi * df['doy'] / 365.25)
df['cos_doy'] = np.cos(2 * np.pi * df['doy'] / 365.25)

# Data Cleaning: Handle missing values from temporal shifts (e.g., SM_yesterday)
# We drop first row with NaN in key surface soil moisture columns to ensure model stability (there is no yesterday for first observation)
df_2 = df.dropna(subset=['SM_surface_AM']).copy()

print(f"Preprocessing complete. Cleaned Dataset Shape: {df_2.shape}")
display(df_2.head())

# 2.2. Physics-Consistent Data Filtering & Tensor Conversion
# To ensure "Responsible & Open Science," we filter for physically consistent hydrological events.

# Filter for precipitation events (Rainy days only) to adhere to SM2RAIN calibration assumptions
df_filtered = df_2[df_2['Pr'] > 1.0].copy()

# Remove inconsistent points where significant rain does not result in expected soil moisture increase
# This prevents the Geo-AI model from learning physically impossible noise
df_clean = df_filtered[~((df_filtered['Pr'] > 1.0) & (df_filtered['dsm_am'] < 0.001))]

print(f"Original records: {len(df_2)}")
print(f"Cleaned records: {len(df_clean)}")
print(f"Removed {len(df_filtered) - len(df_clean)} inconsistent hydrological points.")

# 2. Extract AlphaEarth Features and Physical Constraints
# Features (X): 64-dimensional satellite embeddings (regex '^A')
# Physics Variables: Core inputs for the Inverse-SM2RAIN loss function
X_data = df_clean.filter(regex='^A').values 
physics_data = df_clean[['dsm_am', 'Pr', 'SM_surface_AM', 'ET_Open']].values

# 3. Convert to Tensors for Platform Compatibility
# Note: While original training utilized GPU (Quadro RTX 6000), 
# this execution uses 'device' (CPU) for I-GUIDE Platform evaluation. [cite: 48]
X_tensor = torch.tensor(X_data, dtype=torch.float32).to(device)
P_tensor = torch.tensor(physics_data, dtype=torch.float32).to(device)

print(f"Features shape: {X_tensor.shape}")
print(f"Tensors successfully prepared on: {X_tensor.device}")

# 3.2. ILLUSTRATIVE TRAINING LOOP (CPU Optimized)
optimizer = optim.Adam(model.parameters(), lr=0.005, weight_decay=0)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=400, gamma=0.5)

print(f"Starting Illustrative Training Mode (30 Epochs)...")

for epoch in range(30):
    model.train()
    optimizer.zero_grad()

    try:
        # Forward pass through the Physics-Informed architecture
        Z, K, b, E, M_loss = model(X_tensor)
        
        # Unpack physical constraints from P_tensor
        dSM, pr, SM, ET = P_tensor[:, 0], P_tensor[:, 1], P_tensor[:, 2], P_tensor[:, 3]
        
        # Inverse-SM2RAIN Calculation
        SM_stable = torch.clamp(SM, min=1e-6, max=1.2)
        drainage = K * torch.exp(b * torch.log(SM_stable))
        
        # Physics Equation: IRR = (Z * dSM) - Pr + Drainage + ET
        irr_pred = (Z * dSM) - pr + drainage + ET
      
        # Core Loss: Mean Absolute Error (MAE)
        mae_loss = torch.mean(torch.abs(irr_pred)) 
        
        # Physical Penalty Calculations
        # 1. Negative Penalty: Irrigation cannot be less than zero
        neg_penalty = torch.mean(torch.relu(-irr_pred)) * 10.0
        
        # 2. ET Penalty: Discourages irrigation exceeding 1.2x actual Evapotranspiration
        et_threshold = 1.2 * ET
        et_penalty = torch.mean(torch.relu(irr_pred - et_threshold)) * 1.0
        
        # 3. Cap Penalty: Constrains output to regional physical limits (~10mm/day)
        max_physical_limit = 10.0
        cap_penalty = torch.mean(torch.relu(irr_pred - max_physical_limit)) * 1.0
        
        # 4. Smoothness Penalty: Ensures temporal consistency between samples
        if irr_pred.shape[0] > 1:
            diff = torch.abs(irr_pred[1:] - irr_pred[:-1])
            smooth_penalty = torch.mean(diff) * 1.0 
        else:
            smooth_penalty = 0.0

        # Total Physics-Informed Loss
        total_loss = mae_loss + neg_penalty + et_penalty + cap_penalty
        
        total_loss.backward()
        optimizer.step()
        scheduler.step()

    except RuntimeError as err:
        print(f"Numerical Error encountered: {err}")
        break
    
    # Progress Logging
    if epoch % 10 == 0:
        curr_lr = optimizer.param_groups[0]['lr']
        print(f"Ep {epoch:4d} | MAE: {mae_loss.item():.2f} | Neg-Pen: {neg_penalty.item():.2f} | LR: {curr_lr:.6f}")
        print(f"   AVG Params -> Z: {Z.mean().item():.1f} | K: {K.mean().item():.3f} | b: {b.mean().item():.2f}")
        print("-" * 50)

print("--- Illustrative training phase complete. For final results, load the pre-trained weights. ---")

# =================================================================
# 4.1. SAVING THE MODEL (COMMENTED FOR 30 epoch run)
# =================================================================
# This section is preserved for reference. In the full research 
# pipeline, this creates a comprehensive checkpoint.

# checkpoint = {
#      'epoch': epoch,
#      'model_state_dict': model.state_dict(),
#      'optimizer_state_dict': optimizer.state_dict(),
#      'scheduler_state_dict': scheduler.state_dict(),
#      'loss': mae_loss.item(),
#      'input_dim': X_tensor.shape[1] 
# }

# # Save to working directory for future inference or deployment
# torch.save(checkpoint, 'physics_tabnet_best_Final.pt')
# print("TabNet model and training state saved successfully.")

# Verification of current parameter ranges before full restoration
model.eval()
with torch.no_grad():
    Z, K, b, E, _ = model(X_tensor.to(device))
    print(f"Current Z range (Pre-load): {Z.min().item():.3f} to {Z.max().item():.3f}")
    print(f"Current K range (Pre-load): {K.min().item():.3f} to {K.max().item():.3f}")
    print(f"Current b range (Pre-load): {b.min().item():.3f} to {b.max().item():.3f}")
    print(f"Current E range (Pre-load): {E.min().item():.3f} to {E.max().item():.3f}")

# =================================================================
# 4.2. LOADING THE PRE-TRAINED MODEL
# =================================================================

# 1. Re-initialize the TabNet architecture
# input_dim must match the 64-dimensional AlphaEarth embeddings
model = PhysicsTabNet(input_dim=64).to(device)

# 2. Load the state dictionary from the repository
# We use map_location=device to ensure compatibility with the Platform's CPU
checkpoint_path = 'physics_tabnet_best_Final.pt'

if not os.path.exists(checkpoint_path):
    print(f"Error: {checkpoint_path} not found. Ensure it is in your GitHub repo root.")
else:
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])

    # 3. Set to evaluation mode for final inference
    model.eval()
    print(f"TabNet Model loaded successfully.")

# 5.1. Prepare Features & Execute Forward Pass
master_df = df.copy(deep=True) 

# Ensure data consistency: remove physically impossible records
master_df = master_df[~((master_df['Pr'] > 1.0) & (master_df['dsm_am'] < 0.001))]

# Convert AlphaEarth embeddings to tensor for TabNet processing
X_full = torch.tensor(master_df.filter(regex='^A').values, dtype=torch.float32)

# Ensure model is on CPU for I-GUIDE Platform execution
model.to('cpu')
model.eval()

with torch.no_grad():
    # Model generates the physical parameters for the SM2RAIN equation
    Z_pred, K_pred, b_pred, E_pred, _ = model(X_full)

# Record predicted physical parameters in the main dataframe
master_df['Z_pinn'] = Z_pred.numpy()
master_df['K_pinn'] = K_pred.numpy()
master_df['b_pinn'] = b_pred.numpy()
master_df['E_residual'] = E_pred.numpy()

# 5.2. Calculate Volumetric Irrigation (Inverse-SM2RAIN Physics)
# Convert forcing variables to tensors for vectorized math
dSM_t = torch.tensor(master_df['dsm_am'].values, dtype=torch.float32)
Pr_t  = torch.tensor(master_df['Pr'].values, dtype=torch.float32)
SM_t  = torch.tensor(master_df['SM_surface_AM'].values, dtype=torch.float32)
ET_t  = torch.tensor(master_df['ET_Open'].values, dtype=torch.float32)

# Core Equation: IRR = (Z * dSM) - Pr + (K * SM^b) + ET
# We incorporate a stability constant (1e-6) for the power function
irr_tensor = (Z_pred * dSM_t) - Pr_t + (K_pred * torch.pow(SM_t + 1e-6, b_pred)) + ET_t - 1

# 5.3. Post-Processing & Physical Guardrails
master_df['irr_raw_pinn'] = irr_tensor.numpy()

# Apply physical constraints: 
# 1. Non-negativity
master_df['irr_raw_pinn'] = master_df['irr_raw_pinn'].clip(0)

# 2. Precipitation masking: set to 0 if it is raining
master_df.loc[master_df['Pr'] > 0, 'irr_raw_pinn'] = 0

# 3. Upper-bound regional cap (10mm/day for Central Valley)
master_df['irr_raw_pinn'] = master_df['irr_raw_pinn'].clip(upper=10)

print(f"Inference complete. Total predictions generated: {len(master_df)}")

# 5.4. Distribution & Single-Pixel Visualization
import matplotlib.pyplot as plt
import random

# Create a figure with two subplots: Histogram and Time Series
fig = plt.figure(figsize=(15, 6))
gs = fig.add_gridspec(1, 2)

# Plot A: Irrigation Distribution
ax1 = fig.add_subplot(gs[0, 0])
ax1.hist(master_df['irr_raw_pinn'], bins=50, color='teal', edgecolor='black', alpha=0.7)
ax1.set_title('Global Irrigation Intensity Distribution', fontsize=12)
ax1.set_xlabel('Irrigation (mm/day)')
ax1.set_ylabel('Frequency')
ax1.grid(axis='y', linestyle='--', alpha=0.6)

# Plot B: Random Single-Pixel Time Series
ax2 = fig.add_subplot(gs[0, 1])
unique_pixels = master_df['pixel_id'].unique()
random_pixel = random.choice(unique_pixels)
pixel_data = master_df[master_df['pixel_id'] == random_pixel].copy()
pixel_data['timestamp'] = pd.to_datetime(pixel_data['timestamp'])
pixel_data = pixel_data.sort_values('timestamp')

ax2.plot(pixel_data['timestamp'], pixel_data['irr_raw_pinn'], 
         color='crimson', linewidth=1.5, marker='o', markersize=3, 
         label=f'Pixel {random_pixel}')
ax2.set_title(f'Temporal Signal: Pixel {random_pixel}', fontsize=12)
ax2.set_xlabel('Date')
ax2.set_ylabel('Irrigation (mm/day)')
plt.xticks(rotation=45)
ax2.grid(True, linestyle='--', alpha=0.6)
ax2.legend()

plt.tight_layout()
plt.show()

import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import r2_score

# 6.1. Aggregate Data: HUC12 Monthly Summation
# We group by HUC12 and timestamp to calculate daily means, then move to monthly totals
huc_daily_pinn = master_df.groupby(['huc12', 'timestamp', 'areaacres']).agg({
    'irr_raw_pinn': 'mean',
    'irrwdtot_mgd': 'first',
    'Pr': 'mean'           # Mean daily precipitation depth (mm)
}).reset_index()

huc_daily_pinn['year_month'] = huc_daily_pinn['timestamp'].dt.to_period('M')

huc_monthly_pinn = huc_daily_pinn.groupby(['huc12', 'year_month', 'areaacres']).agg({
    'irr_raw_pinn': 'sum',    # Monthly Total Irrigation Depth (mm)
    'irrwdtot_mgd': 'mean',   # Reported Monthly Daily Rate (MGD)
    'Pr': 'sum'               # Monthly Total Precipitation Depth (mm)
}).reset_index()

# 6.2. Unit Conversion: Depth (mm) to Volume (MGD)
# Logic: (Depth * 0.001 [m] * Area [acres] * 4046.86 [m2/acre] * 264.172 [gal/m3]) / 1e6 / days
huc_monthly_pinn['days_in_month'] = huc_monthly_pinn['year_month'].dt.days_in_month
huc_monthly_pinn['raw_calc_mgd'] = (
    (huc_monthly_pinn['irr_raw_pinn'] * 0.001 * huc_monthly_pinn['areaacres'] * 4046.86 * 264.172) 
    / 1e6 / huc_monthly_pinn['days_in_month']
)

# 6.3. Calculate Regional Time Series
regional_ts_pinn = huc_monthly_pinn.groupby('year_month').agg({
    'irrwdtot_mgd': 'sum',
    'raw_calc_mgd': 'sum',
    'Pr': 'mean' 
}).reset_index()
regional_ts_pinn['date'] = regional_ts_pinn['year_month'].dt.to_timestamp()

# 6.4. Statistical Performance Calculation
y_true = regional_ts_pinn['irrwdtot_mgd']
y_pred = regional_ts_pinn['raw_calc_mgd']
pearson_r2 = y_true.corr(y_pred)**2
nse_score = r2_score(y_true, y_pred)

# 6.5. High-Impact Hydrological Visualization
fig, ax1 = plt.subplots(figsize=(14, 8))

# Plot Irrigation Metrics
ax1.plot(regional_ts_pinn['date'], y_true, 
         label='USGS Reported (Reference)', color='blue', linewidth=2.5, marker='o', markersize=5)
ax1.plot(regional_ts_pinn['date'], y_pred, 
         label='Physics-Informed Geo-AI (Predicted)', color='red', linestyle='--', linewidth=2.5, marker='o', markersize=5)

ax1.set_xlabel('Timeline', fontsize=18)
ax1.set_ylabel('Total Irrigation (MGD)', fontsize=18)
ax1.set_ylim(0, 55000) 
ax1.tick_params(axis='both', labelsize=14)
ax1.grid(True, which='both', linestyle='--', alpha=0.3)

# Overlay Precipitation (Hyetograph)
ax2 = ax1.twinx()
ax2.bar(regional_ts_pinn['date'], regional_ts_pinn['Pr'], width=20, 
        color='skyblue', alpha=0.4, label='Precipitation (mm)')
ax2.set_ylabel('Monthly Precipitation (mm)', fontsize=18)
ax2.tick_params(axis='both', labelsize=14)

# Legend and Metrics Overlay
lines1, labels1 = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper right', fontsize=14, frameon=True)

stats_text = f'Trend (Pearson $r^2$): {pearson_r2:.3f}\nAccuracy (NSE): {nse_score:.3f}'
ax1.text(0.02, 0.96, stats_text, transform=ax1.transAxes, fontsize=16,
         verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

plt.title('Validation: Monthly Irrigation Volume (USGS vs. Physics-Informed TabNet)', fontsize=20)
plt.tight_layout()
plt.show()
print(f"Trend Correlation (Pearson r2): {pearson_r2:.4f}")
print(f"Regional NSE Score: {nse_score:.4f}")

from sklearn.metrics import r2_score, mean_squared_error
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# 7.1. Calculate Metrics per Individual HUC-12 Unit
huc_performance_list = []

# Group by HUC12 and calculate metrics for spatial validation
for huc, group in huc_monthly_pinn.groupby('huc12'):
    # Ensure clean calculation by dropping potential NaNs
    valid_data = group.dropna(subset=['irrwdtot_mgd', 'raw_calc_mgd'])
    
    # Requirement: At least 6 months of data for a meaningful statistical profile
    if len(valid_data) >= 6:
        y_true = valid_data['irrwdtot_mgd']
        y_pred = valid_data['raw_calc_mgd']
        
        # Hydrological Accuracy Metrics
        r2 = r2_score(y_true, y_pred)
        rmse = np.sqrt(mean_squared_error(y_true, y_pred))
        mae = np.abs(y_true - y_pred).mean()
        
        # Relative Error (Bias) calculation
        bias = (y_pred.sum() - y_true.sum()) / (y_true.sum() + 1e-6) * 100
        
        huc_performance_list.append({
            'huc12': huc,
            'data_points': len(valid_data),
            'r2': r2,
            'rmse_mgd': rmse,
            'mae_mgd': mae,
            'pct_bias': bias
        })

# Create the performance summary dataframe
df_huc_results = pd.DataFrame(huc_performance_list)
df_huc_results = df_huc_results.sort_values(by='r2', ascending=False)

print(f"Calculated localized metrics for {len(df_huc_results)} HUC-12 units.")

# 7.2. Visualize MAE Error Distribution
plt.figure(figsize=(10, 8))

# Histogram with Kernel Density Estimate (KDE)
ax = sns.histplot(df_huc_results['mae_mgd'], bins=30, kde=True, color='red', alpha=0.6)

# Title and Labeling for High-Impact Presentation
plt.title('MAE Distribution per HUC-12: Physics-Informed TabNet', fontsize=20)
plt.xlabel('Mean Absolute Error (MGD)', fontsize=18)
plt.ylabel('Count (Number of HUC-12s)', fontsize=18)

# Calculate and plot Median MAE for benchmark reference
median_mae = df_huc_results['mae_mgd'].median()
plt.axvline(median_mae, color='darkred', linestyle='--', 
            label=f'Median MAE: {median_mae:.2f} MGD')

# Formatting for clarity and readability
plt.legend(fontsize=16)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlim(0, 320)
plt.ylim(0, 65)
plt.grid(axis='y', linestyle='--', alpha=0.3)

plt.tight_layout()
plt.show()

# Display the top 10 most accurately modeled HUCs
print("\n--- Highest Reliability Units (Top 10 HUCs by R2) ---")
display(df_huc_results.head(10))

import pandas as pd
import geopandas as gpd
import json
from shapely.geometry import shape
import matplotlib.pyplot as plt

# 8.1. Spatial Data Preparation
# Convert timestamp to datetime and aggregate to annual mean irrigation per pixel
master_df['timestamp'] = pd.to_datetime(master_df['timestamp'])
annual_irr = master_df.groupby(['pixel_id', master_df['timestamp'].dt.year])['irr_raw_pinn'].mean().reset_index()
annual_irr.columns = ['pixel_id', 'year', 'annual_irr_mean']

# Load the grid geometry file (ensure Grid_updated.csv is in the Data folder)
grid = pd.read_csv(data_folder / 'Grid_updated.csv')

# Merge geometric data with our physics-informed predictions
Grid_merged = grid.merge(annual_irr, on='pixel_id', how='left')

# Convert GeoJSON string column into actual geometry objects
Grid_merged['geometry'] = Grid_merged['.geo'].apply(lambda x: shape(json.loads(x)))

# Initialize GeoDataFrame in the standard EPSG:4326 Coordinate Reference System
gdf = gpd.GeoDataFrame(Grid_merged, geometry='geometry', crs="EPSG:4326")

# 8.2. Multi-Year Comparative Mapping (2018 - 2020)
years_to_map = [2018, 2019, 2020]
fig, axes = plt.subplots(1, 3, figsize=(24, 10))

for i, year in enumerate(years_to_map):
    # Filter for the specific year
    gdf_year = gdf[gdf['year'] == year]
    
    # Plotting the annual irrigation intensity
    gdf_year.plot(column='annual_irr_mean', ax=axes[i], legend=True,
                  cmap='YlGnBu', edgecolor='none',
                  legend_kwds={'label': "Mean Irrigation (mm/day)", 'orientation': "horizontal", 'pad': 0.02})
    
    axes[i].set_title(f"Central Valley Irrigation: {year}", fontsize=20)
    axes[i].axis('off')

plt.suptitle("Multi-Year Geospatial Assessment of Predicted Irrigation Intensity", fontsize=26, y=0.95)
plt.tight_layout(rect=[0, 0.03, 1, 0.95])

# Save and show the high-resolution comparison
plt.savefig("central_valley_comparison_2018_2020.png", dpi=300, bbox_inches='tight', facecolor='white')
plt.show()

# 8.3. Interactive Map Generation (Example for 2018)
import folium
m = folium.Map(location=[37.5, -120], zoom_start=7, tiles='CartoDB positron')

gdf_2018 = gdf[gdf['year'] == 2018]
folium.Choropleth(
    geo_data=gdf_2018,
    name='Annual Irrigation 2018',
    data=gdf_2018,
    columns=['pixel_id', 'annual_irr_mean'],
    key_on='feature.properties.pixel_id',
    fill_color='YlGnBu',
    fill_opacity=0.7,
    line_opacity=0.2,
    legend_name='Mean Daily Irrigation (mm)'
).add_to(m)

print("Interactive map and annual comparison plots generated successfully.")

import numpy as np
import pandas as pd
from scipy.optimize import minimize

# 9.1. Traditional SM2RAIN Numerical Calibration (Baseline)
def objective_function(params, dSM, P, SM, ET):
    Z, K, b = params
    # Traditional SM2RAIN Physical Equation
    residual = (Z * dSM) - P + (K * (SM**b)) + ET
    return np.mean(np.abs(residual))

results_list = []
bounds = [(10, 3000), (0.01, 50), (1, 12)]
initial_guess = [500, 5, 5]

unique_pixels = df['pixel_id'].unique()
print(f"Starting traditional calibration for {len(unique_pixels)} pixels...")

for pid in unique_pixels:
    pixel_data = df[df['pixel_id'] == pid]
    
    # Run pixel-by-pixel optimization (Standard Baseline Approach)
    res = minimize(
        objective_function, 
        initial_guess, 
        args=(pixel_data['dsm_am'].values, pixel_data['Pr'].values, 
              pixel_data['SM_surface_AM'].values, pixel_data['ET_Open'].values),
        method='L-BFGS-B', 
        bounds=bounds
    )
    
    if res.success:
        Z_opt, K_opt, b_opt = res.x
        results_list.append({
            'pixel_id': pid,
            'Z_opt': Z_opt, 'K_opt': K_opt, 'b_opt': b_opt,
            'mae_error': res.fun
        })

calibration_df = pd.DataFrame(results_list)
print("Baseline Calibration Complete.")

# 9.2. Baseline Forward Pass & Unit Conversion
# Merge the calibrated parameters back to the main dataframe
df_final = pd.merge(df, calibration_df, on='pixel_id', how='left')

# Calculate irrigation using traditionally calibrated parameters (No Embeddings)
# Equation: Irr = (Z * dSM) - P + (K * SM^b) + ET + Mean_Error
df_final['calculated_irrigation'] = (
    (df_final['Z_opt'] * df_final['dsm_am']) - 
    df_final['Pr'] + 
    (df_final['K_opt'] * (df_final['SM_surface_AM']**df_final['b_opt'])) + 
    df_final['ET_Open'] + 
    df_final['mae_error']
).clip(lower=0)

# Physical Guardrail: Irrigation is zero on rainy days
df_final.loc[df_final['Pr'] > 0, 'calculated_irrigation'] = 0

# 9.3. Regional Baseline Validation vs. USGS
# Aggregate daily results to HUC12 level
huc_daily_num = df_final.groupby(['huc12', 'timestamp', 'areaacres']).agg({
    'calculated_irrigation': 'mean', 
    'irrwdtot_mgd': 'first', 
    'Pr': 'mean'
}).reset_index()

huc_daily_num['year_month'] = huc_daily_num['timestamp'].dt.to_period('M')

# Aggregate to monthly totals
huc_monthly_num = huc_daily_num.groupby(['huc12', 'year_month', 'areaacres']).agg({
    'calculated_irrigation': 'sum', 
    'irrwdtot_mgd': 'mean', 
    'Pr': 'sum'
}).reset_index()

# Convert Depth (mm) to Volume (MGD) for comparison
huc_monthly_num['days_in_month'] = huc_monthly_num['year_month'].dt.days_in_month
huc_monthly_num['calc_irr_mgd'] = (
    (huc_monthly_num['calculated_irrigation'] * 0.001 * huc_monthly_num['areaacres'] * 4046.86 * 264.172) 
    / 1e6 / huc_monthly_num['days_in_month']
)

# 9.4. Regional Performance Metrics & Visualization
regional_ts_num = huc_monthly_num.groupby('year_month').agg({
    'irrwdtot_mgd': 'sum', 
    'calc_irr_mgd': 'sum', 
    'Pr': 'mean'
}).reset_index()
regional_ts_num['date'] = regional_ts_num['year_month'].dt.to_timestamp()

y_true_base = regional_ts_num['irrwdtot_mgd']
y_pred_base = regional_ts_num['calc_irr_mgd']

# Calculate Metrics
pearson_r2_base = y_true_base.corr(y_pred_base)**2
nse_base = r2_score(y_true_base, y_pred_base)

# Plotting
plt.figure(figsize=(14, 8))
plt.plot(regional_ts_num['date'], y_true_base, label='USGS Reported (Reference)', color='blue', linewidth=2, marker='o')
plt.plot(regional_ts_num['date'], y_pred_base, label='Traditional SM2RAIN (Baseline)', color='green', linestyle='--', linewidth=2, marker='o')

# Formatting
plt.title(f'Baseline Validation: USGS vs. Traditional Model', fontsize=20)
plt.ylabel('Total Irrigation (MGD)', fontsize=16)
plt.xlabel('Date', fontsize=16)
plt.grid(True, alpha=0.3)
plt.legend(fontsize=14)

# Display stats on plot
stats_text = f'Baseline $R^2$: {pearson_r2_base:.3f}\nBaseline NSE: {nse_base:.3f}'
plt.text(0.02, 0.95, stats_text, transform=plt.gca().transAxes, fontsize=14, 
         verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

plt.show()

print(f"--- Baseline Performance Results ---")
print(f"Trend Correlation (Pearson R2): {pearson_r2_base:.4f}")
print(f"Absolute Accuracy (NSE): {nse_base:.4f}")

# 9.4. Error Distribution: Physical Model vs. Physics-Informed Geo-AI
huc_performance_baseline = []
for huc, group in huc_monthly_num.groupby('huc12'):
    valid = group.dropna(subset=['irrwdtot_mgd'])
    if len(valid) >= 6:
        mae = np.abs(valid['irrwdtot_mgd'] - valid['calc_irr_mgd']).mean()
        huc_performance_baseline.append({'huc12': huc, 'mae_mgd': mae})

df_huc_baseline = pd.DataFrame(huc_performance_baseline)

plt.figure(figsize=(10, 7))
sns.histplot(df_huc_baseline['mae_mgd'], bins=30, kde=True, color='green', label='Traditional Baseline')
sns.histplot(df_huc_results['mae_mgd'], bins=30, kde=True, color='red', label='Physics-Informed TabNet')

median_base = df_huc_baseline['mae_mgd'].median()
median_pinn = df_huc_results['mae_mgd'].median()

plt.axvline(median_base, color='darkgreen', linestyle='--', label=f'Baseline Median: {median_base:.2f}')
plt.axvline(median_pinn, color='darkred', linestyle='--', label=f'PINN Median: {median_pinn:.2f}')

plt.title('Error Reduction Analysis: PINN vs. Traditional Baseline', fontsize=20)
plt.xlabel('Mean Absolute Error (MGD)', fontsize=16)
plt.legend(fontsize=14)
plt.xlim(0, 320)
plt.show()

improvement = ((median_base - median_pinn) / median_base) * 100
print(f"Our Physics-Informed Geo-AI improved the median MAE by {improvement:.2f}% over the traditional baseline.")