"""
train_rnn_lai_monthly.py - Train an RNN-based Model for LAI Prediction

This script trains a recurrent neural network (RNN) to predict monthly Leaf Area Index (LAI)
values for a given region, using a mix of monthly (temporal) and yearly (static) features.

Supported Features:
- Choice of RNN architecture: vanilla RNN, GRU, or LSTM
- Optional attention mechanism in the decoder
- Dropout regularization
- Teacher forcing for sequence learning
- Output constraint via scaled sigmoid (range [0, 7])
- Log-transform and scaling for ground truth and features
- Imputation strategies for handling missing values
- Persistent water mask filtering to exclude non-land areas
- Early stopping during training
- Automatic experiment logging (CSV)

Pipeline Overview:
1. Loads and preprocesses multi-year LAI data (monthly + yearly features)
2. Trains a backward RNN model to predict current-year LAI using next-year LAI
3. Evaluates predictions on held-out test years with postprocessing
"""

import logging
import os
import argparse
import time
import numpy as np
import pandas as pd
import torch
from torch import nn
from torch import optim
from utils import (
    load_data_for_years, plot_training_curves, log_system_info, log_tensor_stats,
    load_metadata, configure_logging, build_experiment_log
)
from modeling import BackwardLAIPredictor, train_model
from preprocessing import preprocess_lai_data, prepare_dataloaders, load_combined_water_mask
from postprocessing import evaluate_and_postprocess

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a regression model to predict Leaf Area Index.")
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("--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"],
                    help="Type of RNN cell to use: rnn, gru, or lstm", default="rnn")
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")
parser.add_argument("--weight_decay", type=float, default=0.0,
                    help="L2 weight decay (regularization strength) for the optimizer")
parser.add_argument("--log_csv", action="store_true", default=False,
                    help="Whether to log experiment results to a CSV file")
parser.add_argument("--log_csv_file", type=str, default="experiment_log.csv",
                    help="CSV file name for logging experiments")

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
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
log_csv = args.log_csv
log_csv_file = args.log_csv_file

# 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 model name
if TASK_TYPE == "regression":
    MODEL_NAME = "BackwardLAIPredictor"

# Configure logging
log_file_path = configure_logging(start_year, end_year, region, timestamp)

logging.info("Training model of type: %s (%s) with task: %s. Autoregressive: %s.",
             MODEL_TYPE.upper(), cell_type_friendly, TASK_TYPE.upper(), autoregressive)
logging.info("Attention Mechanism Enabled: %s", use_attention)
logging.info("Hidden size: %s", hidden_size)
logging.info("Dropout probability: %s", dropout)
logging.info("Output Constraint Enabled: %s", constrain_output)
logging.info("Teacher Forcing Ratio: %s", teacher_forcing_ratio)
logging.info("Weight decay (L2 normalization): %s", weight_decay)
logging.info("Timestamp: %s", timestamp)

# Log system information
log_system_info()

# Load metadata
metadata = load_metadata(save_data_path, start_year, end_year, region)

# Extract latitude and longitude dimensions from metadata
lat_dim = metadata.get('lat_dim')
lon_dim = metadata.get('lon_dim')

# Extract data path from metadata
data_path = metadata.get('data_path', save_data_path)
logging.info("Using data path from metadata: %s", data_path)
logging.info("Expected dimensions (lat x lon): %s x %s", lat_dim, lon_dim)

# ---- PREPROCESSING ----

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

# Load data for training and testing years, 1 array per year in each list:
# - X_train_list[i]: (num_locations, features) = (15052800, 27)
# - y_train_list[i]: (num_locations, 12)
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("Loaded data for %d training years and %d testing years.",
             num_train_years, num_test_years)
logging.info("Original X_train shapes: %s", [x.shape for x in X_train_list])
logging.info("Original y_train shapes: %s", [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("Applied water mask. Kept %d land pixels out of %d total.",
             np.sum(water_mask), len(water_mask))
    logging.info("Filtered X_train shapes: %s", [x.shape for x in X_train_list])
    logging.info("Filtered y_train shapes: %s", [y.shape for y in y_train_list])

# Check if there are enough training years for validation split
if num_train_years < 2:
    raise ValueError("Need at least 2 years in training data to create a validation split.")

# Split training data, and separate yearly and monthly features
val_years_count = max(1, int(0.2 * num_train_years))
train_years_count = num_train_years - val_years_count

# Stack the lists into an additional dimension: num_years
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)

# For training, take the years from val_years_count on
X_train_monthly = X_train_array[:, val_years_count:, :12] # First 12 features (next year's LAI)
X_train_yearly = X_train_array[:, val_years_count:, 12:] # Last 15 features (yearly features)
y_train = y_train_array[:, val_years_count:, :]

# For validation, take the years up to 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("Training years: %d", train_years_count)
logging.info("Validation years: %d", 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(
    "BEFORE preprocessing - Locations: %d, Yearly Vars: %d, Monthly Vars: %d",
    num_locations, num_yearly_variables, num_monthly_variables)
logging.info(
    "Preprocessing pipeline: scale[0,1]=%s, log=%s, impute=%s",
    ground_truth_scale, ground_truth_log_transform, imputation_method)
logging.info(
    "One-hot encoding HILDA+ (yearly var 15): 14 original + 13 classes => 27 yearly variables")

# 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("%s contains NaNs: %s", name, 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("AFTER preprocessing - Locations: %d, Yearly Variables: %d, Monthly Variables: %d",
             num_locations, num_yearly_variables, num_monthly_variables)

# Build PyTorch DataLoaders
loaders = 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
)

(train_loader, val_loader, test_loader, train_dataset, val_dataset, test_dataset) = loaders

logging.info("Batch size: %d, Train steps/epoch: %d, Val steps: %d, Test steps: %d",
             batch_size, len(train_loader), len(val_loader), len(test_loader))

# Log reprocessing time
preprocessing_duration = time.time() - prep_start_time
logging.info("Preprocessing completed in %.2f seconds.", preprocessing_duration)

# ---- TRAINING ----

# Initialize model, criterion and optimizer
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)
criterion = nn.MSELoss()
optimizer = optim.Adam(MODEL.parameters(), lr=1e-3, weight_decay=weight_decay)

# Count the trainable model parameters
num_model_params = sum(p.numel() for p in MODEL.parameters() if p.requires_grad)
logging.info("Total trainable model parameters: %d", num_model_params)

# Start training step
logging.info("Starting training step for %d epochs", num_epochs)
MODEL, train_losses, val_losses, best_epoch, final_epoch, training_duration = train_model(
    model=MODEL,
    train_loader=train_loader,
    val_loader=val_loader,
    device=device,
    criterion=criterion,
    optimizer=optimizer,
    num_epochs=num_epochs,
    teacher_forcing_ratio=teacher_forcing_ratio,
    use_early_stopping=use_early_stopping,
    early_stopping_patience=early_stopping_patience,
    logger=logging
)

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

# Start evaluation and postprocessing step
logging.info("Starting Evaluation step on test set")
results = evaluate_and_postprocess(
    model=MODEL,
    test_loader=test_loader,
    y_train_post=y_train_post,
    y_test_post=y_test_post,
    scaling_min=scaling_min,
    scaling_max=scaling_max,
    ground_truth_scale=ground_truth_scale,
    ground_truth_log_transform=ground_truth_log_transform,
    region=region,
    timestamp=timestamp,
    model_name=MODEL_NAME,
    save_data_path=save_data_path,
    test_years=test_years,
    lat_dim=lat_dim,
    lon_dim=lon_dim,
    lower_bound=LOWER_BOUND,
    upper_bound=UPPER_BOUND,
    post_plots=post_plots,
    logger=logging
)

# Save model
model_filename = (
    f"{MODEL_NAME}_{rnn_cell_type}_attention_{use_attention}_LAI_"
    f"{start_year}_{end_year}_{region}_{timestamp}.pt"
)
model_path = os.path.join(save_data_path, model_filename)
torch.save(MODEL.state_dict(), model_path)
logging.info("Model saved to: %s", model_path)

if log_csv:
    experiment_log = build_experiment_log(
    args=args,
    model_name=MODEL_NAME,
    model_path=model_path,
    final_epoch=final_epoch,
    best_epoch=best_epoch,
    training_duration=training_duration,
    training_curve_path=(training_curve_path if post_plots else "no_curve"),
    num_model_params=num_model_params,
    num_train_years=train_years_count,
    val_years_count=val_years_count,
    num_test_years=num_test_years,
    num_locations=num_locations,
    num_monthly_variables=num_monthly_variables,
    num_yearly_variables=num_yearly_variables,
    metrics=results["metrics"],
    bias = results["metrics"]["bias"],
    train_losses=train_losses,
    val_losses=val_losses
    )

    experiment_log_path = os.path.join(save_data_path, log_csv_file)
    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("Experiment log created: %s", experiment_log_path)
    else:
        experiment_df.to_csv(experiment_log_path, mode='a', header=False, index=False)
        logging.info("Experiment log appended: %s", experiment_log_path)
