train.py

import argparse
import json
import math
import logging
import sys
import pickle

import numpy as np
from responses import target
import torch
import random
import copy
import pandas as pd
import time
import glob
import random
import matplotlib.pyplot as plt
import os
import shutil
from torch.utils.data import Subset
from datetime import datetime

from torch.utils.data import Dataset, DataLoader, ConcatDataset, random_split
from data_month import LakeMonthlyDataset, AbnormalMonthDataset, DiscriminatorDataset
from data_year import LakeYearlyDataset
from model import BiLSTMModel, LSTM, LSTM_month, Discriminator
import torch.nn.functional as F
from sklearn.metrics import precision_recall_curve, confusion_matrix, roc_curve, classification_report


import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import torch
import os
import torch.nn as nn
import pickle
import sys

import logging
import sys
from datetime import datetime

from torch import cuda
device = 'cuda:0' if cuda.is_available() else 'cpu'


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--steps', type=int, default=3,
                                            help='steps for features')
    parser.add_argument('--feature_size', type=int, default=47,
                                            help='feature size')
    parser.add_argument('--hidden_size', type=int, default=128, # 128
                                            help='hidden dimensions for features')
    parser.add_argument('--output_size', type=int, default=1,
                                            help='output dimensions')
    parser.add_argument('--num_layers_bilstm', type=int, default=2,
                                            help='num_layers of bilstm')
    parser.add_argument('--num_layers_lstm', type=int, default=1,
                                            help='num_layers of lstm')
    parser.add_argument('--batch_size_encoder', type=int, default=16, # 32
                                            help='size for each batch')
    parser.add_argument('--batch_size_yearly_model', type=int, default=1,
                                            help='size for each batch')
    parser.add_argument('--batch_size_monthly_model', type=int, default=4,
                                            help='size for each batch')
    parser.add_argument('--num_epochs_encoder', type=int, default=30,
                                            help='maximum number of epochs for encoder training')
    parser.add_argument('--num_epochs_jointly', type=int, default=6,
                                            help='maximum number of epochs for jointly training')
    parser.add_argument('--num_epochs_pretrain_year', type=int, default=10, 
                                            help='maximum number of epochs for pretraining yearly model')
    parser.add_argument('--num_epochs_year', type=int, default=20,
                                            help='maximum number of epochs for yearly model training')
    parser.add_argument('--num_epochs_month', type=int, default=20, #20
                                        help='maximum number of epochs for monthly model training')
    parser.add_argument('--num_epochs_month_pretrain', type=int, default=30, #5
                                        help='maximum number of epochs for monthly model training')
    parser.add_argument('--discriminator_epochs', type=int, default=50,
                                            help='maximum number of epochs for discriminator training')
    parser.add_argument('--top_N', type=int, default=100,
                                            help='the top N similar embeddings')
    parser.add_argument('--learning_rate_encoder', type=float, default= 0.001,
                                            help='initial learning rate for encoder')
    parser.add_argument('--learning_rate_encoder_ft', type=float, default= 0.0001,
                                            help='initial learning rate for encoder fine-tuning')
    parser.add_argument('--learning_rate_pretrain_yearly_model', type=float, default= 0.01,
                                            help='initial learning rate for pretraining yearly model')
    parser.add_argument('--learning_rate_yearly_model', type=float, default= 0.01,
                                            help='initial learning rate for yearly model')
    parser.add_argument('--learning_rate_monthly_model', type=float, default= 0.0001,
                                            help='initial learning rate for month model') 
    parser.add_argument('--learning_rate_monthly_model_pt', type=float, default= 0.001,
                                            help='initial learning rate for month model') 
    parser.add_argument('--learning_rate_discriminator', type=float, default= 5e-4,
                                            help='initial learning rate for discriminator') 
    parser.add_argument('--weight_decay', type=float, default=1e-8,
                                            help='weight_decay rate')
    parser.add_argument('--seed', type=int, default=21, 
                                            help='seed for random initialisation')
    parser.add_argument('--b', type=float, default=5,
                                            help='weight for loss_epi')
    parser.add_argument('--c', type=float, default=1,
                                            help='weight for loss_hyp')
    parser.add_argument('--a', type=float, default=10.5,             
                                            help='weight for loss_combined')
    parser.add_argument('--task', type=str, default='DO',             
                                            help='DO or Temp')
    parser.add_argument('--save_path', type=str, default='test',
                                            help='all the relative information for a specific training process')
    parser.add_argument('--skip_joint_training', action='store_true',
                                            help='Skip joint training phase and use encoder from Phase 1 with fresh decoder')
    args = parser.parse_args()

    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)
    train(args)

def setup_logging(save_path):
    """Setup logging to both file and console."""
    # Create logs directory
    log_dir = os.path.join(save_path, "logs")
    os.makedirs(log_dir, exist_ok=True)
    
    # Create log filename with timestamp
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_file = os.path.join(log_dir, f"training_log_{timestamp}.log")
    
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)  # Also logger.info to console
        ]
    )
    
    return logging.getLogger(__name__)

def set_random_seeds(seed):
    """Set all random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed) 
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

def setup_data_paths(task):
    if task == 'DO':
        return{
        'month_train': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/DO/processed_month/train',
        'month_val': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/DO/processed_month/val',
        'month_test': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/DO/processed_month/test',
        'year_train': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/DO/processed_year/train',
        'year_val': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/DO/processed_year/val',
        'year_test': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/DO/processed_year/test',
        'cluster_dir': '/home/yuan/codespace/RAG/RAG_lake/Dataset/clusters_9',
        'distance_file': '/home/yuan/codespace/RAG/RAG_lake/Dataset/clusters_9/centroid_distances.csv'
        }
    else:
        return{
        'month_train': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/Temp/processed_month/train',
        'month_val': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/Temp/processed_month/val',
        'month_test': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/Temp/processed_month/test',
        'year_train': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/Temp/processed_year/train',
        'year_val': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/Temp/processed_year/val',
        'year_test': '/home/yuan/codespace/RAG/RAG_lake/Dataset/abnormal_new/Temp/processed_year/test',
        'cluster_dir': '/home/yuan/codespace/RAG/RAG_lake/Dataset/clusters_9',
        'distance_file': '/home/yuan/codespace/RAG/RAG_lake/Dataset/clusters_9/centroid_distances.csv'
        }

def create_datasets_and_loaders(data_config, args):
    all_month_data_dirs = [data_config['month_train'], data_config['month_val'], data_config['month_test']]
    month_data_train_dir = [data_config['month_train']]

    encoder_train_dataset = LakeMonthlyDataset(all_month_data_dirs, data_config['cluster_dir'], data_config['distance_file'], random_seed=args.seed)
    encoder_decoder_train_dataset = LakeMonthlyDataset(month_data_train_dir, data_config['cluster_dir'], data_config['distance_file'], random_seed=args.seed)
    yearly_model_train_dataset = LakeYearlyDataset(data_config['year_train'])
    yearly_model_val_dataset = LakeYearlyDataset(data_config['year_val'])
    yearly_model_test_dataset = LakeYearlyDataset(data_config['year_test'])

    dl_encoder_train = DataLoader(encoder_train_dataset, args.batch_size_encoder, shuffle=True)
    dl_encoder_decoder_train = DataLoader(encoder_decoder_train_dataset, args.batch_size_yearly_model, shuffle=True)
    dl_yearly_model_train = DataLoader(yearly_model_train_dataset, args.batch_size_yearly_model, shuffle=True)
    dl_yearly_model_val = DataLoader(yearly_model_val_dataset, args.batch_size_yearly_model, shuffle=True) # False
    dl_yearly_model_test = DataLoader(yearly_model_test_dataset, args.batch_size_yearly_model, shuffle=False)
    
    datasets = {
        'encoder_train': encoder_train_dataset,
        'encoder_decoder_train': encoder_decoder_train_dataset,
        'yearly_train': yearly_model_train_dataset,
        'yearly_val': yearly_model_val_dataset,
        'yearly_test': yearly_model_test_dataset
    }
    
    dataloaders = {
        'encoder_train': dl_encoder_train,
        'encoder_decoder_train': dl_encoder_decoder_train,
        'yearly_train': dl_yearly_model_train,
        'yearly_val': dl_yearly_model_val,
        'yearly_test': dl_yearly_model_test
    }
    
    return datasets, dataloaders

def initialize_models(args):
    encoder = BiLSTMModel(args.feature_size, args.hidden_size, args.output_size, args.num_layers_bilstm)  # use simulated labels in encoder training, has to learn the joint distribution of x and y
    yearly_model = LSTM(args.feature_size - 2, args.hidden_size, args.output_size, args.num_layers_lstm, args.batch_size_yearly_model)  # don't use simulated labels in yearly model training

    discriminator_epi = Discriminator(args.feature_size + 4) # use simulated labels in discriminator training
    discriminator_hyp = Discriminator(args.feature_size + 4)

    decoder_epi = LSTM_month(args.feature_size-21, args.hidden_size, args.output_size, args.num_layers_lstm, args.batch_size_monthly_model)
    decoder_hyp = LSTM_month(args.feature_size-21, args.hidden_size, args.output_size, args.num_layers_lstm, args.batch_size_monthly_model)

    encoder.to(device)
    yearly_model.to(device)
    discriminator_epi.to(device)
    discriminator_hyp.to(device)
    decoder_epi.to(device)
    decoder_hyp.to(device)

    return {
        "encoder": encoder,
        "yearly_model": yearly_model,
        "discriminator_epi": discriminator_epi,
        "discriminator_hyp": discriminator_hyp,
        "decoder_epi": decoder_epi,
        "decoder_hyp": decoder_hyp
    }

# =============== Phase 1: Train encoder ===========
def train_encoder_phase_with_loading(args, models, dataloaders, logger):
    """
    Handle encoder training with option to skip if already trained.
    
    Args:
        args: Training arguments
        models: Dictionary containing models
        dataloaders: Dictionary containing dataloaders
        skip_if_exists: If True, skip training if encoder already exists
    """
    encoder = models['encoder']
    save_dir0 = os.path.join(args.save_path, "yearly_model_and_pretrained_encoder")
    os.makedirs(save_dir0, exist_ok=True)
    
    # Check if encoder already exists
    encoder_files = [f for f in os.listdir(save_dir0) if f.startswith('encoder_') and f.endswith('.pth')]
    embeddings_exist = (
        os.path.exists(os.path.join(save_dir0, "precomputed_embeddings_epi.pt")) and
        os.path.exists(os.path.join(save_dir0, "precomputed_embeddings_hyp.pt"))
    )
    
    if encoder_files and embeddings_exist:
        logger.info("Found existing encoder and embeddings. Loading from disk...")
        models = load_trained_encoder(args, models, logger)
        return models
    
    # If not skipping or files don't exist, train the encoder
    logger.info("Training encoder from scratch...")
    start_time = time.time()
    try:
        optimizer_encoder = torch.optim.Adam(encoder.parameters(), lr=args.learning_rate_encoder)
        logger.info("Training encoder with standard contrastive loss...")
        encoder = train_encoder(encoder, optimizer_encoder, dataloaders['encoder_train'], args.num_epochs_encoder, save_dir0, logger)
        save_precomputed_samples(encoder, dataloaders['encoder_train'], save_dir0, logger)
    finally:
        torch.cuda.empty_cache()  # Always clear cache
    
    end_time = time.time()
    logger.info(f"encoder_train_time: {end_time - start_time:.2f} seconds")
    
    # Update models dictionary
    models['encoder'] = encoder
    return models

def load_trained_encoder(args, models, logger):
    """
    Load a previously trained encoder from disk.
    
    Args:
        args: Training arguments
        models: Dictionary containing models
        
    Returns:
        Updated models dictionary with loaded encoder
    """
    save_dir0 = os.path.join(args.save_path, "yearly_model_and_pretrained_encoder")
    
    # Look for encoder files in the directory
    encoder_files = [f for f in os.listdir(save_dir0) if f.startswith('encoder_') and f.endswith('.pth')]
    
    if not encoder_files:
        raise FileNotFoundError(f"No encoder files found in {save_dir0}")
    
    # Use the most recent encoder file (or you can specify a particular one)
    encoder_files.sort()  # Sort by name (which includes timestamp)
    latest_encoder_file = encoder_files[-1]
    encoder_path = os.path.join(save_dir0, latest_encoder_file)
    
    logger.info(f"Loading encoder from: {encoder_path}")
    
    # Load the encoder state dict
    encoder_state = torch.load(encoder_path, map_location=device)
    models['encoder'].load_state_dict(encoder_state)
    models['encoder'].to(device)
    # models['encoder'].eval()
    
    logger.info(f"Successfully loaded encoder from {encoder_path}")
    return models

# =============== Phase 2: Pretrained yearly model ==========
def pretrained_yearly_model(args, models, dataloaders=None, logger=None):
    """Load pretrained yearly model based on task, or train one if not found."""
    yearly_model = models['yearly_model']
    
    # Define model paths based on task
    if args.task == 'DO':
        model_path = 'yearly_model_2025_03_19_17_48_23.pth'
    else:
        model_path = 'yearly_model_2025_03_29_16_58_39.pth'

    # Check if the model file exists
    if os.path.exists(model_path):
        logger.info(f"Loading existing pretrained model from: {model_path}")
        try:
            yearly_model.load_state_dict(torch.load(model_path, map_location=device))
            yearly_model.to(device)
            models['yearly_model'] = yearly_model
            logger.info("Successfully loaded pretrained yearly model")
            return models
        except Exception as e:
            logger.info(f"Failed to load pretrained model: {e}")
            logger.info("Proceeding to train a new model...")
    else:
        logger.info(f"Pretrained model not found at: {model_path}")
        logger.info("Training a new yearly model...")
    
    # If model doesn't exist or failed to load, train a new one
    if dataloaders is None:
        raise ValueError("Dataloaders are required for training but were not provided")
    
    save_dir0 = os.path.join(args.save_path, "yearly_model_and_pretrained_encoder")
    os.makedirs(save_dir0, exist_ok=True)
    
    # First, pretrain with simulation data
    logger.info("Phase 1: Pretraining yearly model with simulation data...")
    optimizer_yearly_model = torch.optim.Adagrad(
        yearly_model.parameters(), 
        lr=args.learning_rate_pretrain_yearly_model, 
        weight_decay=args.weight_decay
    )
    yearly_model = pretrain_yearly_model_use_sim(
        yearly_model, 
        optimizer_yearly_model, 
        dataloaders['yearly_train'], 
        dataloaders['yearly_val'], 
        args.num_epochs_pretrain_year, 
        save_dir0,
        logger
    )
    
    # Then, fine-tune with real observations
    logger.info("Phase 2: Fine-tuning yearly model with real observations...")
    optimizer_yearly_model = torch.optim.Adagrad(
        yearly_model.parameters(), 
        lr=args.learning_rate_yearly_model, 
        weight_decay=args.weight_decay
    )
    yearly_model = train_yearly_model(
        yearly_model, 
        optimizer_yearly_model, 
        dataloaders['yearly_train'], 
        dataloaders['yearly_val'], 
        args.num_epochs_year, 
        save_dir0, 
        logger
    )
    
    # Test the trained model
    logger.info("Testing the trained yearly model...")
    test_yearly_model(yearly_model, dataloaders['yearly_test'], args.save_path, logger)
    
    # Update models dictionary
    models['yearly_model'] = yearly_model
    logger.info("Yearly model training completed successfully")
    
    return models

# =============== Phase 3: Jointly train encoder and decoders ==========
def joint_training_phase(args, models, dataloaders, logger):
    encoder = models['encoder']
    decoder_epi = models['decoder_epi']
    decoder_hyp = models['decoder_hyp']
    dl_encoder_train = dataloaders['encoder_train']
    dl_encoder_decoder_train = dataloaders['encoder_decoder_train']
    
    # Setup save directories
    save_dir0 = os.path.join(args.save_path, "yearly_model_and_pretrained_encoder")
    save_dir1 = os.path.join(args.save_path, "jointly_train_encoder_decoder")
    os.makedirs(save_dir1, exist_ok=True)
    
    # Check if jointly trained models already exist
    if check_existing_joint_training(args.save_path, logger):
        logger.info("Found existing jointly trained models. Loading from disk...")
        models = load_jointly_trained_models(args.save_path, models, logger)
        return models
    
    # Load precomputed samples and embeddings
    precomputed_embeddings_path_epi = os.path.join(save_dir0, "precomputed_embeddings_epi.pt")
    precomputed_samples_path_epi = os.path.join(save_dir0, "precomputed_samples_epi.pkl")
    precomputed_embeddings_epi = torch.load(precomputed_embeddings_path_epi)
    with open(precomputed_samples_path_epi, 'rb') as f:
        precomputed_samples_epi = pickle.load(f)
    precomputed_embeddings_epi = precomputed_embeddings_epi.to(device)

    precomputed_embeddings_path_hyp = os.path.join(save_dir0, "precomputed_embeddings_hyp.pt")
    precomputed_samples_path_hyp = os.path.join(save_dir0, "precomputed_samples_hyp.pkl")
    precomputed_embeddings_hyp = torch.load(precomputed_embeddings_path_hyp)
    with open(precomputed_samples_path_hyp, 'rb') as f:
        precomputed_samples_hyp = pickle.load(f)
    precomputed_embeddings_hyp = precomputed_embeddings_hyp.to(device)

    # Joint training
    start_time = time.time()
    optimizer_encoder_decoder = torch.optim.Adam([
        {'params': encoder.parameters(), 'lr': args.learning_rate_encoder_ft},
        {'params': decoder_epi.parameters(), 'lr': args.learning_rate_monthly_model},
        {'params': decoder_hyp.parameters(), 'lr': args.learning_rate_monthly_model}
    ], weight_decay=1e-5)
    
    encoder, decoder_epi, decoder_hyp = jointly_train_encoder_decoder(
        encoder, decoder_epi, decoder_hyp, optimizer_encoder_decoder, 
        dl_encoder_decoder_train, args.num_epochs_jointly, args.top_N, 
        args.a, args.b, args.c, save_dir1, precomputed_embeddings_epi, 
        precomputed_samples_epi, precomputed_embeddings_hyp, precomputed_samples_hyp, logger
    )
    end_time = time.time()
    jointly_train_time = end_time - start_time
    logger.info(f"jointly_train_time: {jointly_train_time:.2f} seconds")
    torch.cuda.empty_cache()
    
    try:
        save_precomputed_samples(encoder, dl_encoder_train, save_dir1, logger)
    finally:
        torch.cuda.empty_cache()  # Always clear cache

    # Save decoder states
    save_dir1_1 = os.path.join(args.save_path, "pretrained_decoder")
    os.makedirs(save_dir1_1, exist_ok=True)
    decoder_epi_path = os.path.join(save_dir1_1, "decoder_epi.pth")
    decoder_hyp_path = os.path.join(save_dir1_1, "decoder_hyp.pth")
    torch.save(decoder_epi.state_dict(), decoder_epi_path)
    torch.save(decoder_hyp.state_dict(), decoder_hyp_path)
    
    # Update models dictionary
    models['encoder'] = encoder
    models['decoder_epi'] = decoder_epi
    models['decoder_hyp'] = decoder_hyp

def check_existing_joint_training(save_path, logger):
    """
    Check if jointly trained models already exist.
    
    Args:
        save_path: Base save path
        logger: Logger instance
        
    Returns:
        bool: True if jointly trained models exist, False otherwise
    """
    save_dir1 = os.path.join(save_path, "jointly_train_encoder_decoder")
    save_dir1_1 = os.path.join(save_path, "pretrained_decoder")
    
    if not os.path.exists(save_dir1):
        logger.info(f"Joint training directory {save_dir1} does not exist")
        return False
    
    if not os.path.exists(save_dir1_1):
        logger.info(f"Pretrained decoder directory {save_dir1_1} does not exist")
        return False
    
    return True

def load_jointly_trained_models(save_path, models, logger):
    """
    Load jointly trained models from disk.
    
    Args:
        save_path: Base save path
        models: Dictionary containing models
        logger: Logger instance
        
    Returns:
        Updated models dictionary with loaded jointly trained models
    """
    save_dir1 = os.path.join(save_path, "jointly_train_encoder_decoder")
    save_dir1_1 = os.path.join(save_path, "pretrained_decoder")
    
    try:
        # Load encoder - first try to find encoder files with timestamp
        encoder_files = [f for f in os.listdir(save_dir1) if f.endswith('.pth') and 'encoder' in f.lower()]
        if encoder_files:
            # Use the most recent encoder file
            encoder_files.sort()
            encoder_file = encoder_files[-1]
            encoder_path = os.path.join(save_dir1, encoder_file)
        else:
            # Try generic encoder.pth
            encoder_path = os.path.join(save_dir1, "encoder.pth")
        
        if os.path.exists(encoder_path):
            logger.info(f"Loading jointly trained encoder from: {encoder_path}")
            encoder_state = torch.load(encoder_path, map_location=device)
            models['encoder'].load_state_dict(encoder_state)
            models['encoder'].to(device)
        else:
            logger.info("No encoder file found, keeping current encoder")
        
        # Load decoders
        decoder_epi_path = os.path.join(save_dir1_1, "decoder_epi.pth")
        decoder_hyp_path = os.path.join(save_dir1_1, "decoder_hyp.pth")
        
        logger.info(f"Loading jointly trained decoder_epi from: {decoder_epi_path}")
        decoder_epi_state = torch.load(decoder_epi_path, map_location=device)
        models['decoder_epi'].load_state_dict(decoder_epi_state)
        models['decoder_epi'].to(device)
        
        logger.info(f"Loading jointly trained decoder_hyp from: {decoder_hyp_path}")
        decoder_hyp_state = torch.load(decoder_hyp_path, map_location=device)
        models['decoder_hyp'].load_state_dict(decoder_hyp_state)
        models['decoder_hyp'].to(device)
        
        logger.info("Successfully loaded jointly trained models")
        
    except Exception as e:
        logger.info(f"Error loading jointly trained models: {e}")
        raise e
        
    return models

def skip_joint_training_phase(args, models, logger):
    """
    Skip joint training phase and use encoder from Phase 1 with fresh decoder initialization.
    
    Args:
        args: Training arguments
        models: Dictionary containing models
        dataloaders: Dictionary containing dataloaders
    """
    encoder = models['encoder']
    decoder_epi = models['decoder_epi']  # Already fresh from initialize_models
    decoder_hyp = models['decoder_hyp']  # Already fresh from initialize_models
    
    # Setup save directories
    save_dir0 = os.path.join(args.save_path, "yearly_model_and_pretrained_encoder")
    save_dir1 = os.path.join(args.save_path, "jointly_train_encoder_decoder")
    os.makedirs(save_dir1, exist_ok=True)
    
    logger.info("Using encoder from Phase 1 and existing fresh decoders...")
    logger.info("Note: Decoders are already initialized fresh from initialize_models function")
    
    # Copy precomputed embeddings and samples from Phase 1 to joint training directory
    files_to_copy = [
        "precomputed_embeddings_epi.pt",
        "precomputed_samples_epi.pkl",
        "precomputed_embeddings_hyp.pt",
        "precomputed_samples_hyp.pkl"
    ]
    
    for file_name in files_to_copy:
        src_path = os.path.join(save_dir0, file_name)
        dst_path = os.path.join(save_dir1, file_name)
        if os.path.exists(src_path):
            logger.info(f"Copying {file_name} from Phase 1 to joint training directory...")
            shutil.copy2(src_path, dst_path)
        else:
            logger.info(f"Warning: {file_name} not found in Phase 1 directory")
    
    # Save the encoder and fresh decoders to joint training directory
    logger.info("Saving encoder and fresh decoders to joint training directory...")
    torch.save(encoder.state_dict(), os.path.join(save_dir1, "encoder.pth"))
    torch.save(decoder_epi.state_dict(), os.path.join(save_dir1, "decoder_epi.pth"))
    torch.save(decoder_hyp.state_dict(), os.path.join(save_dir1, "decoder_hyp.pth"))
    
    # Also save to pretrained_decoder directory for compatibility with later phases
    save_dir1_1 = os.path.join(args.save_path, "pretrained_decoder")
    os.makedirs(save_dir1_1, exist_ok=True)
    torch.save(decoder_epi.state_dict(), os.path.join(save_dir1_1, "decoder_epi.pth"))
    torch.save(decoder_hyp.state_dict(), os.path.join(save_dir1_1, "decoder_hyp.pth"))
    
    logger.info("Joint training phase skipped successfully!")
    logger.info("Using encoder from Phase 1 and existing fresh decoders.")

# =============== Phase 4: Train discriminators ==========
def discriminator_training_phase(args, models, dataloaders, logger):
    """Handle discriminator training phase."""
    yearly_model = models['yearly_model']
    discriminator_epi = models['discriminator_epi']
    discriminator_hyp = models['discriminator_hyp']
    decoder_epi = models['decoder_epi']
    decoder_hyp = models['decoder_hyp']
    encoder = models['encoder']
    dl_yearly_model_val = dataloaders['yearly_val']
    
    # Setup directories
    save_dir1 = os.path.join(args.save_path, "jointly_train_encoder_decoder")
    save_dir1_1 = os.path.join(args.save_path, "pretrained_decoder")
    save_dir3 = os.path.join(args.save_path, "train_discriminator")
    os.makedirs(save_dir3, exist_ok=True)

    # Check if discriminators already exist
    if check_existing_discriminators(args.save_path, logger):
        logger.info("Found existing discriminators. Loading from disk...")
        models = load_trained_discriminators(args.save_path, models, logger)
        return models
    
    # Load precomputed embeddings from joint training phase
    precomputed_embeddings_path_epi = os.path.join(save_dir1, "precomputed_embeddings_epi.pt")
    precomputed_samples_path_epi = os.path.join(save_dir1, "precomputed_samples_epi.pkl")
    precomputed_embeddings_epi = torch.load(precomputed_embeddings_path_epi)
    with open(precomputed_samples_path_epi, 'rb') as f:
        precomputed_samples_epi = pickle.load(f)
    precomputed_embeddings_epi = precomputed_embeddings_epi.to(device)

    precomputed_embeddings_path_hyp = os.path.join(save_dir1, "precomputed_embeddings_hyp.pt")
    precomputed_samples_path_hyp = os.path.join(save_dir1, "precomputed_samples_hyp.pkl")
    precomputed_embeddings_hyp = torch.load(precomputed_embeddings_path_hyp)
    with open(precomputed_samples_path_hyp, 'rb') as f:
        precomputed_samples_hyp = pickle.load(f)
    precomputed_embeddings_hyp = precomputed_embeddings_hyp.to(device)
    
    # Get decoder paths
    decoder_epi_path = os.path.join(save_dir1_1, "decoder_epi.pth")
    decoder_hyp_path = os.path.join(save_dir1_1, "decoder_hyp.pth")
    
    # Generate labels for discriminator training
    start_time = time.time()
    df_epi, df_hyp = run_decoder_training_and_save_labels_sim(
        yearly_model, dl_yearly_model_val, decoder_epi, decoder_hyp, encoder, 
        args.learning_rate_monthly_model, args.learning_rate_monthly_model_pt, 
        args.weight_decay, args.num_epochs_month, args.num_epochs_month_pretrain, 
        args.top_N, save_dir3, precomputed_embeddings_epi, precomputed_samples_epi, 
        precomputed_embeddings_hyp, precomputed_samples_hyp, decoder_epi_path, decoder_hyp_path, logger=logger
    )
    end_time = time.time()
    decoder_generate_label_time = end_time - start_time
    logger.info(f"decoder_generate_label_time: {decoder_generate_label_time:.2f} seconds")

    # Train discriminators
    epi_dataset = DiscriminatorDataset(df_epi, seq_length=30)
    hyp_dataset = DiscriminatorDataset(df_hyp, seq_length=30)
    dl_discriminator_epi = DataLoader(epi_dataset, batch_size=32, shuffle=True, drop_last=True)
    dl_discriminator_hyp = DataLoader(hyp_dataset, batch_size=32, shuffle=True, drop_last=True)
    
    optimizer_discriminator_epi = torch.optim.Adam(discriminator_epi.parameters(), lr=args.learning_rate_discriminator)
    optimizer_discriminator_hyp = torch.optim.Adam(discriminator_hyp.parameters(), lr=args.learning_rate_discriminator)
    
    start_time = time.time()
    discriminator_epi, optimal_threshold_epi = discriminator_training(
        discriminator_epi, optimizer_discriminator_epi, dl_discriminator_epi, 
        args.discriminator_epochs + 20, save_dir3, flag='epi', logger=logger
    )
    discriminator_hyp, optimal_threshold_hyp = discriminator_training(
        discriminator_hyp, optimizer_discriminator_hyp, dl_discriminator_hyp, 
        args.discriminator_epochs, save_dir3, flag='hyp', logger=logger
    )
    end_time = time.time()
    train_discriminator_time = end_time - start_time
    logger.info(f"Train_discriminator_time: {train_discriminator_time:.2f} seconds")
    
    save_optimal_thresholds(optimal_threshold_epi, optimal_threshold_hyp, save_dir3, logger)

    # Update models dictionary with trained discriminators
    models['discriminator_epi'] = discriminator_epi
    models['discriminator_hyp'] = discriminator_hyp
    
    # Store optimal thresholds in models dict for final testing
    models['optimal_threshold_epi'] = optimal_threshold_epi
    models['optimal_threshold_hyp'] = optimal_threshold_hyp

def check_existing_discriminators(save_path, logger):
    """
    Check if trained discriminators already exist.
    
    Args:
        save_path: Base save path
        
    Returns:
        bool: True if both discriminators exist, False otherwise
    """
    save_dir3 = os.path.join(save_path, "train_discriminator")
    
    if not os.path.exists(save_dir3):
        logger.info(f"Discriminator directory {save_dir3} does not exist")
        return False
    
    # Check for discriminator model files
    discriminator_epi_files = [f for f in os.listdir(save_dir3) if f.startswith('discriminator_epi') and f.endswith('.pth')]
    discriminator_hyp_files = [f for f in os.listdir(save_dir3) if f.startswith('discriminator_hyp') and f.endswith('.pth')]
    
    if not discriminator_epi_files or not discriminator_hyp_files:
        logger.info(f"Missing discriminator model files in {save_dir3}")
        return False
    
    logger.info(f"Found existing discriminator models in {save_dir3}")
    return True

def load_trained_discriminators(save_path, models, logger):
    """
    Load previously trained discriminators from disk.
    
    Args:
        save_path: Base save path
        models: Dictionary containing models
        
    Returns:
        Updated models dictionary with loaded discriminators and thresholds
    """
    save_dir3 = os.path.join(save_path, "train_discriminator")
    
    # Look for discriminator files
    discriminator_epi_files = [f for f in os.listdir(save_dir3) if f.startswith('discriminator_epi') and f.endswith('.pth')]
    discriminator_hyp_files = [f for f in os.listdir(save_dir3) if f.startswith('discriminator_hyp') and f.endswith('.pth')]
    
    if not discriminator_epi_files or not discriminator_hyp_files:
        raise FileNotFoundError(f"Discriminator models not found in {save_dir3}")
    
    # Use the most recent discriminator files
    discriminator_epi_files.sort()
    discriminator_hyp_files.sort()
    latest_epi_file = discriminator_epi_files[-1]
    latest_hyp_file = discriminator_hyp_files[-1]
    
    discriminator_epi_path = os.path.join(save_dir3, latest_epi_file)
    discriminator_hyp_path = os.path.join(save_dir3, latest_hyp_file)
    
    logger.info(f"Loading discriminator_epi from: {discriminator_epi_path}")
    logger.info(f"Loading discriminator_hyp from: {discriminator_hyp_path}")
    
    # Load the discriminator state dicts
    discriminator_epi_state = torch.load(discriminator_epi_path, map_location=device)
    discriminator_hyp_state = torch.load(discriminator_hyp_path, map_location=device)
    
    models['discriminator_epi'].load_state_dict(discriminator_epi_state)
    models['discriminator_hyp'].load_state_dict(discriminator_hyp_state)
    models['discriminator_epi'].to(device)
    models['discriminator_hyp'].to(device)
    
    # Try to load optimal thresholds if they exist
    threshold_epi_file = os.path.join(save_dir3, "optimal_threshold_epi.txt")
    threshold_hyp_file = os.path.join(save_dir3, "optimal_threshold_hyp.txt")
    
    if os.path.exists(threshold_epi_file):
        with open(threshold_epi_file, 'r') as f:
            optimal_threshold_epi = float(f.read().strip())
        logger.info(f"Loaded optimal threshold epi: {optimal_threshold_epi}")
    else:
        optimal_threshold_epi = 0.5
        logger.info(f"Optimal threshold epi file not found, using default: {optimal_threshold_epi}")
    
    if os.path.exists(threshold_hyp_file):
        with open(threshold_hyp_file, 'r') as f:
            optimal_threshold_hyp = float(f.read().strip())
        logger.info(f"Loaded optimal threshold hyp: {optimal_threshold_hyp}")
    else:
        optimal_threshold_hyp = 0.5
        logger.info(f"Optimal threshold hyp file not found, using default: {optimal_threshold_hyp}")
    
    # Store optimal thresholds in models dict
    models['optimal_threshold_epi'] = optimal_threshold_epi
    models['optimal_threshold_hyp'] = optimal_threshold_hyp
    
    logger.info(f"Successfully loaded discriminators from {save_dir3}")
    return models

def save_optimal_thresholds(optimal_threshold_epi, optimal_threshold_hyp, save_path, logger):
    """
    Save optimal thresholds to text files for later loading.
    
    Args:
        optimal_threshold_epi: Optimal threshold for epi discriminator
        optimal_threshold_hyp: Optimal threshold for hyp discriminator
        save_path: Directory to save threshold files
    """
    threshold_epi_file = os.path.join(save_path, "optimal_threshold_epi.txt")
    threshold_hyp_file = os.path.join(save_path, "optimal_threshold_hyp.txt")
    
    with open(threshold_epi_file, 'w') as f:
        f.write(str(optimal_threshold_epi))
    
    with open(threshold_hyp_file, 'w') as f:
        f.write(str(optimal_threshold_hyp))
    
    logger.info(f"Saved optimal thresholds: epi={optimal_threshold_epi}, hyp={optimal_threshold_hyp}")

# =============== Phase 5: Final testing phase ==========
def final_testing_phase(args, models, dataloaders, logger):
    """Handle final testing phase with discriminator predictions."""
    yearly_model = models['yearly_model']
    discriminator_epi = models['discriminator_epi']
    discriminator_hyp = models['discriminator_hyp']
    decoder_epi = models['decoder_epi']
    decoder_hyp = models['decoder_hyp']
    encoder = models['encoder']
    dl_yearly_model_test = dataloaders['yearly_test']
    
    optimal_threshold_epi = models.get('optimal_threshold_epi', 0.5)
    optimal_threshold_hyp = models.get('optimal_threshold_hyp', 0.5)
    
    # Setup directories
    save_dir1 = os.path.join(args.save_path, "jointly_train_encoder_decoder")
    save_dir1_1 = os.path.join(args.save_path, "pretrained_decoder")
    
    # Load precomputed embeddings
    precomputed_embeddings_path_epi = os.path.join(save_dir1, "precomputed_embeddings_epi.pt")
    precomputed_samples_path_epi = os.path.join(save_dir1, "precomputed_samples_epi.pkl")
    precomputed_embeddings_epi = torch.load(precomputed_embeddings_path_epi)
    with open(precomputed_samples_path_epi, 'rb') as f:
        precomputed_samples_epi = pickle.load(f)
    precomputed_embeddings_epi = precomputed_embeddings_epi.to(device)

    precomputed_embeddings_path_hyp = os.path.join(save_dir1, "precomputed_embeddings_hyp.pt")
    precomputed_samples_path_hyp = os.path.join(save_dir1, "precomputed_samples_hyp.pkl")
    precomputed_embeddings_hyp = torch.load(precomputed_embeddings_path_hyp)
    with open(precomputed_samples_path_hyp, 'rb') as f:
        precomputed_samples_hyp = pickle.load(f)
    precomputed_embeddings_hyp = precomputed_embeddings_hyp.to(device)
    
    # Get decoder paths
    decoder_epi_path = os.path.join(save_dir1_1, "decoder_epi.pth")
    decoder_hyp_path = os.path.join(save_dir1_1, "decoder_hyp.pth")
    
    # Final testing with optimal thresholds
    thresholds_epi = [optimal_threshold_epi]
    thresholds_hyp = [optimal_threshold_hyp]

    for thresh_epi, thresh_hyp in zip(thresholds_epi, thresholds_hyp):
        save_dir4_1 = os.path.join(args.save_path, "final_results")
        os.makedirs(save_dir4_1, exist_ok=True)
        logger.info(f"Testing with epi threshold: {thresh_epi} and hyp threshold: {thresh_hyp}")
        
        start_time = time.time()
        if args.task == 'DO':
            discriminator_testing(
                yearly_model, dl_yearly_model_test, discriminator_epi, thresh_epi, 
                discriminator_hyp, thresh_hyp, decoder_epi, decoder_hyp, encoder, 
                args.learning_rate_monthly_model, args.learning_rate_monthly_model_pt, 
                args.weight_decay, args.num_epochs_month, args.num_epochs_month_pretrain, 
                args.top_N, save_dir4_1, precomputed_embeddings_epi, precomputed_samples_epi, 
                precomputed_embeddings_hyp, precomputed_samples_hyp, decoder_epi_path, decoder_hyp_path, logger=logger
            )
        else:
            discriminator_testing_temp(
                yearly_model, dl_yearly_model_test, discriminator_epi, thresh_epi, 
                discriminator_hyp, thresh_hyp, decoder_epi, decoder_hyp, encoder, 
                args.learning_rate_monthly_model, args.learning_rate_monthly_model_pt, 
                args.weight_decay, args.num_epochs_month, args.num_epochs_month_pretrain, 
                args.top_N, save_dir4_1, precomputed_embeddings_epi, precomputed_samples_epi, 
                precomputed_embeddings_hyp, precomputed_samples_hyp, decoder_epi_path, decoder_hyp_path, logger=logger
            )
        end_time = time.time()
        elapsed_time = end_time - start_time
        logger.info(f"Inference Time: {elapsed_time:.2f} seconds")

def train(args):
    # Setup logging at the beginning
    logger = setup_logging(args.save_path)
    logger.info(f"Starting training with arguments: {vars(args)}")
    
    set_random_seeds(args.seed)
    logger.info(f"Random seed set to: {args.seed}")
    
    # Setup data paths and datasets
    data_config = setup_data_paths(args.task)
    datasets, dataloaders = create_datasets_and_loaders(data_config, args)
    models = initialize_models(args)
    logger.info("Data and models initialized successfully")

    # Phase 1: Train encoder (now with optional hard negative mining)
    logger.info("Starting Phase 1: Encoder training")
    models = train_encoder_phase_with_loading(args, models, dataloaders, logger)
    logger.info("Phase 1 completed")

    # Phase 2: Load pretrained yearly model
    logger.info("Starting Phase 2: Loading/Training yearly model")
    models = pretrained_yearly_model(args, models, dataloaders, logger)
    logger.info("Phase 2 completed")

    # Phase 3: Joint training of encoder and decoder
    logger.info("Starting Phase 3: Joint training or skipping")
    if args.skip_joint_training:
        logger.info("Skipping joint training as requested")
        skip_joint_training_phase(args, models, logger)
    else:
        logger.info("Performing joint training")
        joint_training_phase(args, models, dataloaders, logger)
    logger.info("Phase 3 completed")

    # Phase 4: Train discriminator
    logger.info("Starting Phase 4: Discriminator training")
    discriminator_training_phase(args, models, dataloaders, logger)
    logger.info("Phase 4 completed")

    # Phase 5: Final testing
    logger.info("Starting Phase 5: Final testing")
    final_testing_phase(args, models, dataloaders, logger)
    logger.info("Phase 5 completed")
    
    logger.info("Training completed successfully!")

# ========== Train Encoder and save embeddings ==========
def contrastive_loss(anchor_embs, positive_embs, negative_embs, semi_positive_embs):
    """
    Args:
        anchor_embs (Tensor): [batch_size, hidden_dim] Anchor embeddings.
        positive_embs (Tensor): [batch_size, hidden_dim] Positive embeddings (similar to anchor).
        negative_embs (Tensor): [batch_size, hidden_dim] Negative embeddings (dissimilar to anchor).
        semi_positive_embs (Tensor, optional): [batch_size, hidden_dim] Semi-positive embeddings.
        margin (float): Margin for separating positives and negatives. Default is 0.5.

    Returns:
        Tensor: The averaged contrastive loss over the batch.
    """
    # Normalize the embeddings
    anchor_embs = F.normalize(anchor_embs, p=2, dim=-1)
    positive_embs = F.normalize(positive_embs, p=2, dim=-1)
    negative_embs = F.normalize(negative_embs, p=2, dim=-1)
    semi_positive_embs = F.normalize(semi_positive_embs, p=2, dim=-1)

    # Calculate cosine similarities
    sim_anchor_pos = torch.sum(anchor_embs * positive_embs, dim=-1)  # Similarity: anchor to positive
    sim_anchor_neg = torch.sum(anchor_embs * negative_embs, dim=-1)  # Similarity: anchor to negative
    sim_anchor_semi = torch.sum(anchor_embs * semi_positive_embs, dim=-1)  # Similarity: anchor to semi-positive

    # Compute positive loss (maximize similarity)
    # Constraint: Positive pairs should aim for maximum similarity (close to 1)
    loss_pos = (1.0 - sim_anchor_pos).clamp(min=0).mean()  # Penalize if below maximum similarity
    # Compute negative loss (minimize similarity)
    # Constraint: Negative pairs should remain lower than a margin (e.g., 0.5)
    loss_neg = (sim_anchor_neg - 0.3).clamp(min=0).mean()  # Penalize negatives close to anchor
    # Constraint: Semi-positive pairs should have intermediate similarity
    loss_semi = ((sim_anchor_semi - 0.5).abs()).mean()  # Keep semi-positives around 0.5

    # Combine losses
    total_loss = loss_pos + loss_neg + loss_semi

    return total_loss

def train_encoder(model, optimizer, dl_train, max_epochs, save_path, logger):
    smallest_train_loss = float('inf')
    best_model_state = None
    losses = []
    for epoch in range(max_epochs):
        model.train()
        epoch_loss = 0.0
        for batch in dl_train:
            anchor, positive, semi_positive, negative = batch['anchor'], batch['positive'], batch['semi_positive'], batch['negative']
            anchor_features, positive_features, semi_positive_features, negative_features = anchor[0], positive[0], semi_positive[0], negative[0]
            anchor_features, positive_features, semi_positive_features, negative_features = anchor_features.to(device), positive_features.to(device), semi_positive_features.to(device), negative_features.to(device)

            # logger.info("anchor_features shape in encoder", anchor_features.shape)

            anchor_embed = model(anchor_features)
            positive_embed = model(positive_features)
            negative_embed = model(negative_features)
            semi_positive_embed = model(semi_positive_features)

            anchor_embed, positive_embed, negative_embed, semi_positive_embed = anchor_embed.to(device), positive_embed.to(device), negative_embed.to(device), semi_positive_embed.to(device)
            # logger.info("embeddings shape in train function", anchor_embed.shape, positive_embed.shape, negative_embed.shape, semi_positive_embed.shape)

            # Compute BPR-inspired contrastive loss
            loss = contrastive_loss(anchor_embed, positive_embed, negative_embed, semi_positive_embed)
            epoch_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        avg_epoch_loss = epoch_loss / len(dl_train)
        losses.append(avg_epoch_loss)
        logger.info(f'Epoch [{epoch + 1}/{max_epochs}], Training Loss: {avg_epoch_loss:.4f}')

        # Save the best model based on the lowest training loss
        if avg_epoch_loss < smallest_train_loss:
            smallest_train_loss = avg_epoch_loss
            best_model_state = copy.deepcopy(model.state_dict())

        nowtime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        logger.info(f"Finished epoch {epoch + 1}/{max_epochs} at {nowtime}")

    logger.info('Finished Training...')
    time_finished = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
    model_path = os.path.join(save_path, f"encoder_{time_finished}.pth")
    torch.save(best_model_state, model_path)
    model.load_state_dict(best_model_state)
    logger.info(f'Model saved to {model_path}')
    
    plt.clf()
    plt.plot(range(1, max_epochs + 1), losses)
    plt.xlabel('Epoch')
    plt.ylabel('Training Loss')
    plt.title('Training Loss Over Epochs')
    file_name = 'loss_encoder.png'
    full_save_path = os.path.join(save_path, file_name)
    plt.savefig(full_save_path)
    plt.close()

    return model

def combine_embeddings(embeddings_folder, output_file, logger):
    # List all .pt files in the embeddings folder and sort them.
    pt_files = sorted(glob.glob(os.path.join(embeddings_folder, '*.pt')))
    all_embeddings = []
    for file in pt_files:
        emb = torch.load(file)
        all_embeddings.append(emb)
    # Concatenate all embeddings along the batch dimension (assumed to be 0).
    combined_embeddings = torch.cat(all_embeddings, dim=0)
    torch.save(combined_embeddings, output_file)
    logger.info(f"Combined embeddings saved to {output_file}")

def combine_samples(samples_folder, output_file, logger):
    # List all pickle (.pkl) files in the samples folder and sort them.
    pkl_files = sorted(glob.glob(os.path.join(samples_folder, '*.pkl')))
    all_samples = []
    for file in pkl_files:
        with open(file, 'rb') as f:
            samples = pickle.load(f)
            # Extend the list if each file contains a list of samples.
            all_samples.extend(samples)
    with open(output_file, 'wb') as f:
        pickle.dump(all_samples, f)
    logger.info(f"Combined samples saved to {output_file}")

def save_precomputed_samples(encoder, dl_train, save_path, logger):
    encoder.eval()

    os.makedirs(os.path.join(save_path, 'precomputed_embedding_epi'), exist_ok=True)
    os.makedirs(os.path.join(save_path, 'precomputed_embedding_hyp'), exist_ok=True)
    os.makedirs(os.path.join(save_path, 'precomputed_samples_epi'), exist_ok=True)
    os.makedirs(os.path.join(save_path, 'precomputed_samples_hyp'), exist_ok=True)

    count_epi = 0
    count_hyp = 0

    logger.info("Start saving precomputed samples and embeddings...")
    for batch in dl_train:
        precomputed_samples_epi = []
        precomputed_samples_hyp = []
        anchor, positive, semi_positive, negative = batch['anchor'], batch['positive'], batch['semi_positive'], batch['negative']
        anchor_features, positive_features, semi_positive_features, negative_features = anchor[0], positive[0], semi_positive[0], negative[0]
        anchor_features, positive_features, semi_positive_features, negative_features = anchor_features.to(device), positive_features.to(device), semi_positive_features.to(device), negative_features.to(device)
        anchor_obs_epi, anchor_obs_hyp, anchor_lake_id, anchor_year, anchor_month = anchor[1], anchor[2], anchor[3], anchor[4], anchor[5]

        mask1 = (anchor_obs_epi != -11)
        mask1 = mask1.to(device)
        mask2 = (anchor_obs_hyp != -11)
        mask2 = mask2.to(device)
        if (torch.sum(mask1) == 0 and torch.sum(mask2) == 0):
            continue

        with torch.no_grad():
            embedding = encoder(anchor_features).detach().cpu()  # Compute embedding

        # only save the features of those samples with valid obs_epi and obs_hyp
        if torch.sum(mask1) > 0:
            file_embed_name = f'precomputed_embedding_epi/precomputed_embeddings_epi_{count_epi}.pt'
            full_embed_save_path = os.path.join(save_path, file_embed_name)
            torch.save(embedding, full_embed_save_path)
            logger.info(f'Precomputed embeddings saved to {full_embed_save_path}')

            for i in range(len(anchor_obs_epi)):  # Loop through the batch
                precomputed_samples_epi.append({
                    'obs': anchor_obs_epi[i].numpy(),  # Save as numpy array
                    'features': anchor_features[i].detach().cpu().numpy(),  # Include original features if needed
                    'anchor_lake_id': anchor_lake_id[i],
                    'anchor_year': anchor_year[i],
                    'anchor_month': anchor_month[i]
                })

            file_sample_name = f'precomputed_samples_epi/precomputed_samples_epi_{count_epi}.pkl'
            full_sample_save_path = os.path.join(save_path, file_sample_name)
            with open(full_sample_save_path, 'wb') as f:
                pickle.dump(precomputed_samples_epi, f)
            logger.info(f'Precomputed samples saved to {full_sample_save_path}')

            count_epi += 1
        
        if torch.sum(mask2) > 0:
            file_embed_name = f'precomputed_embedding_hyp/precomputed_embeddings_hyp_{count_hyp}.pt'
            full_embed_save_path = os.path.join(save_path, file_embed_name)
            torch.save(embedding, full_embed_save_path)
            logger.info(f'Precomputed embeddings saved to {full_embed_save_path}')

            for i in range(len(anchor_obs_hyp)):
                precomputed_samples_hyp.append({
                    'obs': anchor_obs_hyp[i].numpy(),
                    'features': anchor_features[i].detach().cpu().numpy(),
                    'anchor_lake_id': anchor_lake_id[i],
                    'anchor_year': anchor_year[i],
                    'anchor_month': anchor_month[i]
                })
            
            file_sample_name = f'precomputed_samples_hyp/precomputed_samples_hyp_{count_hyp}.pkl'
            full_sample_save_path = os.path.join(save_path, file_sample_name)
            with open(full_sample_save_path, 'wb') as f:
                pickle.dump(precomputed_samples_hyp, f)
            logger.info(f'Precomputed samples saved to {full_sample_save_path}')

            count_hyp += 1
        
    epi_embeddings_folder = os.path.join(save_path, 'precomputed_embedding_epi')
    combined_epi_embed_file = os.path.join(save_path, 'precomputed_embeddings_epi.pt')
    combine_embeddings(epi_embeddings_folder, combined_epi_embed_file, logger)

    hyp_embeddings_folder = os.path.join(save_path, 'precomputed_embedding_hyp')
    combined_hyp_embed_file = os.path.join(save_path, 'precomputed_embeddings_hyp.pt')
    combine_embeddings(hyp_embeddings_folder, combined_hyp_embed_file, logger)

    # Combine epi samples
    epi_samples_folder = os.path.join(save_path, 'precomputed_samples_epi')
    combined_epi_samples_file = os.path.join(save_path, 'precomputed_samples_epi.pkl')
    combine_samples(epi_samples_folder, combined_epi_samples_file, logger)

    # Combine hyp samples
    hyp_samples_folder = os.path.join(save_path, 'precomputed_samples_hyp')
    combined_hyp_samples_file = os.path.join(save_path, 'precomputed_samples_hyp.pkl')
    combine_samples(hyp_samples_folder, combined_hyp_samples_file, logger)

# ========== Train yearly model ==========
def pretrain_yearly_model_use_sim(yearly_model, optimizer, dl_train, dl_val, year_epochs, save_path, logger):
    smallest_val_loss = float('inf')
    best_model_state = None

    train_losses = []
    val_losses = []

    for epoch in range(year_epochs):
        train_epi_pred = []
        train_epi_label = []
        train_hyp_pred = []
        train_hyp_label = []

        yearly_model.train()
        for x, y_epi, y_hyp, lake_id, year in dl_train:
            y_epi = x[:, :, -2]
            y_hyp = x[:, :, -1]
            x, y_epi, y_hyp, year = x[:, :, :-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)  # remove the simulated labels in input features

            y_epi = y_epi.squeeze(-1)
            y_hyp = y_hyp.squeeze(-1)

            pred_epi, pred_hyp = yearly_model(x)

            loss_epi = RMSE(pred_epi, y_epi)
            loss_hyp = RMSE(pred_hyp, y_hyp)
            combined_loss = (loss_epi + loss_hyp) / 2

            optimizer.zero_grad()
            combined_loss.backward()
            optimizer.step()

            train_epi_pred.append(pred_epi.squeeze().cpu().detach())
            train_epi_label.append(y_epi.squeeze().cpu().detach())
            train_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
            train_hyp_label.append(y_hyp.squeeze().cpu().detach())

        train_epi_pred = torch.cat(train_epi_pred, dim=0)
        train_epi_label = torch.cat(train_epi_label, dim=0)
        train_hyp_pred = torch.cat(train_hyp_pred, dim=0)
        train_hyp_label = torch.cat(train_hyp_label, dim=0)

        train_epi_RMSE = RMSE(train_epi_pred.to(device), train_epi_label.to(device))
        train_hyp_RMSE = RMSE(train_hyp_pred.to(device), train_hyp_label.to(device))

        train_losses.append(train_epi_RMSE.cpu().numpy())

        logger.info(f"Epoch [{epoch + 1}/{year_epochs}], Loss Epi: {train_epi_RMSE}, Loss Hyp: {train_hyp_RMSE}")

        val_epi_pred = []
        val_epi_label = []
        val_hyp_pred = []
        val_hyp_label = []

        yearly_model.eval()
        with torch.no_grad():
            for x, y_epi, y_hyp, lake_id, year in dl_val:
                y_epi = x[:, :, -2]
                y_hyp = x[:, :, -1]
                x, y_epi, y_hyp, year = x[:, :, :-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)
                y_epi = y_epi.squeeze(-1)
                y_hyp = y_hyp.squeeze(-1)

                # Decode
                pred_epi, pred_hyp = yearly_model(x)

                val_epi_pred.append(pred_epi.squeeze().cpu().detach())
                val_epi_label.append(y_epi.squeeze().cpu().detach())
                val_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
                val_hyp_label.append(y_hyp.squeeze().cpu().detach())

        val_epi_pred = torch.cat(val_epi_pred, dim=0)
        val_epi_label = torch.cat(val_epi_label, dim=0)
        val_hyp_pred = torch.cat(val_hyp_pred, dim=0)
        val_hyp_label = torch.cat(val_hyp_label, dim=0)


        val_epi_RMSE = RMSE(val_epi_pred.to(device), val_epi_label.to(device))
        val_hyp_RMSE = RMSE(val_hyp_pred.to(device), val_hyp_label.to(device))
        
        logger.info(f'Val RMSE Epi: {val_epi_RMSE:.4f}')
        logger.info(f'Val RMSE Hyp: {val_hyp_RMSE:.4f}')
        val_losses.append(val_epi_RMSE.cpu().numpy())

        # Save the best model based on the lowest training loss
        if val_epi_RMSE < smallest_val_loss:
            smallest_val_loss = val_epi_RMSE
            best_model_state = copy.deepcopy(yearly_model.state_dict())

        nowtime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        logger.info("\n" + "==========" * 8 + f" {nowtime}")

    logger.info('Finished Training...')
    time_finished = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
    model_path = os.path.join(save_path, f"pretrain_yearly_model_{time_finished}.pth")
    torch.save(best_model_state, model_path)
    yearly_model.load_state_dict(best_model_state)
    logger.info(f'Model saved to {model_path}')

    plt.clf()
    # Plot training and val loss
    plt.figure(figsize=(8, 6))
    plt.plot(range(1, year_epochs + 1), train_losses, label='Train Loss')
    plt.plot(range(1, year_epochs + 1), val_losses, label='Val Loss', linestyle='--')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training and Val Loss Over Epochs')
    plt.legend()
    file_name = 'loss_yearly_model_pretrain_val.png'
    full_save_path = os.path.join(save_path, file_name)
    plt.savefig(full_save_path)
    plt.close()

    return yearly_model

def train_yearly_model(yearly_model, optimizer, dl_train, dl_val, year_epochs, save_path, logger):
    smallest_val_loss = float('inf')
    best_model_state = None
    
    train_losses = []
    val_losses = []

    for epoch in range(year_epochs):
        train_epi_pred = []
        train_epi_label = []
        train_hyp_pred = []
        train_hyp_label = []
        train_sim_epi = []
        train_sim_hyp = []

        yearly_model.train()
        for x, y_epi, y_hyp, lake_id, year in dl_train:
            sim_epi = x[:, :, -2].to(device)
            sim_hyp = x[:, :, -1].to(device)
            x, y_epi, y_hyp, year = x[:,:,:-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device) # remove the simulated labels in input features

            y_epi = y_epi.squeeze(-1)
            y_hyp = y_hyp.squeeze(-1)

            mask1 = (y_epi != -11)
            mask1 = mask1.to(device)
            mask2 = (y_hyp != -11)
            mask2 = mask2.to(device)
            if (torch.sum(mask1) == 0) or (torch.sum(mask2) == 0):
                continue

            pred_epi, pred_hyp = yearly_model(x)

            loss_epi = RMSE(pred_epi, y_epi)
            loss_hyp = RMSE(pred_hyp, y_hyp)
            combined_loss = (loss_epi + loss_hyp) / 2

            optimizer.zero_grad()
            combined_loss.backward()
            optimizer.step()

            train_epi_pred.append(pred_epi.squeeze().cpu().detach())
            train_epi_label.append(y_epi.squeeze().cpu().detach())
            train_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
            train_hyp_label.append(y_hyp.squeeze().cpu().detach())  
            train_sim_epi.append(sim_epi.squeeze().cpu().detach())
            train_sim_hyp.append(sim_hyp.squeeze().cpu().detach())    


        train_epi_pred = torch.cat(train_epi_pred, dim=0)
        train_epi_label = torch.cat(train_epi_label, dim=0)
        train_hyp_pred = torch.cat(train_hyp_pred, dim=0)
        train_hyp_label = torch.cat(train_hyp_label, dim=0)
        train_sim_epi = torch.cat(train_sim_epi, dim=0)
        train_sim_hyp = torch.cat(train_sim_hyp, dim=0)

        train_epi_RMSE = RMSE(train_epi_pred.to(device), train_epi_label.to(device))
        train_hyp_RMSE = RMSE(train_hyp_pred.to(device), train_hyp_label.to(device))
        
        train_losses.append(train_epi_RMSE.cpu().numpy())

        logger.info(f"Epoch [{epoch + 1}/{year_epochs}], Loss Epi: {train_epi_RMSE}, Loss Hyp: {train_hyp_RMSE}")

        val_epi_pred = []
        val_epi_label = []
        val_hyp_pred = []
        val_hyp_label = []
        val_sim_epi = []
        val_sim_hyp = []

        yearly_model.eval()
        with torch.no_grad():
            for x, y_epi, y_hyp, lake_id, year in dl_val:
                sim_epi = x[:, :, -2].to(device)
                sim_hyp = x[:, :, -1].to(device)
                x, y_epi, y_hyp, year = x[:,:,:-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)
                y_epi = y_epi.squeeze(-1)
                y_hyp = y_hyp.squeeze(-1)

                # Decode
                pred_epi, pred_hyp = yearly_model(x)

                val_epi_pred.append(pred_epi.squeeze().cpu().detach())
                val_epi_label.append(y_epi.squeeze().cpu().detach())
                val_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
                val_hyp_label.append(y_hyp.squeeze().cpu().detach())
                val_sim_epi.append(sim_epi.squeeze().cpu().detach())
                val_sim_hyp.append(sim_hyp.squeeze().cpu().detach())

        val_epi_pred = torch.cat(val_epi_pred, dim=0)
        val_epi_label = torch.cat(val_epi_label, dim=0)
        val_hyp_pred = torch.cat(val_hyp_pred, dim=0)
        val_hyp_label = torch.cat(val_hyp_label, dim=0)
        val_sim_epi = torch.cat(val_sim_epi, dim=0)
        val_sim_hyp = torch.cat(val_sim_hyp, dim=0)

        val_epi_RMSE = RMSE(val_epi_pred.to(device), val_epi_label.to(device))
        val_hyp_RMSE = RMSE(val_hyp_pred.to(device), val_hyp_label.to(device))
        
        logger.info(f'Val RMSE Epi: {val_epi_RMSE:.4f}')
        logger.info(f'Val RMSE Hyp: {val_hyp_RMSE:.4f}')
        val_losses.append(val_epi_RMSE.cpu().numpy())

        # Save the best model based on the lowest training loss
        if val_epi_RMSE < smallest_val_loss:
            smallest_val_loss = val_epi_RMSE
            best_model_state = copy.deepcopy(yearly_model.state_dict())

        nowtime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        logger.info("\n" + "==========" * 8 + f" {nowtime}")

    logger.info('Finished Training...')
    time_finished = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
    model_path = os.path.join(save_path, f"yearly_model_{time_finished}.pth")
    torch.save(best_model_state, model_path)
    yearly_model.load_state_dict(best_model_state)
    logger.info(f'Model saved to {model_path}')
    
    plt.clf()
    # Plot training and val loss
    plt.figure(figsize=(8, 6))
    plt.plot(range(1, year_epochs + 1), train_losses, label='Train Loss')
    plt.plot(range(1, year_epochs + 1), val_losses, label='Val Loss', linestyle='--')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training and Val Loss Over Epochs')
    plt.legend()
    file_name = 'loss_yearly_model_train_val.png'
    full_save_path = os.path.join(save_path, file_name)
    plt.savefig(full_save_path)
    plt.close()

    return yearly_model

def test_yearly_model(yearly_model, dl_test, save_path, logger):
    test_epi_pred = []
    test_epi_label = []
    test_hyp_pred = []
    test_hyp_label = []
    test_sim_epi = []
    test_sim_hyp = []

    lake_ids = []
    years = []

    yearly_model.eval()
    with torch.no_grad():
        for x, y_epi, y_hyp, lake_id, year in dl_test:
            sim_epi = x[:, :, -2].to(device)
            sim_hyp = x[:, :, -1].to(device)
            x, y_epi, y_hyp, year = x[:,:,:-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)
            y_epi = y_epi.squeeze(-1)
            y_hyp = y_hyp.squeeze(-1)

            mask1 = (y_epi != -11)
            mask1 = mask1.to(device)
            mask2 = (y_hyp != -11)
            mask2 = mask2.to(device)
            if (torch.sum(mask1) == 0) and (torch.sum(mask2) == 0):
                continue

            # Decode
            pred_epi, pred_hyp = yearly_model(x)

            test_epi_pred.append(pred_epi.squeeze().cpu().detach())
            test_epi_label.append(y_epi.squeeze().cpu().detach())
            test_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
            test_hyp_label.append(y_hyp.squeeze().cpu().detach())
            test_sim_epi.append(sim_epi.squeeze().cpu().detach())
            test_sim_hyp.append(sim_hyp.squeeze().cpu().detach())

            lake_ids.append(lake_id[0])
            years.append(year.cpu().item())

    test_epi_pred = torch.cat(test_epi_pred, dim=0)
    test_epi_label = torch.cat(test_epi_label, dim=0)
    test_hyp_pred = torch.cat(test_hyp_pred, dim=0)
    test_hyp_label = torch.cat(test_hyp_label, dim=0)
    test_sim_epi = torch.cat(test_sim_epi, dim=0)
    test_sim_hyp = torch.cat(test_sim_hyp, dim=0)

    test_epi_RMSE = RMSE(test_epi_pred.to(device), test_epi_label.to(device))
    test_hyp_RMSE = RMSE(test_hyp_pred.to(device), test_hyp_label.to(device))

    logger.info(f'Test RMSE Epi: {test_epi_RMSE:.4f}')
    logger.info(f'Test RMSE Hyp: {test_hyp_RMSE:.4f}')

    days_per_year = 360
    lake_ids_expanded = np.repeat(lake_ids, days_per_year)
    years_expanded = np.repeat(years, days_per_year)

    df_epi = pd.DataFrame({
        "Lake_ID": lake_ids_expanded,
        "Year": years_expanded,
        "Day": np.tile(np.arange(1, days_per_year + 1), len(lake_ids)),  # Add Day index (1 to 360)
        "True_Label_Epi": test_epi_label.cpu().numpy(),
        "Prediction_Epi": test_epi_pred.cpu().numpy(),
        "Sim_Epi": test_sim_epi.cpu().numpy(),
    })

    # Create a DataFrame for Hyp
    df_hyp = pd.DataFrame({
        "Lake_ID": lake_ids_expanded,
        "Year": years_expanded,
        "Day": np.tile(np.arange(1, days_per_year + 1), len(lake_ids)),  # Add Day index (1 to 360)
        "True_Label_Hyp": test_hyp_label.cpu().numpy(),
        "Prediction_Hyp": test_hyp_pred.cpu().numpy(),
        "Sim_Hyp": test_sim_hyp.cpu().numpy(),
    })

    # Save DataFrames to CSV files
    df_epi.to_csv(os.path.join(save_path, "epi_predictions_test.csv"), index=False)
    df_hyp.to_csv(os.path.join(save_path, "hyp_predictions_test.csv"), index=False)

# ========== Retrieve and prepare dataloaders ==========
def retrieve_chunks(bad_month_embedding, precomputed_embeddings, precomputed_samples, bad_month_sample, top_n, skip_anchor=True, target_type='epi', logger=None):
    # Get initial similarity-based candidates
    similarities = F.cosine_similarity(precomputed_embeddings, bad_month_embedding)
    sorted_indices = torch.argsort(similarities, descending=True)
    
    initial_candidates = []
    anchor_sample = None
    anchor_similarity = None
    
    # Aggressive early filtering: only consider top candidates with similarity > threshold
    similarity_threshold = 0.3  # Only consider reasonably similar samples
    search_limit = min(len(precomputed_samples), top_n*3)  # Cap at reasonable limit
        
    for idx in sorted_indices[:search_limit]:
        if len(initial_candidates) >= top_n:
            break  # Early termination when we have enough candidates
            
        sample = precomputed_samples[idx.item()]
        similarity = similarities[idx].item()

        # Validate sample has observations
        has_valid_obs = torch.sum(torch.tensor(sample['obs']) != -11) > 0
        if not has_valid_obs:
            continue
        
        # Skip samples with very low similarity (aggressive filtering)
        if similarity < similarity_threshold:
            # logger.info(f"Stopping search at similarity {similarity:.4f} (below threshold {similarity_threshold})")
            break
        
        # Check if this is the anchor sample
        is_anchor = (str(sample['anchor_lake_id']) == str(bad_month_sample[2]) and
                    sample['anchor_year'] == bad_month_sample[3] and
                    sample['anchor_month'] == bad_month_sample[4])
            
        # Handle anchor sample for joint training
        if is_anchor and not skip_anchor and anchor_sample is None:
            anchor_sample = sample
            anchor_similarity = similarity
            # logger.info(f"Found anchor sample for joint training with similarity {anchor_similarity:.4f}")
            continue  # Don't add to initial_candidates yet
            
        # Skip anchor if we're in normal retrieval mode or already have it
        if is_anchor and (skip_anchor or anchor_sample is not None):
            continue
            
        # Add valid candidate
        initial_candidates.append((sample, similarity))
    
    # Add anchor sample to the beginning if we have it (for joint training)
    if anchor_sample is not None:
        initial_candidates.insert(0, (anchor_sample, anchor_similarity))
    
    logger.info(f"Initial candidates: {len(initial_candidates)}")
    
    # Fallback: if no similarity-based candidates found, get any valid samples
    if len(initial_candidates) == 0:
        logger.info(f"No similarity-based candidates found, using fallback retrieval...")
        
        # Fallback: get top samples regardless of similarity, but still validate observations
        fallback_limit = min(len(precomputed_samples), top_n * 2)  # Search more broadly
        for idx in sorted_indices[:fallback_limit]:
            sample = precomputed_samples[idx.item()]
            similarity = similarities[idx].item()

            # Validate sample has observations
            has_valid_obs = torch.sum(torch.tensor(sample['obs']) != -11) > 0
            if not has_valid_obs:
                continue
            
            # Skip anchor sample in fallback mode too
            is_anchor = (str(sample['anchor_lake_id']) == str(bad_month_sample[2]) and
                        sample['anchor_year'] == bad_month_sample[3] and
                        sample['anchor_month'] == bad_month_sample[4])
                        
            if is_anchor and skip_anchor:
                continue
                
            initial_candidates.append((sample, similarity))
            if len(initial_candidates) >= 5:
                break
        
        logger.info(f"Fallback retrieval found {len(initial_candidates)} candidates")
        
        # If still no candidates after fallback, return empty
        if len(initial_candidates) == 0:
            logger.info(f"No valid candidates found even with fallback")
            return []
        
    return initial_candidates

def enhanced_retrieve_and_prepare_dataloaders(bad_month_embedding, precomputed_embeddings, precomputed_samples, bad_month_sample, top_n, for_joint_training=False, target_type='epi', logger=None):
    
    # Get query lake information
    query_lake_id = str(bad_month_sample[2])  # Convert to string for consistency
      
    # Use the reranking function
    retrieved_samples = retrieve_chunks(
        bad_month_embedding,
        precomputed_embeddings, 
        precomputed_samples,
        bad_month_sample,
        top_n,
        skip_anchor=not for_joint_training,  # Don't skip for joint training
        target_type=target_type,
        logger=logger
    )
            
    # Check if no samples were retrieved
    if len(retrieved_samples) == 0:
        logger.info(f"No candidates found for query lake {query_lake_id} - returning None dataloaders")
        return None, None
        
    # Create training samples
    train_samples = retrieved_samples

    train_y = torch.stack([torch.tensor(sample[0]['obs'], dtype=torch.float32).to(device) for sample in train_samples])
    train_x = torch.stack([torch.tensor(sample[0]['features'], dtype=torch.float32).to(device) for sample in train_samples])
    train_similarities = torch.tensor([sample[1] for sample in train_samples], dtype=torch.float32).to(device)
    train_lake_id = [sample[0]['anchor_lake_id'] for sample in train_samples]
    train_year = torch.tensor([sample[0]['anchor_year'] for sample in train_samples], dtype=torch.float32).to(device)
    train_month = torch.tensor([sample[0]['anchor_month'] for sample in train_samples], dtype=torch.float32).to(device)
    
    bad_month_y = torch.stack([torch.tensor(bad_month_sample[1], dtype=torch.float32).to(device)])
    bad_month_x = torch.stack([torch.tensor(bad_month_sample[0], dtype=torch.float32).to(device)])
    bad_month_similarity = torch.tensor([1.0], dtype=torch.float32).to(device)
    bad_month_lake_id = [str(bad_month_sample[2])]
    bad_month_year = torch.tensor([int(bad_month_sample[3])], dtype=torch.float32).to(device)
    bad_month_month = torch.tensor([int(bad_month_sample[4])], dtype=torch.float32).to(device)
    
    train_dataset = AbnormalMonthDataset(train_x, train_y, train_similarities, train_lake_id, train_year, train_month)
    test_dataset = AbnormalMonthDataset(bad_month_x, bad_month_y, bad_month_similarity, bad_month_lake_id, bad_month_year, bad_month_month)
    
    dl_month_train = DataLoader(train_dataset, batch_size=4, shuffle=False)
    dl_month_test = DataLoader(test_dataset, batch_size=4, shuffle=False)
    
    return dl_month_train, dl_month_test

def RMSE(yhat, y):
    yhat = yhat.flatten()
    y = y.flatten()

    mask = (y != -11)

    squared_errors = torch.square((yhat - y) * mask)

    mse = torch.sum(squared_errors) / torch.sum(mask)

    rmse = torch.sqrt(mse)

    return rmse

# ============= Pretrain Decoder ========
def pretrain_decoder(decoder, dl_decoder_train, dl_decoder_test, flag, logger):
    total_loss = 0.0
    for dl in [dl_decoder_train, dl_decoder_test]:
        for x, _, sim, lake_id, year, month in dl:
            if flag == 'epi':
                y = x[:, :, -2]
            else:
                y = x[:, :, -1]

            mask = (y != -11)
            mask = mask.to(device)
            if torch.sum(mask) == 0:
                continue

            x, y, sim = x[:, :, :-2].to(device), y.to(device), sim.to(device)

            x = x[:, :, :26]
        
            pred = decoder(x)
            
            loss = RMSE(pred, y)
            weighted_loss = (loss * sim).mean() # Weighted loss based on similarity
            total_loss += weighted_loss.item()
    
    return total_loss

def jointly_train_encoder_decoder(encoder, decoder_epi, decoder_hyp, optimizer, dl_encoder_decoder_train, max_epochs, top_n, a, b, c, save_path, precomputed_embeddings_epi, precomputed_samples_epi, precomputed_embeddings_hyp, precomputed_samples_hyp, logger):
    loss_history = {'encoder_loss': [], 'decoder_loss_epi': [], 'decoder_loss_hyp': [], 'total_loss': []}

    encoder.train()
    decoder_epi.train()
    decoder_hyp.train()

    for epoch in range(max_epochs):
        logger.info(f"Epoch of jointly training {epoch + 1}/{max_epochs}")
        epoch_loss = {'encoder_loss': 0.0, 'decoder_loss_epi': 0.0, 'decoder_loss_hyp': 0.0, 'total_loss': 0.0}
        num_batches = 0
        epi_updates = 0  # Counter for decoder_loss_epi updates
        hyp_updates = 0  # Counter for decoder_loss_hyp updates

        for batch in dl_encoder_decoder_train:
            num_batches += 1

            anchor, positive, semi_positive, negative = batch['anchor'], batch['positive'], batch['semi_positive'], batch['negative']

            anchor_features, positive_features, semi_positive_features, negative_features = anchor[0], positive[0], semi_positive[0], negative[0]
            anchor_features, positive_features, semi_positive_features, negative_features = anchor_features.to(device), positive_features.to(device), semi_positive_features.to(device), negative_features.to(device)

            anchor_obs_epi, anchor_obs_hyp, anchor_lake_id, anchor_year, anchor_month = anchor[1], anchor[2], anchor[3], anchor[4], anchor[5]
            anchor_obs_epi, anchor_obs_hyp = anchor_obs_epi.to(device), anchor_obs_hyp.to(device)

            # logger.info("anchor_features shape in encoder", anchor_features.shape)

            anchor_embed = encoder(anchor_features)
            positive_embed = encoder(positive_features)
            negative_embed = encoder(negative_features)
            semi_positive_embed = encoder(semi_positive_features)

            anchor_embed, positive_embed, negative_embed, semi_positive_embed = anchor_embed.to(device), positive_embed.to(device), negative_embed.to(device), semi_positive_embed.to(device)
            # logger.info("embeddings shape in train function", anchor_embed.shape, positive_embed.shape, negative_embed.shape, semi_positive_embed.shape)

            encoder_loss = contrastive_loss(anchor_embed, positive_embed, negative_embed, semi_positive_embed)

            # Mask to check if there are valid values for the month
            mask_epi = (anchor_obs_epi != -11).to(device)
            mask_hyp = (anchor_obs_hyp != -11).to(device)
            if torch.sum(mask_epi) == 0 and torch.sum(mask_hyp) == 0:
                continue

            bad_month_sample_epi = (anchor_features.squeeze(0).to(device), anchor_obs_epi.squeeze(0).to(device), anchor_lake_id[0], anchor_year.item(), anchor_month.item())
            bad_month_sample_hyp = (anchor_features.squeeze(0).to(device), anchor_obs_hyp.squeeze(0).to(device), anchor_lake_id[0], anchor_year.item(), anchor_month.item())

            decoder_loss_epi = 0
            decoder_loss_hyp = 0

            if torch.sum(mask_epi) > 0:
                # here the test set is the bad month and it will also be used as training data
                dl_decoder_train_epi, dl_decoder_test_epi = enhanced_retrieve_and_prepare_dataloaders(anchor_embed, precomputed_embeddings_epi, precomputed_samples_epi, bad_month_sample_epi, top_n, for_joint_training=True, target_type='epi', logger=logger)
                if dl_decoder_train_epi is None or dl_decoder_test_epi is None:
                    logger.info("No valid samples found for epi decoder training")
                    continue
                decoder_loss_epi = pretrain_decoder(decoder_epi, dl_decoder_train_epi, dl_decoder_test_epi, flag='epi', logger=logger)
                # decoder_epi_rmse = RMSE(pred_epi_monthly, anchor_obs_epi)
                # yearly_model_rmse_epi = RMSE(pred_epi, anchor_obs_epi)
                epi_updates += 1

            if torch.sum(mask_hyp) > 0:
                # here the test set is the bad month and it will also be used as training data
                dl_decoder_train_hyp, dl_decoder_test_hyp = enhanced_retrieve_and_prepare_dataloaders(anchor_embed, precomputed_embeddings_hyp, precomputed_samples_hyp, bad_month_sample_hyp, top_n, for_joint_training=True, target_type='hyp', logger=logger)
                if dl_decoder_train_hyp is None or dl_decoder_test_hyp is None:
                    logger.info("No valid samples found for hyp decoder training")
                    continue
                decoder_loss_hyp = pretrain_decoder(decoder_hyp, dl_decoder_train_hyp, dl_decoder_test_hyp, flag='hyp', logger=logger)
                # decoder_epi_rmse = RMSE(pred_hyp_monthly, anchor_obs_hyp)
                # yearly_model_rmse_hyp = RMSE(pred_hyp, anchor_obs_hyp)
                hyp_updates += 1  # Increment update counter for hyp decoder    

            all_loss = a * encoder_loss + b * decoder_loss_epi + c * decoder_loss_hyp
            if num_batches % 10000 == 0:
                logger.info("all_loss % 10000 = 0: ", all_loss)

            optimizer.zero_grad()
            all_loss.backward()
            optimizer.step()

            # Accumulate losses for epoch
            epoch_loss['encoder_loss'] += encoder_loss.item()
            epoch_loss['decoder_loss_epi'] += decoder_loss_epi
            epoch_loss['decoder_loss_hyp'] += decoder_loss_hyp
            epoch_loss['total_loss'] += all_loss.item()
        
        epoch_loss['encoder_loss'] /= num_batches
        epoch_loss['decoder_loss_epi'] = (epoch_loss['decoder_loss_epi'] / epi_updates) if epi_updates > 0 else 0.0
        epoch_loss['decoder_loss_hyp'] = (epoch_loss['decoder_loss_hyp'] / hyp_updates) if hyp_updates > 0 else 0.0
        epoch_loss['total_loss'] /= num_batches
        
        # Record epoch losses
        loss_history['encoder_loss'].append(epoch_loss['encoder_loss'])
        loss_history['decoder_loss_epi'].append(epoch_loss['decoder_loss_epi'])
        loss_history['decoder_loss_hyp'].append(epoch_loss['decoder_loss_hyp'])
        loss_history['total_loss'].append(epoch_loss['total_loss'])

        # logger.info epoch summary
        logger.info(f"Epoch {epoch + 1}/{max_epochs} - Encoder Loss: {epoch_loss['encoder_loss']:.4f}, "
              f"Decoder Loss Epi: {epoch_loss['decoder_loss_epi']:.4f}, "
              f"Decoder Loss Hyp: {epoch_loss['decoder_loss_hyp']:.4f}, "
              f"Total Loss: {epoch_loss['total_loss']:.4f}")

    # Save model states
    torch.save(encoder.state_dict(), os.path.join(save_path, "finetuned_encoder.pth"))
    torch.save(decoder_epi.state_dict(), os.path.join(save_path, "pretrained_decoder_epi.pth"))
    torch.save(decoder_hyp.state_dict(), os.path.join(save_path, "pretrained_decoder_hyp.pth"))
    
    epochs = range(1, len(loss_history['encoder_loss']) + 1)
    encoder_loss = loss_history['encoder_loss']
    decoder_loss_epi = loss_history['decoder_loss_epi']
    decoder_loss_hyp = loss_history['decoder_loss_hyp']
    total_loss = loss_history['total_loss']

    # Create the plot
    plt.clf()

    plt.figure(figsize=(10, 6))
    plt.plot(epochs, total_loss, label='Encoder&Decoder Loss', marker='o', color='pink')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Encoder Loss Over Epochs')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_path, 'encoder_decoder_loss.png'))
    plt.close()

    plt.figure(figsize=(10, 6))
    plt.plot(epochs, encoder_loss, label='Encoder Loss', marker='o', color='blue')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Encoder Loss Over Epochs')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_path, 'encoder_loss.png'))
    plt.close()

    # Decoder loss plot
    plt.figure(figsize=(10, 6))
    plt.plot(epochs, decoder_loss_epi, label='Decoder Loss (Epi)', linestyle='--', marker='s', color='orange')
    plt.plot(epochs, decoder_loss_hyp, label='Decoder Loss (Hyp)', linestyle='-.', marker='^', color='green')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Decoder Losses (Epi and Hyp) Over Epochs')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(save_path, 'decoder_losses.png'))
    plt.close()

    return encoder, decoder_epi, decoder_hyp

# ========== Month Model Training and Prediction ==========
def train_and_predict_month_model(month_model, optimizer, dl_month_train, dl_month_test, month_epochs, logger):
    """
    Train the month model and predict for the abnormal month.

    Args:
        month_model: The month model to be trained.
        optimizer: Optimizer for the month model.
        dl_month_train: Dataloader for training.
        dl_month_test: Dataloader for testing.
        month_epochs: Number of epochs to train the month model.

    Returns:
        Test predictions and model state corresponding to the smallest validation RMSE.
    """

    logger.info("Training the month model...")
    for epoch in range(month_epochs):
        train_pred = []
        train_label = []

        valid_batch_count_train = 0

        # Training phase
        month_model.train()
        for x, y, sim, lake_id, year, month in dl_month_train:
            x, y, sim = x[:, :, :-2].to(device), y.to(device), sim.to(device)
            x = x[:, :, :26]

            mask = (y != -11)
            mask = mask.to(device)
            if (torch.sum(mask) == 0):
                continue

            valid_batch_count_train += 1

            pred = month_model(x)

            loss = RMSE(pred, y)
            weighted_loss = (loss * sim).mean()

            optimizer.zero_grad()
            weighted_loss.backward()
            optimizer.step()

            # logger.info("pred shape in month model training: ", pred.shape)
            # logger.info("y shape in month model training: ", y.shape)
            train_pred.append(pred.cpu().detach())
            train_label.append(y.cpu().detach())

        train_pred = torch.cat(train_pred, dim=0)
        train_label = torch.cat(train_label, dim=0)
        train_rmse = RMSE(train_pred.to(device), train_label.to(device))

        # Validation phase
        month_model.eval()
        # Test phase
        test_pred = []
        test_label = []
        with torch.no_grad():
            for x, y, _, lake_id, year, month in dl_month_test:
                x, y = x[:, :, :-2].to(device), y.to(device)
                x = x[:, :, :26]
                
                pred = month_model(x)

                # logger.info("pred shape in month model test: ", pred.shape)
                # logger.info("y shape in month model test: ", y.shape)
                
                test_pred.append(pred.cpu())
                test_label.append(y.cpu())

            test_pred = torch.cat(test_pred, dim=0)
            test_label = torch.cat(test_label, dim=0)
            test_rmse = RMSE(test_pred.to(device), test_label.to(device))

        logger.info(f"Epoch [{epoch + 1}/{month_epochs}], Train RMSE: {train_rmse:.4f}, Test RMSE: {test_rmse:.4f}")

    logger.info(f"Training Complete.")

    return test_pred

def train_and_predict_month_model_sim_pretrain(month_model, optimizer, dl_month_train, dl_month_test, month_epochs, flag='epi', logger=None):
    """
    Train the month model and predict for the abnormal month.

    Args:
        month_model: The month model to be trained.
        optimizer: Optimizer for the month model.
        dl_month_train: Dataloader for training.
        dl_month_test: Dataloader for testing.
        month_epochs: Number of epochs to train the month model.

    Returns:
        Test predictions and model state corresponding to the smallest validation RMSE.
    """
    best_test_rmse = float('inf')
    best_predictions = None
    best_state = None

    logger.info("Pretraining the month model use simulatons...")
    for epoch in range(month_epochs):
        train_pred = []
        train_label = []

        valid_batch_count_train = 0

        # Training phase
        month_model.train()
        for x, _, sim, lake_id, year, month in dl_month_train:
            if flag == 'epi':
                y = x[:, :, -2]
            else:
                y = x[:, :, -1]
                
            x, y, sim = x[:, :, :-2].to(device), y.to(device), sim.to(device)
            x = x[:, :, :26]

            valid_batch_count_train += 1

            pred = month_model(x)

            loss = RMSE(pred, y)
            weighted_loss = (loss * sim).mean()

            optimizer.zero_grad()
            weighted_loss.backward()
            optimizer.step()

            train_pred.append(pred.cpu().detach())
            train_label.append(y.cpu().detach())

        train_pred = torch.cat(train_pred, dim=0)
        train_label = torch.cat(train_label, dim=0)
        train_rmse = RMSE(train_pred.to(device), train_label.to(device))

        # Validation phase
        month_model.eval()
        # Test phase
        test_pred = []
        test_label = []
        with torch.no_grad():
            for x, _, _, lake_id, year, month in dl_month_test:
                if flag == 'epi':
                    y = x[:, :, -2]
                else:
                    y = x[:, :, -1]
                
                x, y = x[:, :, :-2].to(device), y.to(device)
                x = x[:, :, :26]
                
                pred = month_model(x)
                
                test_pred.append(pred.cpu())
                test_label.append(y.cpu())

            test_pred = torch.cat(test_pred, dim=0)
            test_label = torch.cat(test_label, dim=0)
            test_rmse = RMSE(test_pred.to(device), test_label.to(device))

        logger.info(f"Epoch [{epoch + 1}/{month_epochs}], Train RMSE: {train_rmse:.4f}, Test True RMSE: {test_rmse:.4f}")

        if test_rmse < best_test_rmse:
            best_test_rmse = test_rmse
            best_state = copy.deepcopy(month_model)
            best_predictions = test_pred  # Save test predictions corresponding to best validation RMSE    
    
    return best_predictions, best_state

# ==== calculate the optimal threshold using the ROC curve and Youden's Index
def calculate_roc_threshold(true_labels, predicted_probs):
    """
    Calculate the optimal threshold using the ROC curve and Youden's Index.

    Args:
        true_labels (list or np.array): Ground truth binary labels.
        predicted_probs (list or np.array): Predicted probabilities from the model.

    Returns:
        float: Optimal threshold value.
    """
    fpr, tpr, thresholds = roc_curve(true_labels, predicted_probs)
    youden_index = tpr - fpr
    best_idx = youden_index.argmax()
    optimal_threshold = thresholds[best_idx]

    return optimal_threshold

# ========= Discriminator Training =========
def run_decoder_training_and_save_labels_sim(yearly_model, dl_val, decoder_epi, decoder_hyp, encoder, lr_monthly_model, lr_monthly_model_pt, weight_decay_monthly_model, month_epochs, month_epochs_pretrain, top_n, save_path, precomputed_embeddings_epi, precomputed_samples_epi, precomputed_embeddings_hyp, precomputed_samples_hyp, decoder_epi_path, decoder_hyp_path, logger):
    logger.info("Validation Start: ")
    encoder.eval()
    yearly_model.eval()

    val_epi_pred = []
    val_epi_label = []
    val_hyp_pred = []
    val_hyp_label = []
    val_sim_epi = []
    val_sim_hyp = []

    val_original_pred_epi = []
    val_original_pred_hyp = []

    monthly_input_epi = []
    monthly_input_hyp = []
    monthly_labels_epi = []
    monthly_labels_hyp = []


    for x, y_epi, y_hyp, lake_id, year in dl_val:
        all_x = x.to(device)
        sim_epi = x[:, :, -2].to(device)
        sim_hyp = x[:, :, -1].to(device)
        x, y_epi, y_hyp, year = x[:,:,:-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)
        y_epi = y_epi.squeeze(-1)
        y_hyp = y_hyp.squeeze(-1)

        mask1 = (y_epi != -11)
        mask1 = mask1.to(device)
        mask2 = (y_hyp != -11)
        mask2 = mask2.to(device)
        if (torch.sum(mask1) == 0 and torch.sum(mask2) == 0):
            continue

        with torch.no_grad():
            pred_epi, pred_hyp = yearly_model(x)
        
        original_pred_epi, original_pred_hyp = pred_epi.clone(), pred_hyp.clone()

        for i in range(12):
            month_x = all_x[:, i*30:(i+1)*30, :]
            month_y_epi = y_epi[:, i*30:(i+1)*30]
            month_y_hyp = y_hyp[:, i*30:(i+1)*30]

            # Mask to check if there are valid values for the month
            mask_epi = (month_y_epi != -11).to(device)
            mask_hyp = (month_y_hyp != -11).to(device)
            if torch.sum(mask_epi) == 0 and torch.sum(mask_hyp) == 0:
                continue

            embedding = encoder(month_x)
            # month_sample_epi = (month_x, month_y_epi, lake_id[0], year, i+1) 
            # month_sample_hyp = (month_x, month_y_hyp, lake_id[0], year, i+1)
            month_sample_epi = (month_x.squeeze(0).to(device), month_y_epi.squeeze(0).to(device), lake_id[0], year, i+1) 
            month_sample_hyp = (month_x.squeeze(0).to(device), month_y_hyp.squeeze(0).to(device), lake_id[0], year, i+1)

            if torch.sum(mask_epi) > 0:
            # if torch.sum(mask_epi) == 1:
                dl_month_train, dl_month_test = enhanced_retrieve_and_prepare_dataloaders(embedding, precomputed_embeddings_epi, precomputed_samples_epi, month_sample_epi, top_n, for_joint_training=False, target_type='epi', logger=logger)
                # Check if dataloaders were successfully created
                if dl_month_train is None or dl_month_test is None:
                    logger.info(f"Failed to create dataloaders for lake {month_sample_epi[2]}, skipping epi processing")
                    continue
                
                decoder_epi.load_state_dict(torch.load(decoder_epi_path))
                optimizer_decoder_epi = torch.optim.Adam(decoder_epi.parameters(), lr=lr_monthly_model_pt, weight_decay=weight_decay_monthly_model)
                _, decoder_epi_state = train_and_predict_month_model_sim_pretrain(decoder_epi, optimizer_decoder_epi, dl_month_train, dl_month_test, month_epochs_pretrain, flag='epi', logger=logger)
                optimizer_decoder_epi = torch.optim.Adam(decoder_epi_state.parameters(), lr=lr_monthly_model, weight_decay=weight_decay_monthly_model)
                pred_epi_monthly = train_and_predict_month_model(decoder_epi_state, optimizer_decoder_epi, dl_month_train, dl_month_test, month_epochs, logger=logger)
                pred_epi_monthly = pred_epi_monthly.to(device)

                decoder_epi_rmse = RMSE(pred_epi_monthly, month_y_epi)
                yearly_model_rmse_epi = RMSE(pred_epi[:, i * 30:(i + 1) * 30], month_y_epi)

                pred_epi_expanded = pred_epi[:, i * 30:(i + 1) * 30].unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]
                pred_epi_monthly_expanded = pred_epi_monthly.unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]

                # Concatenate them along the last dimension (feature dimension)
                # discriminator_input_epi = torch.cat([month_x[:, :, :26], month_x[:, :, -2:], pred_epi_expanded, pred_epi_monthly_expanded], dim=2)  # Shape: [batchsize, 30, 44]
                discriminator_input_epi = torch.cat([month_x, pred_epi_expanded, pred_epi_monthly_expanded], dim=2)  # Shape: [batchsize, 30, 44]

                if decoder_epi_rmse < yearly_model_rmse_epi:
                    logger.info(f"decoder_epi is better in month {i+1} in year {year} and of lake {lake_id}:")
                    label_epi = 1
                    pred_epi[:, i*30:(i+1)*30] = pred_epi_monthly
                else:
                    logger.info(f"yearly model epi is better in month {i+1} in year {year} and of lake {lake_id}:")
                    label_epi = 0

                label_epi = torch.tensor([label_epi], dtype=torch.float32, device=device)

                monthly_input_epi.append(discriminator_input_epi.squeeze().cpu().detach())
                monthly_labels_epi.append(label_epi.detach().cpu())

            if torch.sum(mask_hyp) > 0:
            # if torch.sum(mask_hyp) == 1:
                dl_month_train, dl_month_test = enhanced_retrieve_and_prepare_dataloaders(embedding, precomputed_embeddings_hyp, precomputed_samples_hyp, month_sample_hyp, top_n, for_joint_training=False, target_type='hyp', logger=logger)
                
                # Check if dataloaders were successfully created
                if dl_month_train is None or dl_month_test is None:
                    logger.info(f"Failed to create dataloaders for lake {month_sample_hyp[2]}, skipping hyp processing")
                    continue
                
                decoder_hyp.load_state_dict(torch.load(decoder_hyp_path))
                optimizer_decoder_hyp = torch.optim.Adam(decoder_hyp.parameters(), lr=lr_monthly_model_pt, weight_decay=weight_decay_monthly_model)
                _, decoder_hyp_state = train_and_predict_month_model_sim_pretrain(decoder_hyp, optimizer_decoder_hyp, dl_month_train, dl_month_test, month_epochs_pretrain, flag='hyp', logger=logger)
                optimizer_decoder_hyp = torch.optim.Adam(decoder_hyp_state.parameters(), lr=lr_monthly_model, weight_decay=weight_decay_monthly_model)
                pred_hyp_monthly = train_and_predict_month_model(decoder_hyp_state, optimizer_decoder_hyp, dl_month_train, dl_month_test, month_epochs, logger=logger)
                pred_hyp_monthly = pred_hyp_monthly.to(device)

                decoder_hyp_rmse = RMSE(pred_hyp_monthly, month_y_hyp)
                yearly_model_rmse_hyp = RMSE(pred_hyp[:, i * 30:(i + 1) * 30], month_y_hyp)

                pred_hyp_expanded = pred_hyp[:, i * 30:(i + 1) * 30].unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]
                pred_hyp_monthly_expanded = pred_hyp_monthly.unsqueeze(-1).repeat(1, 1, 2)

                # discriminator_input_hyp = torch.cat([month_x[:, :, :26], month_x[:, :, -2:], pred_hyp_expanded, pred_hyp_monthly_expanded], dim=2)  # Shape: [batchsize, 30, 44]
                discriminator_input_hyp = torch.cat([month_x, pred_hyp_expanded, pred_hyp_monthly_expanded], dim=2)  # Shape: [batchsize, 30, 44]
                # discriminator_input_hyp = torch.cat([pred_hyp_expanded, pred_hyp_monthly_expanded], dim=2)  # Shape: [batchsize, 30, 44]

                if decoder_hyp_rmse < yearly_model_rmse_hyp:
                    logger.info(f"decoder_hyp is better in month {i+1} in year {year} and of lake {lake_id}:")
                    label_hyp = 1
                    pred_hyp[:, i*30:(i+1)*30] = pred_hyp_monthly
                else:
                    logger.info(f"yearly model hyp is better in month {i+1} in year {year} and of lake {lake_id}:")
                    label_hyp = 0
                
                label_hyp = torch.tensor([label_hyp], dtype=torch.float32, device=device)

                monthly_input_hyp.append(discriminator_input_hyp.squeeze().cpu().detach())
                monthly_labels_hyp.append(label_hyp.detach().cpu())

        val_epi_pred.append(pred_epi.squeeze().cpu().detach())
        val_epi_label.append(y_epi.squeeze().cpu().detach())
        val_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
        val_hyp_label.append(y_hyp.squeeze().cpu().detach())
        val_sim_epi.append(sim_epi.squeeze().cpu().detach())
        val_sim_hyp.append(sim_hyp.squeeze().cpu().detach())
        val_original_pred_epi.append(original_pred_epi.squeeze().cpu().detach())
        val_original_pred_hyp.append(original_pred_hyp.squeeze().cpu().detach())
    

    val_epi_pred = torch.cat(val_epi_pred, dim=0)
    val_epi_label = torch.cat(val_epi_label, dim=0)
    val_hyp_pred = torch.cat(val_hyp_pred, dim=0)
    val_hyp_label = torch.cat(val_hyp_label, dim=0)
    val_sim_epi = torch.cat(val_sim_epi, dim=0)
    val_sim_hyp = torch.cat(val_sim_hyp, dim=0)
    val_original_pred_epi = torch.cat(val_original_pred_epi, dim=0)
    val_original_pred_hyp = torch.cat(val_original_pred_hyp, dim=0)

    val_epi_RMSE = RMSE(val_epi_pred.to(device), val_epi_label.to(device))
    val_hyp_RMSE = RMSE(val_hyp_pred.to(device), val_hyp_label.to(device))
    val_original_epi_RMSE = RMSE(val_original_pred_epi.to(device), val_epi_label.to(device))
    val_original_hyp_RMSE = RMSE(val_original_pred_hyp.to(device), val_hyp_label.to(device))

    logger.info(f'Val RMSE Epi: {val_epi_RMSE:.4f}')
    logger.info(f'Val RMSE Hyp: {val_hyp_RMSE:.4f}')
    logger.info(f'Val Original RMSE Epi: {val_original_epi_RMSE:.4f}')
    logger.info(f'Val Original RMSE Hyp: {val_original_hyp_RMSE:.4f}')

    logger.info("monthly_input_epi shape: ", len(monthly_input_epi))
    logger.info("monthly_labels_epi shape: ", len(monthly_labels_epi))
    logger.info("monthly_input_hyp shape: ", len(monthly_input_hyp))
    logger.info("monthly_labels_hyp shape: ", len(monthly_labels_hyp))

    monthly_input_epi = torch.cat(monthly_input_epi, dim=0)
    monthly_labels_epi = torch.cat(monthly_labels_epi, dim=0)
    monthly_input_hyp = torch.cat(monthly_input_hyp, dim=0)
    monthly_labels_hyp = torch.cat(monthly_labels_hyp, dim=0)

    lable_epi_expanded = np.repeat(monthly_labels_epi, 30)
    lable_hyp_expanded = np.repeat(monthly_labels_hyp, 30)

    logger.info("monthly_input_epi shape: ", monthly_input_epi.shape)
    logger.info("monthly_labels_epi shape: ", lable_epi_expanded.shape)
    logger.info("monthly_input_hyp shape: ", monthly_input_hyp.shape)
    logger.info("monthly_labels_hyp shape: ", lable_hyp_expanded.shape)

    monthly_input_epi_np = monthly_input_epi.cpu().numpy()
    monthly_labels_epi_np = lable_epi_expanded.cpu().numpy()

    monthly_input_hyp_np = monthly_input_hyp.cpu().numpy()
    monthly_labels_hyp_np = lable_hyp_expanded.cpu().numpy()

    # Create DataFrame properly
    df_epi = pd.DataFrame(
        monthly_input_epi_np,  # This will create multiple columns automatically
        columns=[f"Feature_{i}" for i in range(monthly_input_epi_np.shape[1])]
    )
    df_epi["Label"] = monthly_labels_epi_np  # Add label as a single column

    df_hyp = pd.DataFrame(
        monthly_input_hyp_np,
        columns=[f"Feature_{i}" for i in range(monthly_input_hyp_np.shape[1])]
    )
    df_hyp["Label"] = monthly_labels_hyp_np

    df_epi.to_csv(os.path.join(save_path, "discriminator_epi_data.csv"), index=False)
    df_hyp.to_csv(os.path.join(save_path, "discriminator_hyp_data.csv"), index=False)

    return df_epi, df_hyp

def compute_confusion_matrix(all_labels, all_outputs, threshold, logger):
    # Convert to numpy arrays
    all_labels = np.array(all_labels).flatten()
    all_outputs = np.array(all_outputs).flatten()

    # Apply threshold to convert probabilities to binary predictions
    predictions = (all_outputs >= threshold).astype(int)

    # Compute confusion matrix
    cm = confusion_matrix(all_labels, predictions)

    # logger.info confusion matrix and classification report
    logger.info("Confusion Matrix:")
    logger.info(cm)

    TN, FP, FN, TP = cm.ravel()
    logger.info(f"TN: {TN}, FP: {FP}, FN: {FN}, TP: {TP}")
    # Compute the requested metric
    TPR = TP / (TP + FN) if (TP + FN) > 0 else 0  # Recall (Sensitivity)
    FPR = FP / (TN + FP) if (TN + FP) > 0 else 0  # False Positive Rate
    ratio = TPR / FPR if FPR > 0 else float('inf') 

    logger.info(f"Ratio (TPR/FPR): {ratio:.4f} (Compare with 1)")


    return cm

def discriminator_training(discriminator, optimizer_discriminator, dl, discriminator_epochs, save_path, flag='epi', logger=None):
    criterion = nn.BCELoss()
    loss_history = []
    all_labels = []
    all_outputs = []

    for epoch in range(discriminator_epochs):
        logger.info(f"Discriminator training: Epoch {epoch + 1}/{discriminator_epochs}")
        
        epoch_loss = 0.0
        discriminator.train()

        all_labels_confusion = []
        all_outputs_confusion = []


        for inputs, labels in dl:
            all_inputs = inputs.to(device)
            inputs, labels = inputs[:, :, :-2].to(device), labels.to(device)
            labels = labels[:, 0, :]

            optimizer_discriminator.zero_grad()
            outputs = discriminator(all_inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer_discriminator.step()

            epoch_loss += loss.item()
            
            all_outputs.extend(outputs.detach().cpu().numpy())
            all_labels.extend(labels.detach().cpu().numpy())
            all_labels_confusion.extend(labels.detach().cpu().numpy())
            all_outputs_confusion.extend(outputs.detach().cpu().numpy())

        avg_loss = epoch_loss / len(dl)
        loss_history.append(avg_loss)

        logger.info(f"Discriminator training: Epoch {epoch + 1} - Loss: {avg_loss:.4f}")


        # make it in the epoch to track the discriminator performance during training
        optimal_threshold = calculate_roc_threshold(all_labels_confusion, all_outputs_confusion)
        logger.info(f"Optimal Threshold: {optimal_threshold:.4f}")

        cm = compute_confusion_matrix(all_labels_confusion, all_outputs_confusion, optimal_threshold, logger)
    
    # After training, plot the loss history
    plt.clf()
    plt.plot(loss_history, label=f"{flag} Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()

    file_name = f'discriminator_training_loss_{flag}.png'
    full_save_path = os.path.join(save_path, file_name)
    plt.savefig(full_save_path)
    plt.close()

    discriminator_file_name = f"discriminator_{flag}.pth"
    full_discriminator_path = os.path.join(save_path, discriminator_file_name)
    torch.save(discriminator.state_dict(), full_discriminator_path)

    logger.info(f"Loss plot saved to {full_save_path}")
    logger.info(f"Discriminator model saved to {full_discriminator_path}")

    return discriminator, optimal_threshold

# ========== Discriminator Testing ==========
def discriminator_testing(yearly_model, dl_test, discriminator_epi, optimal_threshold_epi, discriminator_hyp, optimal_threshold_hyp, decoder_epi, decoder_hyp, encoder, lr_monthly_model, lr_monthly_model_pt, weight_decay_monthly_model, month_epochs, month_epochs_pretrain, top_n, save_path, precomputed_embeddings_epi, precomputed_samples_epi, precomputed_embeddings_hyp, precomputed_samples_hyp, decoder_epi_path, decoder_hyp_path, logger):
    logger.info("Test Start: ")
    test_epi_pred = []
    test_epi_label = []
    test_hyp_pred = []
    test_hyp_label = []
    test_sim_epi = []
    test_sim_hyp = []

    test_original_epi_pred = []
    test_original_hyp_pred = []
    test_best_epi_pred = []
    test_best_hyp_pred = []
    all_decoder_epi_pred = []
    all_decoder_hyp_pred = []

    all_true_labels_epi = []
    all_predicted_probs_epi = []
    all_true_labels_hyp = []
    all_predicted_probs_hyp = []

    lake_ids = []
    years = []

    decoder_better_count_epi = [0] * 12  # For each month (1-12)
    decoder_discriminator_count_epi = [0] * 12  # For each month (1-12)
    total_months_count_epi = [0] * 12  # Total comparisons for each month

    decoder_better_count_hyp = [0] * 12  # For each month (1-12)
    decoder_discriminator_count_hyp = [0] * 12  # For each month (1-12)
    total_months_count_hyp = [0] * 12  # Total comparisons for each month

    discriminator_epi.eval()
    discriminator_hyp.eval()


    for x, y_epi, y_hyp, lake_id, year in dl_test:
        logger.info(f"Testing lake_id: {lake_id[0]} in year: {year}")
        
        all_x = x.to(device)
        sim_epi = x[:, :, -2].to(device)
        sim_hyp = x[:, :, -1].to(device)
        x, y_epi, y_hyp, year = x[:,:,:-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)
        y_epi = y_epi.squeeze(-1)
        y_hyp = y_hyp.squeeze(-1)

        mask1 = (y_epi != -11)
        mask1 = mask1.to(device)
        mask2 = (y_hyp != -11)
        mask2 = mask2.to(device)
        if (torch.sum(mask1) == 0 and torch.sum(mask2) == 0):
            continue
        
        with torch.no_grad():
            pred_epi, pred_hyp = yearly_model(x)
        
        original_pred_epi, original_pred_hyp = pred_epi.clone(), pred_hyp.clone()
        best_pred_epi, best_pred_hyp = pred_epi.clone(), pred_hyp.clone()
        all_decoder_epi, all_decoder_hyp = pred_epi.clone(), pred_hyp.clone()

        for i in range(12):
            month_x = all_x[:, i*30:(i+1)*30, :]
            month_y_epi = y_epi[:, i*30:(i+1)*30]
            month_y_hyp = y_hyp[:, i*30:(i+1)*30]

            month_sample_epi = (month_x.squeeze(0).to(device), month_y_epi.squeeze(0).to(device), lake_id[0], year, i+1) 
            month_sample_hyp = (month_x.squeeze(0).to(device), month_y_hyp.squeeze(0).to(device), lake_id[0], year, i+1)

            ################
            mask_epi = (month_y_epi != -11).to(device)
            mask_hyp = (month_y_hyp != -11).to(device)
            if torch.sum(mask_epi) == 0 and torch.sum(mask_hyp) == 0:
                continue
            ################

            embedding = encoder(month_x)

            if torch.sum(mask_epi) > 0:
                dl_month_train, dl_month_test = enhanced_retrieve_and_prepare_dataloaders(embedding, precomputed_embeddings_epi, precomputed_samples_epi, month_sample_epi, top_n, for_joint_training=False, target_type='epi', logger=logger)

                if dl_month_train is None or dl_month_test is None:
                    logger.info(f"Failed to create dataloaders for lake {month_sample_epi[2]}, skipping epi processing")
                    continue

                decoder_epi.load_state_dict(torch.load(decoder_epi_path))
                optimizer_decoder_epi = torch.optim.Adam(decoder_epi.parameters(), lr=lr_monthly_model_pt, weight_decay=weight_decay_monthly_model)
                _, decoder_epi_state = train_and_predict_month_model_sim_pretrain(decoder_epi, optimizer_decoder_epi, dl_month_train, dl_month_test, month_epochs_pretrain, flag='epi', logger=logger)
                optimizer_decoder_epi = torch.optim.Adam(decoder_epi_state.parameters(), lr=lr_monthly_model, weight_decay=weight_decay_monthly_model)
                updated_pred_epi = train_and_predict_month_model(decoder_epi_state, optimizer_decoder_epi, dl_month_train, dl_month_test, month_epochs, logger=logger)
                updated_pred_epi = updated_pred_epi.to(device)

                decoder_epi_rmse = RMSE(updated_pred_epi, month_y_epi)
                yearly_model_rmse_epi = RMSE(pred_epi[:, i * 30:(i + 1) * 30], month_y_epi)

                pred_epi_expanded = pred_epi[:, i * 30:(i + 1) * 30].unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]
                updated_pred_epi_expanded = updated_pred_epi.unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]

                discriminator_input_epi = torch.cat([month_x, pred_epi_expanded, updated_pred_epi_expanded], dim=2)  # Shape: [batchsize, 30, 44]

                discriminate_pred_epi = discriminator_epi(discriminator_input_epi)

                if decoder_epi_rmse < yearly_model_rmse_epi:
                    logger.info(f"decoder_epi is better in month {i + 1} in year {year} and of lake {lake_id}:")
                    label_epi = 1
                    decoder_better_count_epi[i] += 1
                    best_pred_epi[:, i*30:(i+1)*30] = updated_pred_epi
                else:
                    label_epi = 0
                    logger.info(f"yearly model epi is better in month {i + 1} in year {year} and of lake {lake_id}:")
            

                if discriminate_pred_epi > optimal_threshold_epi:
                    pred_epi[:, i*30:(i+1)*30] = updated_pred_epi
                    decoder_discriminator_count_epi[i] += 1

                all_decoder_epi[:, i*30:(i+1)*30] = updated_pred_epi

                label_epi = torch.tensor([label_epi], dtype=torch.float32, device=device)

                all_true_labels_epi.append(label_epi.item())
                all_predicted_probs_epi.append(discriminate_pred_epi.item())

                total_months_count_epi[i] += 1
            
            if torch.sum(mask_hyp) > 0:
                dl_month_train, dl_month_test = enhanced_retrieve_and_prepare_dataloaders(embedding, precomputed_embeddings_hyp, precomputed_samples_hyp, month_sample_hyp, top_n, for_joint_training=False, target_type='hyp', logger=logger)

                if dl_month_train is None or dl_month_test is None:
                    logger.info(f"Failed to create dataloaders for lake {month_sample_hyp[2]}, skipping hyp processing")
                    continue

                decoder_hyp.load_state_dict(torch.load(decoder_hyp_path))
                optimizer_decoder_hyp = torch.optim.Adam(decoder_hyp.parameters(), lr=lr_monthly_model_pt, weight_decay=weight_decay_monthly_model)
                _, decoder_hyp_state = train_and_predict_month_model_sim_pretrain(decoder_hyp, optimizer_decoder_hyp, dl_month_train, dl_month_test, month_epochs_pretrain, flag='hyp', logger=logger)
                optimizer_decoder_hyp = torch.optim.Adam(decoder_hyp_state.parameters(), lr=lr_monthly_model, weight_decay=weight_decay_monthly_model)
                updated_pred_hyp = train_and_predict_month_model(decoder_hyp_state, optimizer_decoder_hyp, dl_month_train, dl_month_test, month_epochs, logger=logger)
                updated_pred_hyp = updated_pred_hyp.to(device)

                decoder_hyp_rmse = RMSE(updated_pred_hyp, month_y_hyp)
                yearly_model_rmse_hyp = RMSE(pred_hyp[:, i * 30:(i + 1) * 30], month_y_hyp)

                pred_hyp_expanded = pred_hyp[:, i * 30:(i + 1) * 30].unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]
                updated_pred_hyp_expanded = updated_pred_hyp.unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]

                discriminator_input_hyp = torch.cat([month_x, pred_hyp_expanded, updated_pred_hyp_expanded], dim=2)  # Shape: [batchsize, 30, 44]

                discriminate_pred_hyp = discriminator_hyp(discriminator_input_hyp)


                if decoder_hyp_rmse < yearly_model_rmse_hyp:
                    logger.info(f"decoder_hyp is better in month {i + 1} in year {year} and of lake {lake_id}:")
                    label_hyp = 1
                    best_pred_hyp[:, i*30:(i+1)*30] = updated_pred_hyp
                    decoder_better_count_hyp[i] += 1
                else:
                    label_hyp = 0
                    logger.info(f"yearly model hyp is better in month {i + 1} in year {year} and of lake {lake_id}:")
                ###########


                if discriminate_pred_hyp > optimal_threshold_hyp:
                    pred_hyp[:, i*30:(i+1)*30] = updated_pred_hyp
                    decoder_discriminator_count_hyp[i] += 1
                
                all_decoder_hyp[:, i*30:(i+1)*30] = updated_pred_hyp

                label_hyp = torch.tensor([label_hyp], dtype=torch.float32, device=device)

                all_true_labels_hyp.append(label_hyp.item())
                all_predicted_probs_hyp.append(discriminate_pred_hyp.item())

                total_months_count_hyp[i] += 1
        
        lake_ids.append(lake_id[0])
        years.append(year.cpu().item())

        test_epi_pred.append(pred_epi.squeeze().cpu().detach())
        test_epi_label.append(y_epi.squeeze().cpu().detach())
        test_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
        test_hyp_label.append(y_hyp.squeeze().cpu().detach())
        test_sim_epi.append(sim_epi.squeeze().cpu().detach())
        test_sim_hyp.append(sim_hyp.squeeze().cpu().detach())
        test_original_epi_pred.append(original_pred_epi.squeeze().cpu().detach())
        test_original_hyp_pred.append(original_pred_hyp.squeeze().cpu().detach())
        test_best_epi_pred.append(best_pred_epi.squeeze().cpu().detach())
        test_best_hyp_pred.append(best_pred_hyp.squeeze().cpu().detach())
        all_decoder_epi_pred.append(all_decoder_epi.squeeze().cpu().detach())
        all_decoder_hyp_pred.append(all_decoder_hyp.squeeze().cpu().detach())


    logger.info("Decoder Better Count for each month (Epi):", decoder_better_count_epi)
    logger.info("Total Count for discriminator thinks better for each month (Epi):", decoder_discriminator_count_epi)
    logger.info("Total Count for each month (Epi):", total_months_count_epi)
    logger.info("Decoder Better Count for each month (Hyp):", decoder_better_count_hyp)
    logger.info("Total Count for discriminator thinks better for each month (Hyp):", decoder_discriminator_count_hyp)
    logger.info("Total Count for each month (Hyp):", total_months_count_hyp)

    logger.info(len(lake_ids), len(years), len(test_epi_label), len(test_epi_pred), len(test_original_epi_pred))

    test_epi_pred = torch.cat(test_epi_pred, dim=0)
    test_epi_label = torch.cat(test_epi_label, dim=0)
    test_hyp_pred = torch.cat(test_hyp_pred, dim=0)
    test_hyp_label = torch.cat(test_hyp_label, dim=0)
    test_sim_epi = torch.cat(test_sim_epi, dim=0)
    test_sim_hyp = torch.cat(test_sim_hyp, dim=0)
    test_original_epi_pred = torch.cat(test_original_epi_pred, dim=0)
    test_original_hyp_pred = torch.cat(test_original_hyp_pred, dim=0)
    test_best_epi_pred = torch.cat(test_best_epi_pred, dim=0)
    test_best_hyp_pred = torch.cat(test_best_hyp_pred, dim=0)
    all_decoder_epi_pred = torch.cat(all_decoder_epi_pred, dim=0)
    all_decoder_hyp_pred = torch.cat(all_decoder_hyp_pred, dim=0)

    test_epi_RMSE = RMSE(test_epi_pred.to(device), test_epi_label.to(device))
    test_hyp_RMSE = RMSE(test_hyp_pred.to(device), test_hyp_label.to(device))
    test_origial_epi_RMSE = RMSE(test_original_epi_pred.to(device), test_epi_label.to(device))
    test_original_hyp_RMSE = RMSE(test_original_hyp_pred.to(device), test_hyp_label.to(device))
    test_best_epi_RMSE = RMSE(test_best_epi_pred.to(device), test_epi_label.to(device))
    test_best_hyp_RMSE = RMSE(test_best_hyp_pred.to(device), test_hyp_label.to(device))
    all_decoder_epi_pred_RMSE = RMSE(all_decoder_epi_pred.to(device), test_epi_label.to(device))
    all_decoder_hyp_pred_RMSE = RMSE(all_decoder_hyp_pred.to(device), test_hyp_label.to(device))

    logger.info(f'Test RMSE Epi: {test_epi_RMSE:.4f}')
    logger.info(f'Test RMSE Hyp: {test_hyp_RMSE:.4f}')
    logger.info(f'Test Original RMSE Epi: {test_origial_epi_RMSE:.4f}')
    logger.info(f'Test Original RMSE Hyp: {test_original_hyp_RMSE:.4f}')
    logger.info(f'Test Best RMSE Epi: {test_best_epi_RMSE:.4f}')
    logger.info(f'Test Best RMSE Hyp: {test_best_hyp_RMSE:.4f}')
    logger.info(f'Test All RMSE Epi: {all_decoder_epi_pred_RMSE:.4f}')
    logger.info(f'Test All RMSE Hyp: {all_decoder_hyp_pred_RMSE:.4f}')

    days_per_year = 360
    lake_ids_expanded = np.repeat(lake_ids, days_per_year)
    years_expanded = np.repeat(years, days_per_year)


    df_epi = pd.DataFrame({
        "Lake_ID": lake_ids_expanded,
        "Year": years_expanded,
        "Day": np.tile(np.arange(1, days_per_year + 1), len(lake_ids)),  # Add Day index (1 to 360)
        "True_Label_Epi": test_epi_label.cpu().numpy(),
        "Prediction_Epi": test_epi_pred.cpu().numpy(),
        "Sim_Epi": test_sim_epi.cpu().numpy(),
        "Original_Prediction_Epi": test_original_epi_pred.cpu().numpy(),
        "Best_Prediction_Epi": test_best_epi_pred.cpu().numpy(),
        "All_Decoder_Epi": all_decoder_epi_pred.cpu().numpy()
    })

    # Create a DataFrame for Hyp
    df_hyp = pd.DataFrame({
        "Lake_ID": lake_ids_expanded,
        "Year": years_expanded,
        "Day": np.tile(np.arange(1, days_per_year + 1), len(lake_ids)),  # Add Day index (1 to 360)
        "True_Label_Hyp": test_hyp_label.cpu().numpy(),
        "Prediction_Hyp": test_hyp_pred.cpu().numpy(),
        "Sim_Hyp": test_sim_hyp.cpu().numpy(),
        "Original_Prediction_Hyp": test_original_hyp_pred.cpu().numpy(),
        "Best_Prediction_Hyp": test_best_hyp_pred.cpu().numpy(),
        "All_Decoder_Hyp": all_decoder_hyp_pred.cpu().numpy()
    })

    # Save DataFrames to CSV files
    df_epi.to_csv(os.path.join(save_path, "epi_predictions_test.csv"), index=False)
    df_hyp.to_csv(os.path.join(save_path, "hyp_predictions_test.csv"), index=False)

    result_save_path = os.path.join(save_path, "DO_test")
    os.makedirs(result_save_path, exist_ok=True)

    evaluate_and_save_rmse(epi_csv_path=os.path.join(save_path, "epi_predictions_test.csv"), hyp_csv_path=os.path.join(save_path, "hyp_predictions_test.csv"), save_dir=result_save_path, epi_col='Prediction_Epi', hyp_col='Prediction_Hyp', logger=logger)
    evaluate_and_save_rmse(epi_csv_path=os.path.join(save_path, "epi_predictions_test.csv"), hyp_csv_path=os.path.join(save_path, "hyp_predictions_test.csv"), save_dir=result_save_path, epi_col='All_Decoder_Epi', hyp_col='All_Decoder_Hyp', logger=logger)

def discriminator_testing_temp(yearly_model, dl_test, discriminator_epi, optimal_threshold_epi, discriminator_hyp, optimal_threshold_hyp, decoder_epi, decoder_hyp, encoder, lr_monthly_model, lr_monthly_model_pt, weight_decay_monthly_model, month_epochs, month_epochs_pretrain, top_n, save_path, precomputed_embeddings_epi, precomputed_samples_epi, precomputed_embeddings_hyp, precomputed_samples_hyp, decoder_epi_path, decoder_hyp_path, logger):
    logger.info("Test Start: ")
    test_epi_pred = []
    test_epi_label = []
    test_hyp_pred = []
    test_hyp_label = []
    test_sim_epi = []
    test_sim_hyp = []

    test_original_epi_pred = []
    test_original_hyp_pred = []
    test_best_epi_pred = []
    test_best_hyp_pred = []

    all_true_labels_epi = []
    all_predicted_probs_epi = []
    all_true_labels_hyp = []
    all_predicted_probs_hyp = []

    lake_ids = []
    years = []

    decoder_better_count_epi = [0] * 12  # For each month (1-12)
    decoder_discriminator_count_epi = [0] * 12  # For each month (1-12)
    total_months_count_epi = [0] * 12  # Total comparisons for each month

    decoder_better_count_hyp = [0] * 12  # For each month (1-12)
    decoder_discriminator_count_hyp = [0] * 12  # For each month (1-12)
    total_months_count_hyp = [0] * 12  # Total comparisons for each month

    discriminator_epi.eval()
    discriminator_hyp.eval()

    for x, y_epi, y_hyp, lake_id, year in dl_test:

        all_x = x.to(device)
        sim_epi = x[:, :, -2].to(device)
        sim_hyp = x[:, :, -1].to(device)
        x, y_epi, y_hyp, year = x[:,:,:-2].to(device), y_epi.to(device), y_hyp.to(device), year.to(device)
        y_epi = y_epi.squeeze(-1)
        y_hyp = y_hyp.squeeze(-1)

        mask1 = (y_epi != -11)
        mask1 = mask1.to(device)
        mask2 = (y_hyp != -11)
        mask2 = mask2.to(device)
        if (torch.sum(mask1) == 0 and torch.sum(mask2) == 0):
            continue

        with torch.no_grad():
            pred_epi, pred_hyp = yearly_model(x)
        
        original_pred_epi, original_pred_hyp = pred_epi.clone(), pred_hyp.clone()
        best_pred_epi, best_pred_hyp = pred_epi.clone(), pred_hyp.clone()

        for i in range(12):
            month_x = all_x[:, i*30:(i+1)*30, :]
            month_sim_epi = sim_epi[:, i*30:(i+1)*30]
            month_sim_hyp = sim_hyp[:, i*30:(i+1)*30]
            month_y_epi = y_epi[:, i*30:(i+1)*30]
            month_y_hyp = y_hyp[:, i*30:(i+1)*30]

            month_sample_epi = (month_x.squeeze(0).to(device), month_y_epi.squeeze(0).to(device), lake_id[0], year, i+1) 
            month_sample_hyp = (month_x.squeeze(0).to(device), month_y_hyp.squeeze(0).to(device), lake_id[0], year, i+1)

            ################
            mask_epi = (month_y_epi != -11).to(device)
            mask_hyp = (month_y_hyp != -11).to(device)
            if torch.sum(mask_epi) == 0 and torch.sum(mask_hyp) == 0:
                continue
            ################

            embedding = encoder(month_x)

            if torch.sum(mask_epi) > 0:
                dl_month_train, dl_month_test = enhanced_retrieve_and_prepare_dataloaders(embedding, precomputed_embeddings_epi, precomputed_samples_epi, month_sample_epi, top_n, for_joint_training=False, target_type='epi', logger=logger)
                if dl_month_train is None or dl_month_test is None:
                    logger.info(f"Failed to create dataloaders for lake {month_sample_epi[2]}, skipping epi processing")
                    continue
                decoder_epi.load_state_dict(torch.load(decoder_epi_path))
                optimizer_decoder_epi = torch.optim.Adam(decoder_epi.parameters(), lr=lr_monthly_model_pt, weight_decay=weight_decay_monthly_model)
                _, decoder_epi_state = train_and_predict_month_model_sim_pretrain(decoder_epi, optimizer_decoder_epi, dl_month_train, dl_month_test, month_epochs_pretrain, flag='epi', logger=logger)
                optimizer_decoder_epi = torch.optim.Adam(decoder_epi_state.parameters(), lr=lr_monthly_model, weight_decay=weight_decay_monthly_model)
                updated_pred_epi = train_and_predict_month_model(decoder_epi_state, optimizer_decoder_epi, dl_month_train, dl_month_test, month_epochs, logger=logger)
                updated_pred_epi = updated_pred_epi.to(device)

                decoder_epi_rmse = RMSE(updated_pred_epi, month_y_epi)
                yearly_model_rmse_epi = RMSE(pred_epi[:, i * 30:(i + 1) * 30], month_y_epi)

                decoder_epi_sim_rmse = RMSE(updated_pred_epi, month_sim_epi)
                yearly_model_sim_rmse_epi = RMSE(pred_epi[:, i * 30:(i + 1) * 30], month_sim_epi)

                pred_epi_expanded = pred_epi[:, i * 30:(i + 1) * 30].unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]
                updated_pred_epi_expanded = updated_pred_epi.unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]

                discriminator_input_epi = torch.cat([month_x, pred_epi_expanded, updated_pred_epi_expanded], dim=2)  # Shape: [batchsize, 30, 44]

                discriminate_pred_epi = discriminator_epi(discriminator_input_epi)

                if decoder_epi_rmse < yearly_model_rmse_epi:
                    logger.info(f"decoder_epi is better in month {i + 1} in year {year} and of lake {lake_id}:")
                    label_epi = 1
                    decoder_better_count_epi[i] += 1
                    best_pred_epi[:, i*30:(i+1)*30] = updated_pred_epi
                else:
                    label_epi = 0
                    logger.info(f"yearly model epi is better in month {i + 1} in year {year} and of lake {lake_id}:")
            

                if discriminate_pred_epi > optimal_threshold_epi and decoder_epi_sim_rmse < yearly_model_sim_rmse_epi:
                    logger.info(f"discriminator of decoder_epi is better in month {i + 1} in year {year} and of lake {lake_id} for sim:")
                    pred_epi[:, i*30:(i+1)*30] = updated_pred_epi
                    decoder_discriminator_count_epi[i] += 1


                label_epi = torch.tensor([label_epi], dtype=torch.float32, device=device)

                all_true_labels_epi.append(label_epi.item())
                all_predicted_probs_epi.append(discriminate_pred_epi.item())

                total_months_count_epi[i] += 1
            
            if torch.sum(mask_hyp) > 0:
                dl_month_train, dl_month_test = enhanced_retrieve_and_prepare_dataloaders(embedding, precomputed_embeddings_hyp, precomputed_samples_hyp, month_sample_hyp, top_n, for_joint_training=False, target_type='hyp', logger=logger)

                if dl_month_train is None or dl_month_test is None:
                    logger.info(f"Failed to create dataloaders for lake {month_sample_hyp[2]}, skipping hyp processing")
                    continue

                decoder_hyp.load_state_dict(torch.load(decoder_hyp_path))
                optimizer_decoder_hyp = torch.optim.Adam(decoder_hyp.parameters(), lr=lr_monthly_model_pt, weight_decay=weight_decay_monthly_model)
                _, decoder_hyp_state = train_and_predict_month_model_sim_pretrain(decoder_hyp, optimizer_decoder_hyp, dl_month_train, dl_month_test, month_epochs_pretrain, flag='hyp', logger=logger)
                optimizer_decoder_hyp = torch.optim.Adam(decoder_hyp_state.parameters(), lr=lr_monthly_model, weight_decay=weight_decay_monthly_model)
                updated_pred_hyp = train_and_predict_month_model(decoder_hyp_state, optimizer_decoder_hyp, dl_month_train, dl_month_test, month_epochs, logger=logger)
                updated_pred_hyp = updated_pred_hyp.to(device)

                decoder_hyp_rmse = RMSE(updated_pred_hyp, month_y_hyp)
                yearly_model_rmse_hyp = RMSE(pred_hyp[:, i * 30:(i + 1) * 30], month_y_hyp)

                decoder_hyp_sim_rmse = RMSE(updated_pred_hyp, month_sim_hyp)
                yearly_model_sim_rmse_hyp = RMSE(pred_hyp[:, i * 30:(i + 1) * 30], month_sim_hyp)

                pred_hyp_expanded = pred_hyp[:, i * 30:(i + 1) * 30].unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]
                updated_pred_hyp_expanded = updated_pred_hyp.unsqueeze(-1).repeat(1, 1, 2)  # Shape: [batchsize, 30, 1]

                discriminator_input_hyp = torch.cat([month_x, pred_hyp_expanded, updated_pred_hyp_expanded], dim=2)  # Shape: [batchsize, 30, 44]

                discriminate_pred_hyp = discriminator_hyp(discriminator_input_hyp)

                if decoder_hyp_rmse < yearly_model_rmse_hyp:
                    logger.info(f"decoder_hyp is better in month {i + 1} in year {year} and of lake {lake_id}:")
                    label_hyp = 1
                    best_pred_hyp[:, i*30:(i+1)*30] = updated_pred_hyp
                    decoder_better_count_hyp[i] += 1
                else:
                    label_hyp = 0
                    logger.info(f"yearly model hyp is better in month {i + 1} in year {year} and of lake {lake_id}:")
                ###########


                if discriminate_pred_hyp > optimal_threshold_hyp and decoder_hyp_sim_rmse < yearly_model_sim_rmse_hyp:
                    pred_hyp[:, i*30:(i+1)*30] = updated_pred_hyp
                    decoder_discriminator_count_hyp[i] += 1

                label_hyp = torch.tensor([label_hyp], dtype=torch.float32, device=device)

                all_true_labels_hyp.append(label_hyp.item())
                all_predicted_probs_hyp.append(discriminate_pred_hyp.item())

                total_months_count_hyp[i] += 1
        
        lake_ids.append(lake_id[0])
        years.append(year.cpu().item())

        test_epi_pred.append(pred_epi.squeeze().cpu().detach())
        test_epi_label.append(y_epi.squeeze().cpu().detach())
        test_hyp_pred.append(pred_hyp.squeeze().cpu().detach())
        test_hyp_label.append(y_hyp.squeeze().cpu().detach())
        test_sim_epi.append(sim_epi.squeeze().cpu().detach())
        test_sim_hyp.append(sim_hyp.squeeze().cpu().detach())
        test_original_epi_pred.append(original_pred_epi.squeeze().cpu().detach())
        test_original_hyp_pred.append(original_pred_hyp.squeeze().cpu().detach())
        test_best_epi_pred.append(best_pred_epi.squeeze().cpu().detach())
        test_best_hyp_pred.append(best_pred_hyp.squeeze().cpu().detach())


    logger.info("Decoder Better Count for each month (Epi):", decoder_better_count_epi)
    logger.info("Total Count for discriminator thinks better for each month (Epi):", decoder_discriminator_count_epi)
    logger.info("Total Count for each month (Epi):", total_months_count_epi)
    logger.info("Decoder Better Count for each month (Hyp):", decoder_better_count_hyp)
    logger.info("Total Count for discriminator thinks better for each month (Hyp):", decoder_discriminator_count_hyp)
    logger.info("Total Count for each month (Hyp):", total_months_count_hyp)

    logger.info(len(lake_ids), len(years), len(test_epi_label), len(test_epi_pred), len(test_original_epi_pred))

    test_epi_pred = torch.cat(test_epi_pred, dim=0)
    test_epi_label = torch.cat(test_epi_label, dim=0)
    test_hyp_pred = torch.cat(test_hyp_pred, dim=0)
    test_hyp_label = torch.cat(test_hyp_label, dim=0)
    test_sim_epi = torch.cat(test_sim_epi, dim=0)
    test_sim_hyp = torch.cat(test_sim_hyp, dim=0)
    test_original_epi_pred = torch.cat(test_original_epi_pred, dim=0)
    test_original_hyp_pred = torch.cat(test_original_hyp_pred, dim=0)
    test_best_epi_pred = torch.cat(test_best_epi_pred, dim=0)
    test_best_hyp_pred = torch.cat(test_best_hyp_pred, dim=0)

    test_epi_RMSE = RMSE(test_epi_pred.to(device), test_epi_label.to(device))
    test_hyp_RMSE = RMSE(test_hyp_pred.to(device), test_hyp_label.to(device))
    test_origial_epi_RMSE = RMSE(test_original_epi_pred.to(device), test_epi_label.to(device))
    test_original_hyp_RMSE = RMSE(test_original_hyp_pred.to(device), test_hyp_label.to(device))
    test_best_epi_RMSE = RMSE(test_best_epi_pred.to(device), test_epi_label.to(device))
    test_best_hyp_RMSE = RMSE(test_best_hyp_pred.to(device), test_hyp_label.to(device))

    logger.info(f'Test RMSE Epi: {test_epi_RMSE:.4f}')
    logger.info(f'Test RMSE Hyp: {test_hyp_RMSE:.4f}')
    logger.info(f'Test Original RMSE Epi: {test_origial_epi_RMSE:.4f}')
    logger.info(f'Test Original RMSE Hyp: {test_original_hyp_RMSE:.4f}')
    logger.info(f'Test Best RMSE Epi: {test_best_epi_RMSE:.4f}')
    logger.info(f'Test Best RMSE Hyp: {test_best_hyp_RMSE:.4f}')

    days_per_year = 360
    lake_ids_expanded = np.repeat(lake_ids, days_per_year)
    years_expanded = np.repeat(years, days_per_year)


    df_epi = pd.DataFrame({
        "Lake_ID": lake_ids_expanded,
        "Year": years_expanded,
        "Day": np.tile(np.arange(1, days_per_year + 1), len(lake_ids)),  # Add Day index (1 to 360)
        "True_Label_Epi": test_epi_label.cpu().numpy(),
        "Prediction_Epi": test_epi_pred.cpu().numpy(),
        "Sim_Epi": test_sim_epi.cpu().numpy(),
        "Original_Prediction_Epi": test_original_epi_pred.cpu().numpy(),
        "Best_Prediction_Epi": test_best_epi_pred.cpu().numpy()
    })

    # Create a DataFrame for Hyp
    df_hyp = pd.DataFrame({
        "Lake_ID": lake_ids_expanded,
        "Year": years_expanded,
        "Day": np.tile(np.arange(1, days_per_year + 1), len(lake_ids)),  # Add Day index (1 to 360)
        "True_Label_Hyp": test_hyp_label.cpu().numpy(),
        "Prediction_Hyp": test_hyp_pred.cpu().numpy(),
        "Sim_Hyp": test_sim_hyp.cpu().numpy(),
        "Original_Prediction_Hyp": test_original_hyp_pred.cpu().numpy(),
        "Best_Prediction_Hyp": test_best_hyp_pred.cpu().numpy()
    })

    # Save DataFrames to CSV files
    df_epi.to_csv(os.path.join(save_path, "epi_predictions_test_sim.csv"), index=False)
    df_hyp.to_csv(os.path.join(save_path, "hyp_predictions_test_sim.csv"), index=False)

def evaluate_and_save_rmse(epi_csv_path, hyp_csv_path, save_dir, epi_col='Prediction_Epi', hyp_col='Prediction_Hyp', logger=None):
    """
    Load epi and hyp CSV files, compute RMSEs per lake, and save results.

    Args:
        epi_csv_path (str): Path to epi prediction CSV.
        hyp_csv_path (str): Path to hyp prediction CSV.
        save_dir (str): Directory to save RMSE CSV.
        epi_col (str): Column name for epi prediction.
        hyp_col (str): Column name for hyp prediction.
    """

    def to_tensor(x):
        if isinstance(x, pd.Series):
            return torch.tensor(x.values, dtype=torch.float32)
        elif isinstance(x, np.ndarray):
            return torch.from_numpy(x).float()
        elif isinstance(x, torch.Tensor):
            return x
        else:
            raise TypeError(f"Unsupported data type: {type(x)}")

    def RMSE(yhat, y):
        mask = (y != -11)
        squared_errors = torch.square((yhat - y) * mask)
        mse = torch.sum(squared_errors) / torch.sum(mask)
        return torch.sqrt(mse)

    def cal_rmse(epi_pred, hyp_pred, sim_epi, sim_hyp, obs_epi, obs_hyp):
        mixed_mask = torch.eq(sim_epi, sim_hyp)

        epi_rmse = RMSE(epi_pred[~mixed_mask], obs_epi[~mixed_mask]) if torch.sum(~mixed_mask) > 0 else torch.tensor(float('nan'))
        hyp_rmse = RMSE(hyp_pred[~mixed_mask], obs_hyp[~mixed_mask]) if torch.sum(~mixed_mask) > 0 else torch.tensor(float('nan'))
        total_rmse = RMSE(epi_pred[mixed_mask], obs_epi[mixed_mask]) if torch.sum(mixed_mask) > 0 else torch.tensor(float('nan'))

        return epi_rmse.item(), hyp_rmse.item(), total_rmse.item()

    # Load and clean data
    epi_df = pd.read_csv(epi_csv_path)
    hyp_df = pd.read_csv(hyp_csv_path)
    epi_df['Lake_ID'] = epi_df['Lake_ID'].astype(str).str.strip().str.lower()
    hyp_df['Lake_ID'] = hyp_df['Lake_ID'].astype(str).str.strip().str.lower()

    # Merge on common keys
    df = pd.merge(epi_df, hyp_df, on=['Lake_ID', 'Year', 'Day'], suffixes=('_epi', '_hyp'))

    results = []

    for lake_id, group in df.groupby('Lake_ID'):
        epi_pred = to_tensor(group[epi_col])
        hyp_pred = to_tensor(group[hyp_col])
        sim_epi = to_tensor(group['Sim_Epi'])
        sim_hyp = to_tensor(group['Sim_Hyp'])
        obs_epi = to_tensor(group['True_Label_Epi'])
        obs_hyp = to_tensor(group['True_Label_Hyp'])

        epi_rmse, hyp_rmse, total_rmse = cal_rmse(epi_pred, hyp_pred, sim_epi, sim_hyp, obs_epi, obs_hyp)

        results.append({
            'Lake_ID': lake_id,
            'Epi_RMSE': epi_rmse,
            'Hyp_RMSE': hyp_rmse,
            'Total_RMSE': total_rmse
        })

    # Create and save results DataFrame
    results_df = pd.DataFrame(results)
    avg_rmse = results_df[['Epi_RMSE', 'Hyp_RMSE', 'Total_RMSE']].mean().to_dict()
    avg_rmse['Lake_ID'] = 'Average'
    results_df = pd.concat([results_df, pd.DataFrame([avg_rmse])], ignore_index=True)

    # Construct dynamic filename
    safe_epi_col = epi_col.lower().replace(" ", "_")
    safe_hyp_col = hyp_col.lower().replace(" ", "_")
    out_filename = f'lake_rmse_results_{safe_epi_col}_{safe_hyp_col}.csv'
    out_path = os.path.join(save_dir, out_filename)

    results_df.to_csv(out_path, index=False)
    logger.info(results_df)
    logger.info(f"Results saved to: {out_path}")


if __name__ == '__main__':
    main()