#!/usr/bin/env python3
"""
Script to filter glycan monosaccharide composition data.
Keeps only rows where 'aldi', 'Xxx', and 'X' columns all have value 0.
"""

import pandas as pd
import os
from datetime import datetime

def generate_definition(row):
    """
    Generate the DEFINITION text for a glycan based on its type and composition.
    
    Args:
        row: A pandas Series containing glytoucan_type and monosaccharide counts
    
    Returns:
        str: The formatted DEFINITION text
    """
    # Mapping for monosaccharide types
    mono_mapping = {
        'Hex': {'definition': 'hexosyl', 'composition': 'Hex'},
        'HexNAc': {'definition': 'acetaminohexosyl', 'composition': 'HexNAc'},
        'dHex': {'definition': '6-deoxyhexosyl', 'composition': 'dHex'},
        'NeuAc': {'definition': 'N-acetyl-neuraminic acid', 'composition': 'NeuAc'},
        'NeuGc': {'definition': 'N-glycolyl-neuraminic acid', 'composition': 'NeuGc'},
        'HexA': {'definition': 'hexuronic acid', 'composition': 'HexA'},
        'HexN': {'definition': 'hexosamine', 'composition': 'HexN'},
        'S': {'definition': 'sulfate', 'composition': 'Sulpho'},
        'P': {'definition': 'phosphoryl', 'composition': 'Phospho'}
    }
    
    # Templates based on glytoucan_type
    templates = {
        'Topology': {
            'prefix': 'A partially-defined glycan consisting of',
            'suffix': '(identities are known but connectivity is partially known).'
        },
        'BaseComposition': {
            'prefix': 'A composition-only glycan consisting of',
            'suffix': '(identities and connectivity unknown).'
        },
        'Composition': {
            'prefix': 'A composition-only glycan consisting of',
            'suffix': '(identities are either known or partially known, and connectivity is unknown).'
        },
        'Saccharide': {
            'prefix': 'A fully or partially-defined glycan consisting of',
            'suffix': '(identities are known and connectivity is known or partially known).'
        }
    }
    
    glycan_type = row['glytoucan_type']
    
    if glycan_type not in templates:
        return ''
    
    # Collect non-zero monosaccharides
    components_def = []  # For definition text
    components_comp = []  # For composition text
    
    for mono, mapping in mono_mapping.items():
        count = int(row[mono])
        if count > 0:
            # For definition part
            if count == 1:
                components_def.append(f"one {mapping['definition']} group")
            else:
                components_def.append(f"{count} {mapping['definition']} groups")
            
            # For composition part
            components_comp.append(f"{mapping['composition']}({count})")
    
    if not components_def:
        return ''
    
    # Join components with commas and "and"
    if len(components_def) == 1:
        definition_text = components_def[0]
    elif len(components_def) == 2:
        definition_text = f"{components_def[0]} and {components_def[1]}"
    else:
        definition_text = ', '.join(components_def[:-1]) + f" and {components_def[-1]}"
    
    # Build final definition
    template = templates[glycan_type]
    composition_text = ''.join(components_comp)
    
    definition = f"{template['prefix']} {definition_text} {template['suffix']} Composition: {composition_text}."
    
    return definition

def generate_relationship(glytoucan_type):
    """
    Generate the RELATIONSHIP (ISA number) based on glytoucan_type.
    
    Args:
        glytoucan_type: The type of glycan (Topology, BaseComposition, Composition, Saccharide)
    
    Returns:
        str: The ISA number
    """
    relationship_mapping = {
        'BaseComposition': 'ISA167481',
        'Topology': 'ISA167503',
        'Composition': 'ISA167502',
        'Saccharide': 'ISA167559'
    }
    
    return relationship_mapping.get(glytoucan_type, '')

# Define file paths
# Data files are in the reviewed datasets directory
data_dir = "/data/shared/glygen/releases/data/current/reviewed/"
input_file = os.path.join(data_dir, "glycan_monosaccharide_composition.csv")
chebi_file = os.path.join(data_dir, "glycan_xref_chebi.csv")
pubchem_file = os.path.join(data_dir, "glycan_xref_pubchem.csv")
masterlist_file = os.path.join(data_dir, "glycan_masterlist.csv")
wurcs_file = os.path.join(data_dir, "glycan_sequences_wurcs.csv")
iupac_file = os.path.join(data_dir, "glycan_sequences_iupac_extended.csv")

# Images are in the shared releases directory
images_folder = "/data/shared/glygen/releases/data/current/glycanimages_snfg_svg"

# Citation files are also in the reviewed datasets directory
citations_pattern = os.path.join(data_dir, "glycan_citations_*.csv")

# Output directory with date-based folder
today = datetime.now().strftime("%Y_%m_%d")
output_dir = os.path.join("/data/projects/glygen/downloads/chebi_bulk_sub", today)

# Create output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)

output_file = os.path.join(output_dir, "glycans_for_chebi.csv")

print(f"Output will be saved to: {output_file}")

# Read the CSV file
print(f"Reading data from: {input_file}")
if not os.path.exists(input_file):
    raise FileNotFoundError(f"ERROR: Input file not found: {input_file}")

df = pd.read_csv(input_file)

if len(df) == 0:
    raise ValueError(f"ERROR: Input file is empty: {input_file}")

# Display initial information
print(f"Initial dataset shape: {df.shape}")
print(f"Total rows before filtering: {len(df)}")

# Filter rows: keep only those where aldi, Xxx, and X columns all equal 0
# Note: aldi and X are read as integers, Xxx is read as string
filtered_df = df[(df['aldi'] == 0) & (df['Xxx'] == '0') & (df['X'] == 0)]

# Display filtering results
rows_removed = len(df) - len(filtered_df)
print(f"\nFiltering complete!")
print(f"Rows removed (aldi/Xxx/X filter): {rows_removed}")
print(f"Rows remaining: {len(filtered_df)}")
print(f"Filtered dataset shape: {filtered_df.shape}")

if rows_removed == 0:
    print("WARNING: aldi/Xxx/X filter did not remove any rows. Check if filter criteria are correct.")

if len(filtered_df) == 0:
    raise ValueError("ERROR: No rows remaining after aldi/Xxx/X filtering. All glycans were filtered out.")

# Remove glycans that have ChEBI IDs
print(f"\nReading ChEBI cross-reference data from: {chebi_file}")
if not os.path.exists(chebi_file):
    raise FileNotFoundError(f"ERROR: ChEBI file not found: {chebi_file}")

chebi_df = pd.read_csv(chebi_file)

if len(chebi_df) == 0:
    raise ValueError(f"ERROR: ChEBI file is empty: {chebi_file}")

print(f"ChEBI dataset shape: {chebi_df.shape}")

# Get list of glytoucan_ac that have ChEBI IDs
chebi_ids = set(chebi_df['glytoucan_ac'].unique())
print(f"Number of unique glytoucan_ac with ChEBI IDs: {len(chebi_ids)}")

# Remove rows where glytoucan_ac is in the ChEBI list
initial_count = len(filtered_df)
filtered_df = filtered_df[~filtered_df['glytoucan_ac'].isin(chebi_ids)]
chebi_removed = initial_count - len(filtered_df)

print(f"Rows removed (ChEBI filter): {chebi_removed}")
print(f"Rows remaining after ChEBI filter: {len(filtered_df)}")

if chebi_removed == 0:
    print("WARNING: ChEBI filter did not remove any rows. This may be unexpected.")

if len(filtered_df) == 0:
    raise ValueError("ERROR: No rows remaining after ChEBI filtering. All glycans have ChEBI IDs.")

# Display final filtering results
total_removed = len(df) - len(filtered_df)
print(f"\n{'='*60}")
print(f"SUMMARY AFTER INITIAL FILTERS:")
print(f"Total rows removed: {total_removed}")
print(f"Total rows remaining: {len(filtered_df)}")
print(f"{'='*60}")

# Add PubChem IDs to the filtered dataset
print(f"\nReading PubChem cross-reference data from: {pubchem_file}")
if not os.path.exists(pubchem_file):
    raise FileNotFoundError(f"ERROR: PubChem file not found: {pubchem_file}")

pubchem_df = pd.read_csv(pubchem_file)

if len(pubchem_df) == 0:
    raise ValueError(f"ERROR: PubChem file is empty: {pubchem_file}")

print(f"PubChem dataset shape: {pubchem_df.shape}")

# Rename columns for clarity
pubchem_df = pubchem_df.rename(columns={
    'xref_id': 'pubchem_id',
    'xref_key': 'pubchem_type'
})

# Merge with filtered dataset
# Using 'inner' join to keep only glycans that have PubChem IDs
print(f"\nMerging PubChem data with filtered glycans...")
before_pubchem = len(filtered_df)
merged_df = filtered_df.merge(
    pubchem_df[['glytoucan_ac', 'pubchem_id', 'pubchem_type']], 
    on='glytoucan_ac', 
    how='inner'
)

print(f"Rows after PubChem merge: {len(merged_df)}")
print(f"Rows removed (no PubChem ID): {before_pubchem - len(merged_df)}")

if len(merged_df) == 0:
    raise ValueError("ERROR: No rows remaining after PubChem merge. No glycans have PubChem IDs.")

# Check for duplicates (glycans with both compound and substance)
duplicates = merged_df.groupby('glytoucan_ac').size()
glycans_with_multiple = (duplicates > 1).sum()
print(f"Glycans with multiple PubChem entries: {glycans_with_multiple}")

# Count by pubchem_type
type_counts = merged_df['pubchem_type'].value_counts()
print(f"\nPubChem type distribution:")
for ptype, count in type_counts.items():
    print(f"  {ptype}: {count}")

filtered_df = merged_df

# Final summary
print(f"\n{'='*60}")
print(f"SUMMARY BEFORE PUBMED FILTERING:")
print(f"Total rows in dataset: {len(filtered_df)}")
print(f"Dataset shape: {filtered_df.shape}")
print(f"{'='*60}")

# Collect PubMed IDs from all glycan_citations_*.csv files
print(f"\nCollecting PubMed IDs from citation files...")
import glob

# Get all citation files but exclude .stat.csv files
all_citation_files = glob.glob(citations_pattern)
citation_files = [f for f in all_citation_files if not f.endswith('.stat.csv')]

if len(citation_files) == 0:
    raise FileNotFoundError(f"ERROR: No citation files found matching '{citations_pattern}' (excluding .stat.csv files)")

print(f"Found {len(citation_files)} citation files:")
for f in citation_files:
    print(f"  - {os.path.basename(f)}")

# Dictionary to store pubmed_ids for each glytoucan_ac
pubmed_dict = {}

# Track stats per file
file_stats = {}

for citation_file in citation_files:
    file_basename = os.path.basename(citation_file)
    print(f"\nProcessing {file_basename}...")
    
    file_stats[file_basename] = {
        'rows_loaded': 0,
        'pubmed_entries': 0,
        'unique_glycans': 0,
        'unique_pubmed_ids': 0,
        'error': None
    }
    
    try:
        # Try reading with UTF-8 encoding (most common)
        try:
            citation_df = pd.read_csv(citation_file, encoding='utf-8')
        except UnicodeDecodeError:
            # Fallback to latin-1 if UTF-8 fails
            print(f"  UTF-8 encoding failed, trying latin-1...")
            citation_df = pd.read_csv(citation_file, encoding='latin-1')
        
        file_stats[file_basename]['rows_loaded'] = len(citation_df)
        
        if len(citation_df) == 0:
            print(f"  WARNING: File is empty")
            continue
        
        print(f"  File loaded successfully. Total rows: {len(citation_df)}")
        print(f"  Columns: {list(citation_df.columns)}")
        
        # Check for required columns
        if 'xref_key' not in citation_df.columns:
            print(f"  WARNING: Missing 'xref_key' column. Skipping file.")
            file_stats[file_basename]['error'] = "Missing xref_key column"
            continue
        
        if 'xref_id' not in citation_df.columns:
            print(f"  WARNING: Missing 'xref_id' column. Skipping file.")
            file_stats[file_basename]['error'] = "Missing xref_id column"
            continue
        
        if 'glytoucan_ac' not in citation_df.columns:
            print(f"  WARNING: Missing 'glytoucan_ac' column. Skipping file.")
            file_stats[file_basename]['error'] = "Missing glytoucan_ac column"
            continue
        
        # Filter for rows where xref_key is 'glycan_xref_pubmed'
        pubmed_rows = citation_df[citation_df['xref_key'] == 'glycan_xref_pubmed']
        file_stats[file_basename]['pubmed_entries'] = len(pubmed_rows)
        print(f"  Found {len(pubmed_rows)} PubMed entries")
        
        if len(pubmed_rows) == 0:
            unique_keys = citation_df['xref_key'].unique()
            print(f"  Unique xref_key values in file: {unique_keys}")
            continue
        
        unique_glycans_in_file = pubmed_rows['glytoucan_ac'].nunique()
        unique_pubmed_in_file = pubmed_rows['xref_id'].nunique()
        file_stats[file_basename]['unique_glycans'] = unique_glycans_in_file
        file_stats[file_basename]['unique_pubmed_ids'] = unique_pubmed_in_file
        print(f"  Unique glycans with PubMed IDs: {unique_glycans_in_file}")
        print(f"  Unique PubMed IDs: {unique_pubmed_in_file}")
        
        # Group by glytoucan_ac and collect unique pubmed IDs
        for glytoucan_ac, group in pubmed_rows.groupby('glytoucan_ac'):
            unique_pubmed_ids = group['xref_id'].unique().tolist()
            
            if glytoucan_ac not in pubmed_dict:
                pubmed_dict[glytoucan_ac] = set()
            
            # Add unique pubmed IDs to the set
            pubmed_dict[glytoucan_ac].update(unique_pubmed_ids)
    
    except Exception as e:
        print(f"  ERROR processing file: {e}")
        file_stats[file_basename]['error'] = str(e)
        # Don't raise - continue processing other files
        continue

# Print summary of all citation files
print(f"\n{'='*60}")
print(f"CITATION FILE SUMMARY:")
print(f"{'='*60}")
for fname, stats in file_stats.items():
    print(f"\n{fname}:")
    print(f"  Rows loaded: {stats['rows_loaded']}")
    print(f"  PubMed entries: {stats['pubmed_entries']}")
    print(f"  Unique glycans: {stats['unique_glycans']}")
    print(f"  Unique PubMed IDs: {stats['unique_pubmed_ids']}")
    if stats['error']:
        print(f"  ERROR: {stats['error']}")

print(f"\nTotal glycans with PubMed IDs: {len(pubmed_dict)}")

if len(pubmed_dict) == 0:
    raise ValueError("ERROR: No PubMed IDs found in any citation files. Check that citation files contain 'glycan_xref_pubmed' entries.")

# Convert sets to pipe-separated strings (convert IDs to strings first)
pubmed_series = pd.Series({k: '|'.join(sorted([str(id) for id in v])) for k, v in pubmed_dict.items()})

# Add pubmed_ids column to filtered_df
filtered_df['pubmed_ids'] = filtered_df['glytoucan_ac'].map(pubmed_series)

# Count glycans with and without PubMed IDs
with_pubmed = filtered_df['pubmed_ids'].notna().sum()
without_pubmed = filtered_df['pubmed_ids'].isna().sum()

print(f"\nGlycans with PubMed IDs: {with_pubmed}")
print(f"Glycans without PubMed IDs (will have empty citations): {without_pubmed}")
print(f"Total glycans (with and without citations): {len(filtered_df)}")

# NOTE: Not removing glycans without citations - they will have empty pubmed_ids column

# Filter for glycans with SVG images
# Filter for glycans with SVG images
print(f"\n{'='*60}")
print(f"IMAGE FILTERING:")
print(f"{'='*60}")
print(f"Checking for SVG images in: {images_folder}")
print(f"Folder exists: {os.path.exists(images_folder)}")

if not os.path.exists(images_folder):
    raise FileNotFoundError(f"ERROR: Images folder not found at '{images_folder}'. Please create the folder and add SVG files.")

# Get all SVG files in the folder
svg_files = glob.glob(os.path.join(images_folder, "*.svg"))
print(f"Found {len(svg_files)} SVG files")

if len(svg_files) == 0:
    raise ValueError(f"ERROR: No SVG files found in images folder: {images_folder}")

if len(svg_files) > 0:
    print(f"First few SVG files found:")
    for f in svg_files[:5]:
        print(f"  - {os.path.basename(f)}")

# Extract glytoucan_ac from filenames (remove .svg extension)
glycans_with_images = set([os.path.splitext(os.path.basename(f))[0] for f in svg_files])
print(f"Unique glycans with images: {len(glycans_with_images)}")

if len(glycans_with_images) > 0:
    print(f"First few glycan IDs from images:")
    for gid in list(glycans_with_images)[:5]:
        print(f"  - {gid}")

# Filter dataset to only include glycans with images
before_image_filter = len(filtered_df)
print(f"\nGlycans before image filter: {before_image_filter}")
filtered_df = filtered_df[filtered_df['glytoucan_ac'].isin(glycans_with_images)]
removed_no_image = before_image_filter - len(filtered_df)

print(f"Glycans removed (no image): {removed_no_image}")
print(f"Glycans remaining: {len(filtered_df)}")

if removed_no_image == 0:
    print("WARNING: Image filter did not remove any rows. All glycans in dataset have images.")

if len(filtered_df) == 0:
    raise ValueError("ERROR: No rows remaining after image filtering. No glycans in the dataset have corresponding SVG images.")

# Add glytoucan_type from glycan_masterlist
print(f"\nReading glycan masterlist from: {masterlist_file}")
if not os.path.exists(masterlist_file):
    raise FileNotFoundError(f"ERROR: Masterlist file not found: {masterlist_file}")

masterlist_df = pd.read_csv(masterlist_file)

if len(masterlist_df) == 0:
    raise ValueError(f"ERROR: Masterlist file is empty: {masterlist_file}")

print(f"Masterlist dataset shape: {masterlist_df.shape}")

# Select only the columns we need
type_mapping = masterlist_df[['glytoucan_ac', 'glytoucan_type']].drop_duplicates()

# Merge glytoucan_type
filtered_df = filtered_df.merge(
    type_mapping,
    on='glytoucan_ac',
    how='left'
)

print(f"Added glytoucan_type column")

# Check for any glycans without glytoucan_type
missing_type = filtered_df['glytoucan_type'].isna().sum()
if missing_type > 0:
    print(f"WARNING: {missing_type} glycans do not have glytoucan_type in masterlist")
    print(f"Removing glycans without glytoucan_type...")
    filtered_df = filtered_df[filtered_df['glytoucan_type'].notna()]
    print(f"Glycans remaining after removing those without type: {len(filtered_df)}")

if len(filtered_df) == 0:
    raise ValueError("ERROR: No rows remaining after removing glycans without glytoucan_type.")

# Reorder columns to make glytoucan_type the second column
cols = filtered_df.columns.tolist()
# Remove glytoucan_type from wherever it is
cols.remove('glytoucan_type')
# Insert it as the second column (after glytoucan_ac)
cols.insert(1, 'glytoucan_type')
filtered_df = filtered_df[cols]

print(f"Reordered columns - glytoucan_type is now the second column")

# Add WURCS sequences
print(f"\nReading WURCS sequences from: {wurcs_file}")
if not os.path.exists(wurcs_file):
    raise FileNotFoundError(f"ERROR: WURCS file not found: {wurcs_file}")

wurcs_df = pd.read_csv(wurcs_file)

if len(wurcs_df) == 0:
    raise ValueError(f"ERROR: WURCS file is empty: {wurcs_file}")

print(f"WURCS dataset shape: {wurcs_df.shape}")

# Merge WURCS sequences
filtered_df = filtered_df.merge(
    wurcs_df[['glytoucan_ac', 'sequence_wurcs']],
    on='glytoucan_ac',
    how='left'
)

wurcs_missing = filtered_df['sequence_wurcs'].isna().sum()
wurcs_present = filtered_df['sequence_wurcs'].notna().sum()
print(f"Glycans with WURCS: {wurcs_present}")
print(f"Glycans without WURCS: {wurcs_missing}")

# Add IUPAC extended sequences
print(f"\nReading IUPAC extended sequences from: {iupac_file}")
if not os.path.exists(iupac_file):
    raise FileNotFoundError(f"ERROR: IUPAC file not found: {iupac_file}")

iupac_df = pd.read_csv(iupac_file)

if len(iupac_df) == 0:
    raise ValueError(f"ERROR: IUPAC file is empty: {iupac_file}")

print(f"IUPAC dataset shape: {iupac_df.shape}")

# Merge IUPAC sequences
filtered_df = filtered_df.merge(
    iupac_df[['glytoucan_ac', 'sequence_iupac_extended']],
    on='glytoucan_ac',
    how='left'
)

iupac_missing = filtered_df['sequence_iupac_extended'].isna().sum()
iupac_present = filtered_df['sequence_iupac_extended'].notna().sum()
print(f"Glycans with IUPAC: {iupac_present}")
print(f"Glycans without IUPAC: {iupac_missing}")

# Generate DEFINITION for each glycan
print(f"\nGenerating DEFINITION text for each glycan...")
filtered_df['definition'] = filtered_df.apply(generate_definition, axis=1)
print(f"DEFINITION generated for all glycans")

# Check for any failures
empty_definitions = (filtered_df['definition'] == '').sum()
if empty_definitions > 0:
    print(f"WARNING: {empty_definitions} glycans have empty DEFINITION (likely unknown glytoucan_type)")

# Generate RELATIONSHIP (ISA number) for each glycan
print(f"\nGenerating RELATIONSHIP (ISA number) for each glycan...")
filtered_df['relationship'] = filtered_df['glytoucan_type'].apply(generate_relationship)
print(f"RELATIONSHIP generated for all glycans")

# Check relationship distribution
relationship_counts = filtered_df['relationship'].value_counts()
print(f"\nRELATIONSHIP distribution:")
for rel, count in relationship_counts.items():
    print(f"  {rel}: {count}")

# Final summary
print(f"\n{'='*60}")
print(f"FINAL SUMMARY:")
print(f"Total rows in final dataset: {len(filtered_df)}")
print(f"Final dataset shape: {filtered_df.shape}")
print(f"{'='*60}")

# Save the filtered dataset
print(f"\nSaving filtered data to: {output_file}")
try:
    filtered_df.to_csv(output_file, index=False)
    print("Done!")
except PermissionError:
    print(f"\nERROR: Permission denied when trying to save the file.")
    print(f"The file may be open in Excel or another program.")
    print(f"Please close the file and run the script again.")
    print(f"\nAlternatively, saving to a backup file with timestamp...")
    
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = output_file.replace('.csv', f'_backup_{timestamp}.csv')
    filtered_df.to_csv(backup_file, index=False)
    print(f"Saved to backup file: {backup_file}")
    print("Done!")
except Exception as e:
    print(f"\nERROR: Failed to save file: {e}")
    raise

# Display a preview of the filtered data
print("\nPreview of filtered data (first 10 rows):")
print(filtered_df.head(10))