import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import xgboost as xgb  # Adding XGBoost
from sklearn.linear_model import LinearRegression
from sklearn.multioutput import MultiOutputRegressor   
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, accuracy_score, f1_score, precision_score, recall_score, classification_report
import logging
import os
import multiprocessing
from datetime import datetime
import argparse
import joblib
import platform
import json
from sklearn.impute import SimpleImputer, KNNImputer
import zarr
import time
import psutil  # For hardware details
from sklearn.preprocessing import MinMaxScaler, StandardScaler
import gc
import sys
import torch
from utils import (
    load_data_for_years, inverse_scale_data,
    sanitize_params, convert_numpy_types, evaluate,
    plot_postprocessing_outputs, plot_training_curves,
    log_system_info, log_tensor_stats, log_stats
)
from modeling import BackwardLAIPredictor, LAIDataset
from preprocessing import preprocess_lai_data, prepare_dataloaders, load_combined_water_mask    
from postprocessing import apply_water_mask

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a regression model to predict Leaf Area Index (LAI).")
parser.add_argument("--start_year", type=int, required=True, help="Start year for the data range")
parser.add_argument("--end_year", type=int, required=True, help="End year for the data range")
parser.add_argument("--region", type=str, required=True, help="Region for the data slice (ES, EU, UK, AU, GLOBE)")
parser.add_argument("--imputation", type=str, choices=['mean', 'median', 'most_frequent', 'knn'], default='mean', help="Imputation method for handling NaN values")
parser.add_argument("--train_years", nargs='+', type=int, required=True, help="Years to use for training")
parser.add_argument("--test_years", nargs='+', type=int, required=True, help="Years to use for testing")
parser.add_argument("--save_data_path", type=str, required=True, help="Path to save data and model outputs")
parser.add_argument("--autoregressive", action='store_true', help="Include autoregressive feature from ground truth data (year +1)")
parser.add_argument("--autoregressive_smoothing", action='store_true', help="Smooth the autoregression values for reducing the importance of the LAI LKS", default=False)
parser.add_argument("--feature_scale", action='store_true', help="scale the features to the max value", default=False)
parser.add_argument("--ground_truth_scale", action='store_true', help="scale the ground truth to the max value", default=False)
parser.add_argument("--ground_truth_log_transform", action='store_true', help="Apply log transform to ground truth", default=False)
parser.add_argument("--water_filter", action='store_true', help="Filter out water data from training with the water masks")
parser.add_argument("--model_timestamp", type=str, required=True, help="Model timestamp from training phase")
parser.add_argument("--batch_size", type=int, default=1024, help="Batch size for training")
parser.add_argument("--post_plots", action="store_true", default=False, help="Enable post training plots")

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
region = args.region
model_type = "rnn"
task_type = "regression"
imputation_method = args.imputation
train_years = args.train_years
test_years = args.test_years
save_data_path = args.save_data_path
autoregressive = args.autoregressive
autoregressive_smoothing = args.autoregressive_smoothing
timestamp = args.model_timestamp
feature_scale = args.feature_scale
ground_truth_scale = args.ground_truth_scale
ground_truth_log_transform = args.ground_truth_log_transform
water_filter = args.water_filter
batch_size = args.batch_size
post_plots = args.post_plots

# Define  default scaling range
lower_bound = 0
upper_bound = 7

# Define smoothing parameters
sigma = 10
mode='nearest'

# Define model name
if task_type == "regression":
    model_name = "XGBoost"


# Configure logging
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)
log_file_name = f"train_log_{start_year}_{end_year}_{region}_{timestamp}.log"
log_file_path = os.path.join(logs_path, log_file_name)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[logging.FileHandler(log_file_path),logging.StreamHandler()])

logging.info(f"Training model: {model_name} with MultiOutputRegressor wrapper.")

logging.info(f"Timestamp: {timestamp}")

# Log system information
log_system_info()

# Set paths for saving/loading processed data and metadata
metadata_path = os.path.join(save_data_path, f'metadata_training_{start_year}_{end_year}_{region}.json')

# Load metadata
with open(metadata_path, 'r') as f:
    metadata = json.load(f)

# Log the loaded metadata
logging.info(f"Loaded metadata: {metadata}")

# Extract additional configurations
lat_dim = metadata.get('lat_dim')
lon_dim = metadata.get('lon_dim')
   
# Use the metadata's data path
data_path = metadata.get('data_path', save_data_path)
logging.info(f"Using data path from metadata: {data_path}")
logging.info(f"Expected dimensions (lat x lon): {lat_dim} x {lon_dim}")

# ---- PREPROCESSING ----

# Start tracking the preprocessing time
logging.info(f"Starting preprocessing step")
prep_start_time = time.time()

# Load data for training and testing years (now as lists)
X_train_list, y_train_list = load_data_for_years(train_years, region, save_data_path)
X_test_list, y_test_list = load_data_for_years(test_years, region, save_data_path)

# Sanity check on loaded lists
num_train_years = len(train_years)
num_test_years = len(test_years)
assert len(X_train_list) == num_train_years, "Mismatch in training years and data length"
assert len(X_test_list) == num_test_years, "Mismatch in testing years and data length"

# Log original shape info
logging.info(f"Loaded data for {num_train_years} training years and {num_test_years} testing years.")
logging.info(f"Original X_train shapes: {[x.shape for x in X_train_list]}")
logging.info(f"Original y_train shapes: {[y.shape for y in y_train_list]}")

if water_filter:
    # Load water mask and apply to training features/targets
    water_mask = load_combined_water_mask(train_years, region, save_data_path, logger=logging)

    # Filter only land points from each year
    X_train_list = [X[water_mask] for X in X_train_list]
    y_train_list = [y[water_mask] for y in y_train_list]

    logging.info(f"Applied water mask. Kept {np.sum(water_mask):,} land pixels out of {len(water_mask):,} total.")
    logging.info(f"Filtered X_train shapes: {[x.shape for x in X_train_list]}")
    logging.info(f"Filtered y_train shapes: {[y.shape for y in y_train_list]}")
    
# Validate enough training years
if num_train_years < 2:
    raise ValueError("Need at least 2 years in training data to create a validation split.")

# Data splitting
val_years_count = max(1, int(0.2 * num_train_years))
train_years_count = num_train_years - val_years_count

X_train_array = np.stack(X_train_list, axis=1)  # Shape: (num_locations, num_years, features)
y_train_array = np.stack(y_train_list, axis=1)

X_test_array = np.stack(X_test_list, axis=1)
y_test_array = np.stack(y_test_list, axis=1)

X_train_monthly = X_train_array[:, val_years_count:, :12] # First 12 variables (next year's LAI)
X_train_yearly = X_train_array[:, val_years_count:, 12:] # Last 15 variables (yearly static variables)
y_train = y_train_array[:, val_years_count:, :]

X_val_monthly = X_train_array[:, :val_years_count, :12]
X_val_yearly = X_train_array[:, :val_years_count, 12:]
y_val = y_train_array[:, :val_years_count, :]

X_test_monthly = X_test_array[:, :, :12]
X_test_yearly = X_test_array[:, :, 12:]
y_test = y_test_array

# Log split and variables info
num_locations = X_train_yearly.shape[0]
num_train_years = X_train_yearly.shape[1]
num_yearly_variables = X_train_yearly.shape[2]
num_monthly_variables = X_train_monthly.shape[2]
logging.info(f"Training years: {train_years_count}")
logging.info(f"Validation years: {val_years_count}")

# Compute and log stats BEFORE preprocessing
log_tensor_stats(torch.tensor(y_train), torch.tensor(y_val), torch.tensor(y_test), torch.tensor(X_train_monthly), torch.tensor(X_val_monthly),
    torch.tensor(X_test_monthly), message_prefix="LAI stats BEFORE preprocessing:"
)
logging.info(f"Data summary BEFORE preprocessing - Locations: {num_locations}, "f"Yearly Variables: {num_yearly_variables}, " f"Monthly Variables: {num_monthly_variables}")

# Start preprocessing
logging.info(f"Applying preprocessing pipeline: scaling ([0,1] range)={ground_truth_scale}, log_transform={ground_truth_log_transform}, imputation={imputation_method}")
logging.info(f"Preprocessing includes one-hot encoding for HILDA+ variable (yearly variable 15), expected to get 14 + 13(HILDA+ classes) = 27 yearly variables after")

# Preprocess data
X_train_monthly, X_train_yearly, y_train, \
X_val_monthly, X_val_yearly, y_val, \
X_test_monthly, X_test_yearly, y_test, \
y_train_post, y_val_post, y_test_post, \
scaling_min, scaling_max = preprocess_lai_data(X_train_monthly, X_train_yearly, y_train, X_val_monthly, X_val_yearly, y_val,
    X_test_monthly, X_test_yearly, y_test, scale_gt=ground_truth_scale, log_transform_gt=ground_truth_log_transform,
    lower_bound=lower_bound, upper_bound=upper_bound, plot=True, save_data_path=save_data_path, region=region, timestamp=timestamp
)

# Check for NaNs after preprocessing
for name, tensor in {
    "X_train_monthly": X_train_monthly,
    "X_train_yearly": X_train_yearly,
    "y_train": y_train,
    "X_val_monthly": X_val_monthly,
    "X_val_yearly": X_val_yearly,
    "y_val": y_val,
    "X_test_monthly": X_test_monthly,
    "X_test_yearly": X_test_yearly,
    "y_test": y_test,
}.items():
    logging.info(f"{name} contains NaNs: {torch.isnan(tensor).any().item()}")

# Log stats AFTER preprocessing
num_yearly_variables = X_train_yearly.shape[-1] # Update num_yearly_vars after one-hot encoding
num_monthly_variables = X_train_monthly.shape[-1]
log_tensor_stats(y_train, y_val, y_test, X_train_monthly, X_val_monthly, X_test_monthly, message_prefix="LAI stats AFTER preprocessing:")
logging.info(f"Data summary AFTER preprocessing - Locations: {num_locations}, "f"Yearly Variables: {num_yearly_variables}, " f"Monthly Variables: {num_monthly_variables}")

# Convert torch tensors to NumPy arrays (required for XGBoost)
X_train_monthly = X_train_monthly.numpy()
X_train_yearly = X_train_yearly.numpy()
y_train = y_train.numpy()
X_val_monthly = X_val_monthly.numpy()
X_val_yearly = X_val_yearly.numpy()
y_val = y_val.numpy()
X_test_monthly = X_test_monthly.numpy()
X_test_yearly = X_test_yearly.numpy()
y_test = y_test.numpy()

# Flatten monthly and yearly data and combine
X_train_flat = X_train_monthly.reshape(-1, X_train_monthly.shape[-1])
X_train_yearly_flat = X_train_yearly.reshape(-1, X_train_yearly.shape[-1])
X_train_combined = np.concatenate([X_train_flat, X_train_yearly_flat], axis=1)

y_train_flat = y_train.reshape(-1, 12)

X_val_flat = X_val_monthly.reshape(-1, X_val_monthly.shape[-1])
X_val_yearly_flat = X_val_yearly.reshape(-1, X_val_yearly.shape[-1])
X_val_combined = np.concatenate([X_val_flat, X_val_yearly_flat], axis=1)

y_val_flat = y_val.reshape(-1, 12)

X_test_flat = X_test_monthly.reshape(-1, X_test_monthly.shape[-1])
X_test_yearly_flat = X_test_yearly.reshape(-1, X_test_yearly.shape[-1])
X_test_combined = np.concatenate([X_test_flat, X_test_yearly_flat], axis=1)

y_test_flat = y_test.reshape(-1, 12)


# Preprocessing time
preprocessing_duration = time.time() - prep_start_time
logging.info(f"Preprocessing completed in {preprocessing_duration:.2f} seconds.")

# ---- TRAINING ---- 

# Train XGBoost
logging.info(f"Starting XGBoost training...")
start_time = time.time()

base_model = xgb.XGBRegressor(
    n_estimators=700,
    max_depth=12,
    learning_rate=0.01,
    subsample=0.8,
    reg_alpha=0.5,
    reg_lambda=1.0,
    colsample_bytree=0.8,
    tree_method='hist',
    random_state=42,
    device='cuda',
    verbosity=2,
    eval_metric='rmse',
    #early_stopping_rounds=50 # Disable early stopping since we cannot pass a validation set with MultiOutputRegressor
)

model = MultiOutputRegressor(base_model)

model.fit(X_train_combined, y_train_flat)

training_duration = time.time() - start_time
logging.info(f"Training completed in {training_duration:.2f} seconds.")

# Evaluation
logging.info(f"Starting Evaluation step on test set")
logging.info(f"Model architecture:\n{model}")

# Predict
logging.info(f"Starting prediction...")
y_pred = model.predict(X_test_combined)
logging.info(f"Prediction done. y_pred shape: {y_pred.shape}")

# Flatten y_test_post to match y_pred
y_train = y_train_post.reshape(-1, 12)  # shape: (num_locations * num_train_years, 12)
y_test = y_test_post.reshape(-1, 12)  # shape: (num_locations * num_test_years, 12)

# Log shapes of y_pred and y_test_post
logging.info(f"Predictions shape: {y_pred.shape}")
logging.info(f"Ground Truth shape: {y_test.shape}")

assert y_pred.shape == y_test.shape, "Shape mismatch between predictions and ground truth"

# Postprocessing
if ground_truth_scale:
    logging.info("Applying inverse scaling...")
    y_pred = inverse_scale_data(y_pred, scaling_min, scaling_max)
    y_test = inverse_scale_data(y_test, scaling_min, scaling_max)

    # Logging stats
    logging.info(f"After inverse scaling — y_pred min: {np.nanmin(y_pred):.4f}, max: {np.nanmax(y_pred):.4f}, mean: {np.nanmean(y_pred):.4f}")
    logging.info(f"After inverse scaling — y_test min: {np.nanmin(y_test):.4f}, max: {np.nanmax(y_test):.4f}, mean: {np.nanmean(y_test):.4f}")
    logging.info(f"NaNs in y_pred after inverse scaling: {np.isnan(y_pred).sum()}, Infs: {np.isinf(y_pred).sum()}")
    logging.info(f"NaNs in y_test after inverse scaling: {np.isnan(y_test).sum()}, Infs: {np.isinf(y_test).sum()}")

if ground_truth_log_transform:
    logging.info("Applying inverse log transform...")
    y_pred = np.expm1(y_pred)
    y_test = np.expm1(y_test)

    logging.info(f"After inverse log — y_pred min: {np.nanmin(y_pred):.4f}, max: {np.nanmax(y_pred):.4f}")
    logging.info(f"After inverse log — y_test min: {np.nanmin(y_test):.4f}, max: {np.nanmax(y_test):.4f}")
    logging.info(f"NaNs in y_pred after inverse log: {np.isnan(y_pred).sum()}, Infs: {np.isinf(y_pred).sum()}")
    logging.info(f"NaNs in y_test after inverse log: {np.isnan(y_test).sum()}, Infs: {np.isinf(y_test).sum()}")

# Get water mask for test set (same dimensions as full test region)
water_mask_test = load_combined_water_mask(test_years, region, save_data_path, logger=logging)

# Flatten the 2D spatial mask to apply to flattened prediction/gt arrays
water_mask_test_flat = water_mask_test.flatten()

# Apply mask to y_pred and y_test
y_pred_land = y_pred[water_mask_test_flat]
y_test_land = y_test[water_mask_test_flat]

# Log shapes of y_pred and y_test_post
logging.info(f"Predictions shape (land-only): {y_pred_land.shape}")
logging.info(f"Ground Truth shape (land-only): {y_test_land.shape}")

# Log number of predictions outside valid LAI range (0-7)
num_below_0 = np.sum(y_pred_land < lower_bound)
num_above_7 = np.sum(y_pred_land > upper_bound)
logging.info(f"Predictions outside valid LAI range: below 0 = {num_below_0}, above 7 = {num_above_7}")

# Compute bias based on land-only residuals
residuals_land = y_test_land - y_pred_land
bias = np.mean(residuals_land)
logging.info(f"Bias Correction (land-only): {abs(bias):.6f}")

# Apply clipping alone and bias correction  and clipping to land-only and full-grid predictions
y_pred_land_clipped = np.clip(y_pred_land, lower_bound, upper_bound)
y_pred_corrected = y_pred + abs(bias)
y_pred_corrected = np.clip(y_pred_corrected, lower_bound, upper_bound)
y_pred_land_corrected = y_pred_land + abs(bias)
y_pred_land_corrected = np.clip(y_pred_land_corrected, lower_bound, upper_bound)
logging.info("Bias and clipping ([0, 7]) applied to land-only and full predictions.")

# Water mask
y_pred, y_pred_corrected, water_masks = apply_water_mask(y_pred, y_pred_corrected, test_years=test_years, region=region, save_data_path=save_data_path, logger=logging)

# Apply clipping on y_pred (full grid) to have final product (assuming no bias correction)
y_pred = np.clip(y_pred, lower_bound, upper_bound)

# Summary stats
logging.info("Summary of Evaluation (land-only):")
log_stats("Ground Truth", y_test_land)
log_stats("Raw Predictions", y_pred_land)
log_stats("Clipped Predictions", y_pred_land_clipped)
log_stats("Corrected Predictions", y_pred_land_corrected)

# Generate post processing plots 
if post_plots:
    plot_paths = plot_postprocessing_outputs(y_train, y_test, y_pred, y_pred_corrected, save_data_path, region, timestamp, model_name,lat_dim, lon_dim, lower_bound, upper_bound)

# Raw metrics
mse = mean_squared_error(y_test_land, y_pred_land)
rmse = np.sqrt(mse)
mae = mean_absolute_error(y_test_land, y_pred_land)
r2 = r2_score(y_test_land, y_pred_land)
logging.info(f'Raw XGBoost Metrics (land-only) — MSE: {mse:.5f}, RMSE: {rmse:.5f}, MAE: {mae:.5f}, R²: {r2:.5f}')

# Clipped metrics
mse_clipped = mean_squared_error(y_test_land, y_pred_land_clipped)
rmse_clipped = np.sqrt(mse_clipped)
mae_clipped = mean_absolute_error(y_test_land, y_pred_land_clipped)
r2_clipped = r2_score(y_test_land, y_pred_land_clipped)
logging.info(f'Clipped XGBoost Metrics (land-only) — MSE: {mse_clipped:.5f}, RMSE: {rmse_clipped:.5f}, MAE: {mae_clipped:.5f}, R²: {r2_clipped:.5f}')

# Bias corrected and clipped metrics
mse_c = mean_squared_error(y_test_land, y_pred_land_corrected)
rmse_c = np.sqrt(mse_c)
mae_c = mean_absolute_error(y_test_land, y_pred_land_corrected)
r2_c = r2_score(y_test_land, y_pred_land_corrected)
logging.info(f'Bias Corrected and Clipped XGBoost Metrics (land-only) — MSE: {mse_c:.5f}, RMSE: {rmse_c:.5f}, MAE: {mae_c:.5f}, R²: {r2_c:.5f}')

# Save model
model_path = os.path.join(save_data_path, f'XGBoost_LAI_{start_year}_{end_year}_{region}_{timestamp}.joblib')

joblib.dump(model, model_path)
logging.info(f"Model saved to: {model_path}")


# # Write to CSV
# experiment_log_path = os.path.join(save_data_path, "experiment_logs_v2.csv")
# experiment_df = pd.DataFrame([experiment_log])

# if not os.path.exists(experiment_log_path):
#     experiment_df.to_csv(experiment_log_path, index=False)
#     logging.info(f"Experiment log created: {experiment_log_path}")
# else:
#     experiment_df.to_csv(experiment_log_path, mode='a', header=False, index=False)
#     logging.info(f"Experiment log appended: {experiment_log_path}")
