#!/usr/bin/env python3
"""
Validation script for ChEBI SDF file.
Checks that all entries contain correct data and match the source CSV.
"""

import os
import pandas as pd
import sys
import glob
from datetime import datetime
from collections import defaultdict

# PRODUCTION MODE - Files in dated folder
# Usage: python3 validate_chebi_sdf.py [YYYY_MM_DD]
# If no date provided, uses today's date, then falls back to most recent folder

base_dir = "/data/projects/glygen/downloads/chebi_bulk_sub"

# Check for command line argument
if len(sys.argv) > 1:
    date_str = sys.argv[1]
    data_dir = os.path.join(base_dir, date_str)
    if not os.path.exists(data_dir):
        print(f"ERROR: Directory not found: {data_dir}")
        sys.exit(1)
else:
    # Try today's date first
    today = datetime.now().strftime("%Y_%m_%d")
    data_dir = os.path.join(base_dir, today)
    
    # If today's folder doesn't exist, find the most recent one
    if not os.path.exists(data_dir):
        print(f"Today's folder not found ({data_dir})")
        print(f"Searching for most recent folder in {base_dir}...")
        
        # Find all date folders
        all_folders = sorted(glob.glob(os.path.join(base_dir, "20*_*_*")))
        if not all_folders:
            print(f"ERROR: No date folders found in {base_dir}")
            sys.exit(1)
        
        data_dir = all_folders[-1]  # Most recent
        print(f"Using most recent folder: {data_dir}")

sdf_file = os.path.join(data_dir, "glycans_chebi_bulk_submission_cleaned.sdf")
csv_file = os.path.join(data_dir, "glycans_for_chebi.csv")

print("="*70)
print("ChEBI SDF VALIDATION")
print("="*70)
print(f"\nValidating SDF file: {sdf_file}")
print(f"Against CSV file: {csv_file}")
print()

# Check files exist
if not os.path.exists(sdf_file):
    print(f"ERROR: SDF file not found: {sdf_file}")
    sys.exit(1)

if not os.path.exists(csv_file):
    print(f"ERROR: CSV file not found: {csv_file}")
    sys.exit(1)

# Read CSV
print("Reading CSV file...")
csv_df = pd.read_csv(csv_file)
print(f"  Loaded {len(csv_df)} entries")

# Create lookup for quick access
csv_lookup = {}
for idx, row in csv_df.iterrows():
    glytoucan_ac = row['glytoucan_ac']
    csv_lookup[glytoucan_ac] = row

print(f"  Created lookup for {len(csv_lookup)} unique glytoucan IDs")
print()

# Parse SDF file
print("Parsing SDF file...")
entries = []
current_entry = {}
current_field = None
current_field_value = []

def save_current_field(entry, field, value_lines):
    """Save the current field to the entry if it has content."""
    if field and value_lines:
        entry[field] = '\n'.join(value_lines).strip()

with open(sdf_file, 'r', encoding='utf-8', errors='ignore') as f:
    for line in f:
        line_stripped = line.rstrip('\n')
        
        # End of entry
        if line_stripped == '$$$$':
            # Save the last field before ending the entry
            save_current_field(current_entry, current_field, current_field_value)
            
            if current_entry:
                entries.append(current_entry)
            current_entry = {}
            current_field = None
            current_field_value = []
            continue
        
        # Field marker
        if line_stripped.startswith('> <') or line_stripped.startswith('>  <'):
            # Save previous field
            save_current_field(current_entry, current_field, current_field_value)
            
            # Parse field name
            field_name = line_stripped.replace('> <', '').replace('>  <', '').rstrip('>')
            current_field = field_name
            current_field_value = []
            continue
        
        # Structure block (lines before first field marker)
        if current_field is None:
            if 'STRUCTURE_BLOCK' not in current_entry:
                current_entry['STRUCTURE_BLOCK'] = []
            # Include all lines (including blank) in structure block
            current_entry['STRUCTURE_BLOCK'].append(line_stripped)
            continue
        
        # Field value
        if current_field:
            if line_stripped:  # Only add non-empty lines
                current_field_value.append(line_stripped)

# Save any remaining entry (in case file doesn't end with $$$$)
if current_entry:
    save_current_field(current_entry, current_field, current_field_value)
    entries.append(current_entry)

print(f"  Parsed {len(entries)} entries from SDF")
print()

# Validation
print("="*70)
print("VALIDATION RESULTS")
print("="*70)

errors = []
warnings = []
stats = {
    'total_entries': len(entries),
    'entries_with_definition': 0,
    'entries_with_synonym': 0,
    'entries_with_iupac': 0,
    'entries_with_reference': 0,
    'entries_with_empty_structure': 0,
    'entries_with_structure': 0,
}

# Check each entry
for entry_idx, entry in enumerate(entries):
    entry_num = entry_idx + 1
    
    # Extract ID
    if 'ID' not in entry:
        errors.append(f"Entry {entry_num}: Missing ID field")
        continue
    
    entry_id = entry['ID']
    
    # Parse glytoucan_ac from ID
    if not entry_id.startswith('TEMP_'):
        errors.append(f"Entry {entry_num}: ID '{entry_id}' does not start with 'TEMP_'")
        continue
    
    glytoucan_ac = entry_id.replace('TEMP_', '')
    
    # Check if glytoucan_ac exists in CSV
    if glytoucan_ac not in csv_lookup:
        errors.append(f"Entry {entry_num}: GlyTouCan ID '{glytoucan_ac}' not found in CSV")
        continue
    
    csv_row = csv_lookup[glytoucan_ac]
    
    # Validate each field
    
    # NAME field
    if 'NAME' not in entry:
        errors.append(f"Entry {entry_num} ({glytoucan_ac}): Missing NAME field")
    else:
        expected_name = f"GlyTouCan {glytoucan_ac}"
        if entry['NAME'] != expected_name:
            errors.append(f"Entry {entry_num} ({glytoucan_ac}): NAME '{entry['NAME']}' != expected '{expected_name}'")
    
    # DEFINITION field
    if 'DEFINITION' not in entry:
        errors.append(f"Entry {entry_num} ({glytoucan_ac}): Missing DEFINITION field")
    elif not entry['DEFINITION']:
        errors.append(f"Entry {entry_num} ({glytoucan_ac}): DEFINITION field is empty")
    else:
        stats['entries_with_definition'] += 1
        # Check if CSV definition matches
        if pd.notna(csv_row['definition']):
            if entry['DEFINITION'] != csv_row['definition']:
                errors.append(f"Entry {entry_num} ({glytoucan_ac}): DEFINITION in SDF doesn't match CSV")
    
    # SYNONYM field (WURCS - optional)
    if 'SYNONYM' in entry:
        stats['entries_with_synonym'] += 1
        if pd.notna(csv_row['sequence_wurcs']):
            if entry['SYNONYM'] != csv_row['sequence_wurcs']:
                errors.append(f"Entry {entry_num} ({glytoucan_ac}): SYNONYM doesn't match CSV sequence_wurcs")
    
    # IUPAC_NAME field (optional)
    if 'IUPAC_NAME' in entry:
        stats['entries_with_iupac'] += 1
        if pd.notna(csv_row['sequence_iupac_extended']):
            if entry['IUPAC_NAME'] != csv_row['sequence_iupac_extended']:
                errors.append(f"Entry {entry_num} ({glytoucan_ac}): IUPAC_NAME doesn't match CSV sequence_iupac_extended")
    
    # RELATIONSHIP field
    if 'RELATIONSHIP' not in entry:
        errors.append(f"Entry {entry_num} ({glytoucan_ac}): Missing RELATIONSHIP field")
    elif not entry['RELATIONSHIP']:
        errors.append(f"Entry {entry_num} ({glytoucan_ac}): RELATIONSHIP field is empty")
    else:
        # Check if it's a valid ISA number (ISA followed by digits)
        rel = entry['RELATIONSHIP'].strip()
        if not rel.startswith('ISA'):
            errors.append(f"Entry {entry_num} ({glytoucan_ac}): RELATIONSHIP '{rel}' doesn't start with 'ISA'")
        
        # Check against CSV
        if pd.notna(csv_row['relationship']):
            if entry['RELATIONSHIP'] != csv_row['relationship']:
                errors.append(f"Entry {entry_num} ({glytoucan_ac}): RELATIONSHIP doesn't match CSV")
    
    # DATABASE_ACCESSION field
    if 'DATABASE_ACCESSION' not in entry:
        errors.append(f"Entry {entry_num} ({glytoucan_ac}): Missing DATABASE_ACCESSION field")
    else:
        accession = entry['DATABASE_ACCESSION']
        # Check format: should contain GlyGen, GlyTouCan, and PubChem IDs
        if f"GlyGen:{glytoucan_ac}" not in accession:
            errors.append(f"Entry {entry_num} ({glytoucan_ac}): DATABASE_ACCESSION missing GlyGen ID")
        if f"GlyTouCan:{glytoucan_ac}" not in accession:
            errors.append(f"Entry {entry_num} ({glytoucan_ac}): DATABASE_ACCESSION missing GlyTouCan ID")
        if f"PubChem:{int(csv_row['pubchem_id'])}" not in accession:
            errors.append(f"Entry {entry_num} ({glytoucan_ac}): DATABASE_ACCESSION missing PubChem ID or ID mismatch")
    
    # REFERENCE field (PubMed IDs - optional)
    if 'REFERENCE' in entry:
        stats['entries_with_reference'] += 1
        ref = entry['REFERENCE'].strip()
        # Check that PubMed IDs are semicolon-space separated
        if ref and ';' in ref:
            # Check against CSV (should be converted from pipe to semicolon-space)
            if pd.notna(csv_row['pubmed_ids']):
                csv_pubmed = csv_row['pubmed_ids'].replace('|', '; ')
                if ref != csv_pubmed:
                    errors.append(f"Entry {entry_num} ({glytoucan_ac}): REFERENCE doesn't match CSV pubmed_ids")
    
    # Structure block
    if 'STRUCTURE_BLOCK' in entry:
        structure_block = entry['STRUCTURE_BLOCK']
        # Find the counts line (contains atom/bond counts, ends with V2000)
        counts_line = None
        for sb_line in structure_block:
            if 'V2000' in sb_line:
                counts_line = sb_line
                break
        
        if counts_line:
            # Parse atom count (first number in the counts line)
            parts = counts_line.split()
            if parts and parts[0] == '0':
                stats['entries_with_empty_structure'] += 1
            else:
                stats['entries_with_structure'] += 1
        else:
            stats['entries_with_empty_structure'] += 1
    else:
        warnings.append(f"Entry {entry_num} ({glytoucan_ac}): No structure block found")

print()
print(f"Total entries in SDF: {stats['total_entries']}")
print(f"Entries with DEFINITION: {stats['entries_with_definition']}")
print(f"Entries with SYNONYM (WURCS): {stats['entries_with_synonym']}")
print(f"Entries with IUPAC_NAME: {stats['entries_with_iupac']}")
print(f"Entries with REFERENCE (PubMed): {stats['entries_with_reference']}")
print(f"Entries with structures: {stats['entries_with_structure']}")
print(f"Entries with empty structures: {stats['entries_with_empty_structure']}")
print()

# Check for missing glycans (in CSV but not in SDF)
print("Checking for missing entries...")
sdf_ids = set([entry.get('ID', '').replace('TEMP_', '') for entry in entries if 'ID' in entry])
csv_ids = set(csv_lookup.keys())
missing_in_sdf = csv_ids - sdf_ids

if missing_in_sdf:
    print(f"WARNING: {len(missing_in_sdf)} glycans in CSV but missing in SDF:")
    for gid in list(missing_in_sdf)[:10]:
        warnings.append(f"Missing in SDF: {gid}")
    if len(missing_in_sdf) > 10:
        print(f"  ... and {len(missing_in_sdf) - 10} more")
else:
    print(f"✓ All CSV entries are in the SDF file")

print()
print("="*70)
print("SUMMARY")
print("="*70)
print(f"Total Errors: {len(errors)}")
print(f"Total Warnings: {len(warnings)}")
print()

if errors:
    print("ERRORS:")
    for i, error in enumerate(errors[:20], 1):
        print(f"  {i}. {error}")
    if len(errors) > 20:
        print(f"  ... and {len(errors) - 20} more errors")
    print()

if warnings:
    print("WARNINGS:")
    for i, warning in enumerate(warnings[:20], 1):
        print(f"  {i}. {warning}")
    if len(warnings) > 20:
        print(f"  ... and {len(warnings) - 20} more warnings")
    print()

# Final verdict
print("="*70)
if len(errors) == 0:
    print("✓ VALIDATION PASSED - All entries are valid!")
    sys.exit(0)
else:
    print("✗ VALIDATION FAILED - Please review errors above")
    sys.exit(1)