data_preprocessing.py

import pandas as pd
import json
import re

def normalize_marker_name(name):
    """
    Normalizes a marker name by:
    1. Removing content in parentheses (e.g. "CD96 (TACTILE)" -> "CD96")
    2. Converting to lowercase
    3. Removing delimiters (spaces, dashes)
    4. Handling common suffixes
    """
    if not isinstance(name, str):
        return ""
    
    # Remove parentheses and content inside them
    name = re.sub(r'\s*\(.*?\)', '', name)
    
    name = name.lower().replace("-", "").replace(" ", "")
    if name.endswith('a') or name.endswith('b'): # Handle CD8a, CD8b
        name = name[:-1]
    return name

def parse_target_aliases(target_string):
    """
    Parses a target string like "CD274 (B7-H1, PD-L1)" and returns a list of all normalized aliases.
    Handles comma and slash separators, e.g., CD45 (LCA/T200).
    """
    if not isinstance(target_string, str):
        return []

    # Extract the main name (text before parentheses) and normalize it
    main_name = re.sub(r'\s*\(.*\)', '', target_string).strip()
    all_names = {normalize_marker_name(main_name)}

    # Extract content from within parentheses
    aliases_in_parens = re.findall(r'\((.*?)\)', target_string)
    
    # Process each group of aliases found
    for group in aliases_in_parens:
        # Split by comma or slash, and strip whitespace from each part
        aliases = [name.strip() for name in re.split(r'[,/]', group)]
        for alias in aliases:
            if alias:  # Ensure not an empty string
                all_names.add(normalize_marker_name(alias))

    return list(all_names)

def load_antibody_data(file_input, mapping_file=None, column_mapping=None):
    """
    Loads antibody inventory from a CSV file (path or buffer), maps columns if needed,
    parses target aliases, and optionally adds a System_Code based on a mapping file.
    
    Args:
        file_input: File path (str) or file-like object.
        mapping_file: Path to the channel mapping JSON.
        column_mapping: Dict mapping user columns to standard columns 
                        e.g. {'Antigen': 'Target', 'Fluor': 'Fluorescein'}
    """
    df = None

    # 1. Load spreadsheet — Excel (.xlsx) uses read_excel; CSV uses encoding detection.
    is_buffer = not isinstance(file_input, str)
    is_xlsx = isinstance(file_input, str) and file_input.lower().endswith((".xlsx", ".xls"))

    if is_xlsx:
        # Excel files are binary; no encoding probing needed.
        try:
            df = pd.read_excel(file_input)
        except Exception as e:
            print(f"Error: Failed to read Excel file: {e}")
            return None
    else:
        encodings_to_try = ['utf-8', 'gbk', 'gb18030', 'latin1']

        # If it's a file-like object, we need to be careful about seeking to 0 if we retry
        for encoding in encodings_to_try:
            try:
                if is_buffer:
                    file_input.seek(0)
                df = pd.read_csv(file_input, encoding=encoding)
                break # Success
            except (UnicodeDecodeError, pd.errors.ParserError):
                continue

    if df is None:
        print("Error: Failed to decode file with supported encodings.")
        return None

    # 2. Apply Column Mapping
    if column_mapping:
        df.rename(columns=column_mapping, inplace=True)
    
    # 3. Drop rows where Target or Fluorescein is NaN (trailing empty rows, etc.)
    df = df.dropna(subset=['Target', 'Fluorescein'])

    # 4. Validation: Check for critical columns
    required_cols = ['Target', 'Fluorescein'] 
    missing = [c for c in required_cols if c not in df.columns]
    if missing:
        print(f"Error: Missing required columns after mapping: {missing}")
        return None

    try:
        # --- NEW: Parse Aliases into a new column ---
        df['Target_Aliases'] = df['Target'].apply(parse_target_aliases)

        if mapping_file:
            with open(mapping_file, 'r', encoding='utf-8') as f:
                channel_map = json.load(f)
            
            # Normalize map keys to lowercase for robust matching
            channel_map = {k.lower(): v for k, v in channel_map.items()}

            # Use lower() for case-insensitive mapping
            df['System_Code'] = df['Fluorescein'].str.lower().map(channel_map)
            df['System_Code'] = df['System_Code'].fillna('UNKNOWN')

        return df
    except Exception as e:
        print(f"Error processing data: {e}")
        return None

def format_antibodies_for_llm(df):
    """
    Formats the antibody DataFrame into a list of dictionaries suitable for LLM input.
    Each dictionary will contain 'Fluorescein', 'Target', 'System_Code', 
    'Brand', 'Catalog Number', 'Clone', and the new 'Target_Aliases'.
    """
    if df is None:
        return []
    
    # Ensure 'Target_Aliases' column exists; create it if it doesn't
    if 'Target_Aliases' not in df.columns:
        df['Target_Aliases'] = df['Target'].apply(parse_target_aliases)
    
    # Select and rename columns for clarity for the LLM
    llm_data = df[['Fluorescein', 'Target', 'System_Code', 'Brand', 'Catalog Number', 'Clone', 'Target_Aliases']].copy()
    
    # Convert DataFrame to a list of dictionaries
    return llm_data.to_dict(orient='records')

def aggregate_antibodies_by_marker(antibody_df, brightness_data):
    """
    Aggregates antibody data by marker, simplifying the information and adding brightness.
    This is the new "information hub" for antigens.
    """
    antibodies_by_marker = {}
    marker_expression = {}

    # Ensure brightness keys are lowercase for case-insensitive matching
    brightness_data_lower = {k.lower(): v for k, v in brightness_data.items()}

    for _, row in antibody_df.iterrows():
        # The primary, un-normalized marker name from the 'Target' column
        main_marker = row['Target'].split('(')[0].strip()
        
        if not row['Target_Aliases']:
            continue

        # Store the simplified antibody info
        fluorochrome = row['Fluorescein']
        brightness = brightness_data_lower.get(fluorochrome.lower(), 3)

        stock = None
        for col in ('现有数目', 'Quantity', 'Stock', 'stock', '库存'):
            if col in row.index and pd.notna(row[col]):
                try:
                    stock = int(float(row[col]))
                except (ValueError, TypeError):
                    stock = None
                break

        antibody_info = {
            "clone": row['Clone'],
            "fluorochrome": fluorochrome,
            "brightness": brightness,
            "system_code": row.get('System_Code', 'UNKNOWN'),
            "brand": row.get('Brand', 'N/A'),
            "catalog_number": row.get('Catalog Number', 'N/A'),
            "stock": stock,
        }

        # --- FIX: Index antibody under ALL aliases ---
        # Instead of just picking the first alias, we add this antibody to the list 
        # for EVERY alias found. This creates a comprehensive inverted index.
        for alias in row['Target_Aliases']:
            if alias not in antibodies_by_marker:
                antibodies_by_marker[alias] = []
            antibodies_by_marker[alias].append(antibody_info)

    return antibodies_by_marker, marker_expression


if __name__ == "__main__":
    csv_file = "流式抗体库-20250625小鼠.csv"
    mapping_file = "channel_mapping.json"
    brightness_file = "fluorochrome_brightness.json"
    
    antibody_df = load_antibody_data(csv_file, mapping_file)

    if antibody_df is not None:
        print("Antibody data loaded successfully.")
        
        # llm_formatted_data = format_antibodies_for_llm(antibody_df)
        # print("\nFormatted data for LLM (first 2 entries):")
        # print(json.dumps(llm_formatted_data[:2], indent=2, ensure_ascii=False))

        with open(brightness_file, 'r', encoding='utf-8') as f:
            brightness_data = json.load(f)

        # --- Test the new aggregation function ---
        print("\n--- Testing New Aggregation Function ---")
        antibodies_by_marker, marker_expression = aggregate_antibodies_by_marker(antibody_df, brightness_data)
        
        print("\nMarker Expression Levels Found:")
        print(json.dumps(marker_expression, indent=2))

        print("\nAntibodies for 'cd3':")
        if 'cd3' in antibodies_by_marker:
            print(json.dumps(antibodies_by_marker['cd3'], indent=2, ensure_ascii=False))
        else:
            print("No data found for 'cd3'. Check marker normalization.")

        print("\nAntibodies for 'nk1.1':")
        if 'nk1.1' in antibodies_by_marker:
            print(json.dumps(antibodies_by_marker['nk1.1'], indent=2, ensure_ascii=False))
        else:
            print("No data found for 'nk1.1'. Check marker normalization.")