import torch
from torch.utils.data import Dataset, DataLoader
from transformers import LongformerTokenizer, LongformerModel
import pandas as pd
import logging
import argparse
import numpy as np
from tqdm import tqdm
import time  # For time tracking
import os


# Set up logging configuration
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s:%(message)s')

class ClinicalNotesDataset(Dataset):
    def __init__(self, events_df, tokenizer, max_len=4096):
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.notes = []
        self.note_ids = []
        self.field_types = []  # To track whether the chunk is from 'value' or 'feature'

        logging.info("Processing notes and splitting into sub-chunks...")
        start_time = time.time()

        for idx, row in tqdm(events_df.iterrows(), total=events_df.shape[0], desc='Notes'):
            note_id = idx  # Assume index is the unique note ID

            # Process 'value' chunks
            value_chunks = row.get('value', [])
            for chunk in value_chunks:
                if isinstance(chunk, str):
                    chunk_text = chunk
                else:
                    chunk_text = str(chunk)
                # Tokenize and split if necessary
                tokens = self.tokenizer.tokenize(chunk_text)
                if len(tokens) <= self.max_len - 2:
                    self.notes.append(chunk_text)
                    self.note_ids.append(note_id)
                    self.field_types.append('value')
                else:
                    # Split tokens into sub-chunks
                    for i in range(0, len(tokens), self.max_len - 2):
                        sub_tokens = tokens[i:i + self.max_len - 2]
                        sub_chunk = self.tokenizer.convert_tokens_to_string(sub_tokens)
                        self.notes.append(sub_chunk)
                        self.note_ids.append(note_id)
                        self.field_types.append('value')

            # Process 'feature' chunks
            feature_chunks = row.get('feature', [])
            for chunk in feature_chunks:
                if isinstance(chunk, str):
                    chunk_text = chunk
                else:
                    chunk_text = str(chunk)
                # Tokenize and split if necessary
                tokens = self.tokenizer.tokenize(chunk_text)
                if len(tokens) <= self.max_len - 2:
                    self.notes.append(chunk_text)
                    self.note_ids.append(note_id)
                    self.field_types.append('feature')
                else:
                    # Split tokens into sub-chunks
                    for i in range(0, len(tokens), self.max_len - 2):
                        sub_tokens = tokens[i:i + self.max_len - 2]
                        sub_chunk = self.tokenizer.convert_tokens_to_string(sub_tokens)
                        self.notes.append(sub_chunk)
                        self.note_ids.append(note_id)
                        self.field_types.append('feature')

        logging.info(f"Total sub-chunks after splitting: {len(self.notes)}")
        end_time = time.time()
        logging.info(f"Time taken for processing and splitting notes: {end_time - start_time:.2f} seconds")

    def __len__(self):
        return len(self.notes)

    def __getitem__(self, idx):
        text = self.notes[idx]
        note_id = self.note_ids[idx]
        field_type = self.field_types[idx]
        # Tokenize the text
        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )
        input_ids = encoding['input_ids'].squeeze(0)  # Remove batch dimension
        attention_mask = encoding['attention_mask'].squeeze(0)
        # Generate global attention mask
        global_attention_mask = torch.zeros_like(attention_mask)
        global_attention_mask[0] = 1  # Set global attention on 
 token
        return {
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'global_attention_mask': global_attention_mask,
            'note_id': note_id,
            'field_type': field_type
        }

def collate_fn(batch):
    input_ids = torch.stack([item['input_ids'] for item in batch])
    attention_mask = torch.stack([item['attention_mask'] for item in batch])
    global_attention_mask = torch.stack([item['global_attention_mask'] for item in batch])
    note_ids = [item['note_id'] for item in batch]
    field_types = [item['field_type'] for item in batch]
    return {
        'input_ids': input_ids,
        'attention_mask': attention_mask,
        'global_attention_mask': global_attention_mask,
        'note_ids': note_ids,
        'field_types': field_types
    }

def main():
    parser = argparse.ArgumentParser(description='Embed clinical notes using Clinical Longformer')
    parser.add_argument('--events_path', type=str, required=True, help='Path to the events pickle file')
    parser.add_argument('--batch_size', type=int, default=2, help='Batch size for DataLoader')
    parser.add_argument('--num_workers', type=int, default=4, help='Number of workers for DataLoader')
    parser.add_argument('--max_len', type=int, default=4096, help='Maximum sequence length for tokenizer')
    parser.add_argument('--output_dir', type=str, default='.', help='Directory to save outputs')
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f"Using device: {device}")

    # Load the tokenizer and model
    logging.info("Loading tokenizer and model...")
    start_time = time.time()
    tokenizer = LongformerTokenizer.from_pretrained('yikuan8/Clinical-Longformer')
    model = LongformerModel.from_pretrained('yikuan8/Clinical-Longformer')
    model.to(device)
    end_time = time.time()
    logging.info(f"Time taken to load tokenizer and model: {end_time - start_time:.2f} seconds")

    # Use DataParallel if multiple GPUs are available
    if torch.cuda.device_count() > 1:
        logging.info(f"Using {torch.cuda.device_count()} GPUs")
        model = torch.nn.DataParallel(model)
    else:
        logging.info(f"Using single GPU or CPU")

    # Load events DataFrame
    if args.events_path is not None:
        logging.info(f"Loading events from {args.events_path}...")
        start_time = time.time()
        events_df = pd.read_pickle(args.events_path)
        end_time = time.time()
        logging.info(f"Time taken to load events DataFrame: {end_time - start_time:.2f} seconds")
    else:
        raise ValueError('events_path must be provided')

    # Rename 'value' and 'feature' columns
    events_df.rename(columns={'value': 'text_chunked', 'feature': 'cat_desc'}, inplace=True)

    # Create dataset and dataloader
    dataset = ClinicalNotesDataset(events_df, tokenizer, max_len=args.max_len)
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        collate_fn=collate_fn
    )

    # Switch model to evaluation mode
    model.eval()

    # Dictionaries to hold embeddings for 'value' and 'feature'
    embeddings_dict_value = {}
    embeddings_dict_feature = {}
    total_batches = len(dataloader)
    logging.info(f"Total batches: {total_batches}")

    start_time_embedding = time.time()
    with torch.no_grad():
        for batch_idx, batch in enumerate(tqdm(dataloader, desc='Embedding', total=total_batches)):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            global_attention_mask = batch['global_attention_mask'].to(device)
            note_ids = batch['note_ids']
            field_types = batch['field_types']
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                global_attention_mask=global_attention_mask
            )
            # Take the 
 token representation as the embedding
            cls_embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy()
            # Aggregate embeddings per note_id and field_type
            for idx, (note_id, field_type) in enumerate(zip(note_ids, field_types)):
                embedding = cls_embeddings[idx]
                if field_type == 'value':
                    if note_id in embeddings_dict_value:
                        embeddings_dict_value[note_id].append(embedding)
                    else:
                        embeddings_dict_value[note_id] = [embedding]
                elif field_type == 'feature':
                    if note_id in embeddings_dict_feature:
                        embeddings_dict_feature[note_id].append(embedding)
                    else:
                        embeddings_dict_feature[note_id] = [embedding]

            # Optional: Save intermediate embeddings periodically
            # if (batch_idx + 1) % 1000 == 0:
            #     logging.info(f"Saving intermediate embeddings at batch {batch_idx + 1}")
            #     # Save intermediate embeddings to avoid data loss in case of errors
            #     # Save embeddings_dict_value and embeddings_dict_feature as temporary files
            #     temp_value_path = os.path.join(args.output_dir, f'embeddings_value_temp.pkl')
            #     temp_feature_path = os.path.join(args.output_dir, f'embeddings_feature_temp.pkl')
            #     # Implement saving logic here if needed

    end_time_embedding = time.time()
    logging.info(f"Time taken for embedding: {end_time_embedding - start_time_embedding:.2f} seconds")

    # Compute the mean embedding for each note and field type
    logging.info("Aggregating embeddings for each note and field type...")

    # Process 'value' embeddings
    logging.info("Processing 'value' embeddings...")
    start_time_value = time.time()
    embeddings_df_value = pd.DataFrame({
        note_id: np.mean(embeddings_list, axis=0)
        for note_id, embeddings_list in tqdm(embeddings_dict_value.items(), desc='Aggregating Value')
    }).T  # Transpose to have note_id as index
    end_time_value = time.time()
    logging.info(f"Time taken to process 'value' embeddings: {end_time_value - start_time_value:.2f} seconds")

    # Save 'value' embeddings separately
    value_embeddings_path = os.path.join(args.output_dir, 'embeddings_value.pkl')
    embeddings_df_value.to_pickle(value_embeddings_path)
    logging.info(f"'value' embeddings saved to {value_embeddings_path}")

    # Process 'feature' embeddings
    logging.info("Processing 'feature' embeddings...")
    start_time_feature = time.time()
    embeddings_df_feature = pd.DataFrame({
        note_id: np.mean(embeddings_list, axis=0)
        for note_id, embeddings_list in tqdm(embeddings_dict_feature.items(), desc='Aggregating Feature')
    }).T  # Transpose to have note_id as index
    end_time_feature = time.time()
    logging.info(f"Time taken to process 'feature' embeddings: {end_time_feature - start_time_feature:.2f} seconds")

    # Save 'feature' embeddings separately
    feature_embeddings_path = os.path.join(args.output_dir, 'embeddings_feature.pkl')
    embeddings_df_feature.to_pickle(feature_embeddings_path)
    logging.info(f"'feature' embeddings saved to {feature_embeddings_path}")

    # Map embeddings back to events_df
    logging.info("Mapping embeddings back to events DataFrame...")
    start_time_mapping = time.time()

    # Ensure the indices are of the same type
    events_df.index = events_df.index.astype(embeddings_df_value.index.dtype)

    # Map embeddings back to events_df
    def get_embedding(embeddings_df, note_id):
        try:
            embedding = embeddings_df.loc[note_id].values
        except KeyError:
            embedding = np.zeros(model.config.hidden_size)
        return embedding

    events_df['value'] = events_df.index.map(lambda idx: get_embedding(embeddings_df_value, idx))
    events_df['feature'] = events_df.index.map(lambda idx: get_embedding(embeddings_df_feature, idx))

    end_time_mapping = time.time()
    logging.info(f"Time taken for mapping embeddings: {end_time_mapping - start_time_mapping:.2f} seconds")

    # Save the updated events_df
    output_events_path = os.path.join(args.output_dir, 'events_with_embeddings.pkl')
    events_df.to_pickle(output_events_path)
    logging.info(f"Updated events DataFrame saved to {output_events_path}")

    
if __name__ == '__main__':
    main()


any flaws or issues? ask me anything