import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import logging
import os
import argparse
import json
import time
import torch
import torch.nn as nn
import torch.optim as optim
from utils import (
    load_data_for_years, inverse_scale_data, evaluate,
    plot_postprocessing_outputs, plot_training_curves,
    log_system_info, log_tensor_stats, log_stats
)
from modeling import BackwardLAIPredictor
from preprocessing import preprocess_lai_data, prepare_dataloaders, load_combined_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("--regions", nargs='+', type=str, required=True, help="Regions for training (e.g. ES EU UK)")
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("--hidden_size", type=int, default=32, help="Hidden size for the RNN cell")
parser.add_argument("--batch_size", type=int, default=1024, help="Batch size for training")
parser.add_argument("--num_epochs", type=int, default=5, help="Number of epochs for training")
parser.add_argument("--early_stopping", action='store_true', default=False, help="Enable early stopping during training")
parser.add_argument("--early_stopping_patience", type=int, default=3, help="Patience for early stopping")
parser.add_argument("--rnn_cell_type", type=str, choices=["rnn", "gru", "lstm"], default="rnn", help="Type of RNN cell to use: rnn, gru, or lstm")
parser.add_argument("--use_attention", action="store_true", default=False, help="Enable attention mechanism in the RNN decoder")
parser.add_argument("--dropout", type=float, default=0.0, help="Dropout probability for RNN hidden state")
parser.add_argument("--constrain_output", action="store_true", help="Apply scaled sigmoid to constrain outputs in [0, 7]")
parser.add_argument("--post_plots", action="store_true", default=False, help="Enable post training plots")
parser.add_argument("--teacher_forcing_ratio", type=float, default=0.0, help="Probability of using teacher forcing during training (using ground truth as input instead of model's prediction)")
parser.add_argument("--weight_decay", type=float, default=0.0, help="L2 weight decay (regularization strength) for the optimizer")

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
regions = args.regions
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
hidden_size = args.hidden_size
batch_size = args.batch_size
num_epochs = args.num_epochs
use_early_stopping = args.early_stopping
early_stopping_patience = args.early_stopping_patience
rnn_cell_type = args.rnn_cell_type
use_attention = args.use_attention
dropout = args.dropout
constrain_output = args.constrain_output
post_plots = args.post_plots
teacher_forcing_ratio = args.teacher_forcing_ratio
weight_decay = args.weight_decay

# Map internal cell type to friendly names
cell_type_friendly = {
    "rnn": "Vanilla RNN",
    "gru": "GRU",
    "lstm": "LSTM"
}.get(rnn_cell_type.lower(), "Unknown RNN Type")

# 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 = "BackwardLAIPredictor"


# Configure logging
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)
region_str = "_".join(regions)
log_file_name = f"train_log_{start_year}_{end_year}_{region_str}_{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 of type: {model_type.upper()} ({cell_type_friendly}) with task: {task_type.upper()}. Autoregressive: {autoregressive}.")
logging.info(f"Attention Mechanism Enabled: {use_attention}")
logging.info(f"Hidden size: {hidden_size}")
logging.info(f"Dropout probability: {dropout}")
logging.info(f"Output Constraint Enabled: {constrain_output}")
logging.info(f"Teacher Forcing Ratio: {teacher_forcing_ratio}")
logging.info(f"Weight decay (L2 normalization): {weight_decay}")
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}_{regions[0]}.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()

X_train_monthly_list = []
X_val_monthly_list = []
X_train_yearly_list = []
X_val_yearly_list = []
y_train_list = []
y_val_list = []

X_test_monthly_list = []
X_test_yearly_list = []
y_test_list = []

# Init a dict to track test region slices
region_slices = {}
test_idx_counter = 0  # counter for indices in the concatenated test array

for region in regions:
    
    logging.info(f"Loading data for region: {region}")
    X_train_list, y_train_list_region = load_data_for_years(train_years, region, save_data_path)
    X_test_list, y_test_list_region = load_data_for_years(test_years, region, save_data_path)
    
    logging.info(f"Training data shapes for {region}: {[x.shape for x in X_train_list]}")
    logging.info(f"Testing data shapes for {region}: {[x.shape for x in X_test_list]}")

    if water_filter:
        water_mask_train = load_combined_water_mask(train_years, region, save_data_path, logger=logging)
        X_train_list = [X[water_mask_train] for X in X_train_list]
        y_train_list_region = [y[water_mask_train] for y in y_train_list_region]

        logging.info(f"Applied water mask to {region}. Kept {np.sum(water_mask_train):,} land points.")
        
        water_mask_test = load_combined_water_mask(test_years, region, save_data_path, logger=logging)
        X_test_list = [X[water_mask_test] for X in X_test_list]
        y_test_list_region = [y[water_mask_test] for y in y_test_list_region]

    X_train_array_region = np.stack(X_train_list, axis=1)
    y_train_array_region = np.stack(y_train_list_region, axis=1)

    X_test_array_region = np.stack(X_test_list, axis=1)
    y_test_array_region = np.stack(y_test_list_region, axis=1)

    # Track region slice for individual region evaluation
    num_test_samples = X_test_array_region.shape[0]
    region_slices[region] = (test_idx_counter, test_idx_counter + num_test_samples)
    test_idx_counter += num_test_samples

    num_train_years = X_train_array_region.shape[1]
    val_years_count = max(1, int(0.2 * num_train_years))

    # Split for this region
    X_train_monthly_list.append(X_train_array_region[:, val_years_count:, :12])
    X_train_yearly_list.append(X_train_array_region[:, val_years_count:, 12:])
    y_train_list.append(y_train_array_region[:, val_years_count:, :])

    X_val_monthly_list.append(X_train_array_region[:, :val_years_count, :12])
    X_val_yearly_list.append(X_train_array_region[:, :val_years_count, 12:])
    y_val_list.append(y_train_array_region[:, :val_years_count, :])

    # Test data (no split)
    X_test_monthly_list.append(X_test_array_region[:, :, :12])
    X_test_yearly_list.append(X_test_array_region[:, :, 12:])
    y_test_list.append(y_test_array_region)

X_train_monthly = np.concatenate(X_train_monthly_list, axis=0)
X_train_yearly = np.concatenate(X_train_yearly_list, axis=0)
y_train = np.concatenate(y_train_list, axis=0)

X_val_monthly = np.concatenate(X_val_monthly_list, axis=0)
X_val_yearly = np.concatenate(X_val_yearly_list, axis=0)
y_val = np.concatenate(y_val_list, axis=0)

X_test_monthly = np.concatenate(X_test_monthly_list, axis=0)
X_test_yearly = np.concatenate(X_test_yearly_list, axis=0)
y_test = np.concatenate(y_test_list, axis=0)

logging.info(f"Concatenated Training set shape (monthly): {X_train_monthly.shape}")
logging.info(f"Concatenated Training set shape (yearly): {X_train_yearly.shape}")
logging.info(f"Concatenated Validation set shape (monthly): {X_val_monthly.shape}")
logging.info(f"Concatenated Validation set shape (yearly): {X_val_yearly.shape}")
logging.info(f"Concatenated Testing set shape (monthly): {X_test_monthly.shape}")
logging.info(f"Concatenated Testing set shape (yearly): {X_test_yearly.shape}")

# Log basic split 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: {num_train_years}")
logging.info(f"Validation years: {X_val_yearly.shape[1]}")
logging.info(f"Test years: {X_test_yearly.shape[1]}")
logging.info(f"Data summary BEFORE preprocessing - Locations: {num_locations}, Yearly Variables: {num_yearly_variables}, Monthly Variables: {num_monthly_variables}")

# Log tensor stats
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:"
)

# Announce preprocessing start
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)")

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_str,
    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 after preprocessing (one-hot encoding could change it!)
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}, Yearly Variables: {num_yearly_variables}, Monthly Variables: {num_monthly_variables}")

# Prepare the dataloaders
train_loader, val_loader, test_loader, train_dataset, val_dataset, test_dataset = prepare_dataloaders(
    X_train_monthly, X_train_yearly, y_train,
    X_val_monthly, X_val_yearly, y_val,
    X_test_monthly, X_test_yearly, y_test,
    batch_size=batch_size,
    num_workers=32,
    num_yearly_variables=num_yearly_variables
)
logging.info(f"Batch size: {batch_size}, Train steps/epoch: {len(train_loader)}, Val steps: {len(val_loader)}, Test steps: {len(test_loader)}")

# Initialize model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = BackwardLAIPredictor(
    hidden_size=hidden_size,
    rnn_cell_type=rnn_cell_type,
    use_attention=use_attention,
    dropout=dropout,
    num_yearly_variables=num_yearly_variables,
    constrain_output=constrain_output
).to(device)

# Define criterion and optimizer
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=weight_decay)

# Start tracking the training time
logging.info(f"Starting training step for {num_epochs} epochs")
start_time = time.time()

train_losses = []
val_losses = []

best_model_state = None
best_val_loss = float('inf')
patience_counter = 0

for epoch in range(num_epochs):
    model.train()
    total_loss = 0

    for lai_next_year, yearly_vars, lai_target in train_loader:
        lai_next_year = lai_next_year.to(device)
        yearly_vars = yearly_vars.to(device)
        lai_target = lai_target.to(device)

        optimizer.zero_grad()
        output = model(lai_next_year, yearly_vars, targets=lai_target, teacher_forcing_ratio=teacher_forcing_ratio)
        
        loss = criterion(output, lai_target)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    avg_train_loss = total_loss / len(train_loader)
    avg_val_loss = evaluate(model, val_loader, device, criterion)
    train_losses.append(avg_train_loss)
    val_losses.append(avg_val_loss)

    logging.info(f"Epoch [{epoch+1}/{num_epochs}] | Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f}")

    # Early stopping
    if use_early_stopping:
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            patience_counter = 0
            best_model_state = model.state_dict()
            logging.info(f"Validation loss improved. Resetting patience.")
        else:
            patience_counter += 1
            logging.info(f"No improvement. Early stopping patience: {patience_counter}/{early_stopping_patience}")
            if patience_counter >= early_stopping_patience:
                logging.info(f"Early stopping triggered after {epoch+1} epochs.")
                break

# Get the best model if early stopping is used
if use_early_stopping and best_model_state is not None:
    model.load_state_dict(best_model_state)
    logging.info("Loaded best model from early stopping checkpoint.")

# Save training curve
training_curve_path = os.path.join(save_data_path, f"training_curve_{region_str}_{timestamp}.png")
plot_training_curves(train_losses, val_losses, filename=training_curve_path)
logging.info(f"Training curve saved to: {training_curve_path}")

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

# Make predictions on the test set
y_pred = []
model.eval()  # Set model to evaluation mode
with torch.no_grad():
    for lai_next_year, yearly_vars, _ in test_loader:
        lai_next_year = lai_next_year.to(device)
        yearly_vars = yearly_vars.to(device)
        output = model(lai_next_year, yearly_vars)
        y_pred.append(output.cpu().numpy())

# Concatenate predictions
y_pred = np.concatenate(y_pred, axis=0)

y_pred = np.flip(y_pred, axis=1)

y_train = y_train_post.reshape(-1, 12)  # Optional if you later plot training predictions
y_test = y_test_post.reshape(-1, 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 ----

# Inverse scaling
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.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()}")

# Inverse log transform
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()}")

# ---- GLOBAL EVALUATION ----
logging.info("Performing global evaluation (all regions combined).")

# Global metrics
mse_global = mean_squared_error(y_test, y_pred)
rmse_global = np.sqrt(mse_global)
mae_global = mean_absolute_error(y_test, y_pred)
r2_global = r2_score(y_test, y_pred)

logging.info(f"Global Test Set Metrics — MSE: {mse_global:.4f}, RMSE: {rmse_global:.4f}, MAE: {mae_global:.4f}, R²: {r2_global:.4f}")

# ---- PER REGION EVALUATION ----
logging.info("Starting per-region evaluation.")

for region in regions:
    logging.info(f"Evaluating region: {region}")

    start, end = region_slices[region]

    y_pred_region = y_pred[start:end]
    y_test_region = y_test[start:end]

    mse_region = mean_squared_error(y_test_region, y_pred_region)
    rmse_region = np.sqrt(mse_region)
    mae_region = mean_absolute_error(y_test_region, y_pred_region)
    r2_region = r2_score(y_test_region, y_pred_region)

    logging.info(f"[{region}] Test Set Metrics — MSE: {mse_region:.4f}, RMSE: {rmse_region:.4f}, MAE: {mae_region:.4f}, R²: {r2_region:.4f}")
