# Setup
import sys
import os
sys.path.append('/Users/sarahurbut/aladynoulli2/pyScripts/new_oct_revision')

import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

# Set style
sns.set_style("whitegrid")
plt.rcParams['figure.dpi'] = 300

print("Setup complete")

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/generate_prs_signature_plots.py --batch_dir='/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/' --output_dir='/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/paper_figs/prs_signatures_corrected_E_PCS_SEX' --n_top 30

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/paper_figs/rap/visualize_genetic_loci.py

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/reviewer_responses/notebooks/R1/check_loci_uniqueness.py

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/main_paper_figures/generate_genetic_validation_multipanel.py

"""
Compare loci from individual SIG files with all_loci_annotated.tsv
to verify they match and identify unique loci.
"""

import pandas as pd
import numpy as np
from pathlib import Path
from collections import defaultdict

# Path to SIG files directory (user mentioned it's in Desktop)
sig_files_dir = Path("/Users/sarahurbut/Desktop/leadforpaper")

# Load all_loci_annotated.tsv
loci_file = Path("/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/all_loci_annotated.tsv")
print(f"Loading all_loci_annotated.tsv from {loci_file}...")
all_loci_df = pd.read_csv(loci_file, sep='\t')
all_loci_df['SIG_NUM'] = all_loci_df['SIG'].str.replace('SIG', '').astype(int)

# Create unique identifier: CHR:POS:EA:OA (UID) or CHR:POS if UID not available
if 'UID' in all_loci_df.columns:
    all_loci_df['unique_id'] = all_loci_df['UID']
else:
    # Try different column name variations
    chr_col = None
    for col in all_loci_df.columns:
        if col.upper().replace('#', '').replace('_', '').strip() == 'CHR' or col == '#CHR':
            chr_col = col
            break
    
    if chr_col:
        all_loci_df['unique_id'] = all_loci_df[chr_col].astype(str) + ':' + all_loci_df['POS'].astype(str)
    else:
        print("Warning: Could not find CHR column in all_loci_annotated.tsv")
        print(f"Available columns: {all_loci_df.columns.tolist()}")

print(f"Loaded {len(all_loci_df)} loci from all_loci_annotated.tsv")
print(f"Signatures: {sorted(all_loci_df['SIG_NUM'].unique())}")

# Check if directory exists
if not sig_files_dir.exists():
    print(f"⚠️  Directory does not exist: {sig_files_dir}")
else:
    print(f"✓ Found directory: {sig_files_dir}")
    print(f"Files in directory:")
    for item in sorted(sig_files_dir.iterdir()):
        if item.is_file():
            size = item.stat().st_size
            print(f"  {item.name} ({size} bytes)")

# Load individual SIG files
sig_files_data = {}
sig_file_paths = {}

for sig_num in range(21):
    # Try different file naming patterns
    patterns = [
        f"SIG{sig_num}_AUC_ukb_eur*.lead.sumstats.txt",
        f"*SIG{sig_num}_AUC_ukb_eur*.lead.sumstats.txt",
        f"*SIG{sig_num}*.txt"
    ]
    
    sig_df = None
    file_path = None
    
    for pattern in patterns:
        matches = list(sig_files_dir.glob(pattern))
        if matches:
            file_path = matches[0]
            break
    
    if file_path and file_path.exists():
        # Check if file is empty
        if file_path.stat().st_size == 0:
            print(f"  ⚠ SIG{sig_num}: File is empty - {file_path.name}")
            continue
            
        try:
            # Read the file - header starts with #CHR, so don't use comment='#'
            # Read normally and pandas will handle the # in the column name
            sig_df = pd.read_csv(file_path, sep='\t')
            
            if sig_df is None or len(sig_df) == 0:
                print(f"  ⚠ SIG{sig_num}: No data rows - {file_path.name}")
                continue
            
            # Print column names for debugging (first file only)
            if sig_num == 0:
                print(f"\n  Column names in SIG files:")
                print(f"    {list(sig_df.columns)}")
            
            # Create unique identifier - use UID if available, otherwise CHR:POS
            if 'UID' in sig_df.columns:
                sig_df['unique_id'] = sig_df['UID']
            else:
                # Try to find CHR column (might be #CHR)
                chr_col = None
                for col in sig_df.columns:
                    if col == '#CHR' or col.upper().replace('#', '').strip() == 'CHR':
                        chr_col = col
                        break
                
                if chr_col and 'POS' in sig_df.columns:
                    sig_df['unique_id'] = sig_df[chr_col].astype(str) + ':' + sig_df['POS'].astype(str)
                else:
                    print(f"  ✗ SIG{sig_num}: Could not find CHR/POS columns")
                    print(f"    Available columns: {list(sig_df.columns)}")
                    continue
            
            sig_df['SIG_NUM'] = sig_num
            sig_files_data[sig_num] = sig_df
            sig_file_paths[sig_num] = file_path
            print(f"  ✓ SIG{sig_num}: {len(sig_df)} loci from {file_path.name}")
            
        except Exception as e:
            print(f"  ✗ SIG{sig_num}: Error processing file - {e}")
            import traceback
            traceback.print_exc()

print(f"\nLoaded {len(sig_files_data)} signature files")

if len(sig_files_data) == 0:
    print("⚠️  No signature files were successfully loaded. Please check:")
    print("  1. File paths and naming conventions")
    print("  2. File formats (should be tab-separated)")
    print("  3. File contents (some may be empty)")
    exit(1)

# Compare loci
comparison_results = []

for sig_num in range(21):
    if sig_num not in sig_files_data:
        continue
    
    sig_df = sig_files_data[sig_num]
    all_loci_sig = all_loci_df[all_loci_df['SIG_NUM'] == sig_num]
    
    # Get unique IDs
    sig_unique_ids = set(sig_df['unique_id'].values)
    all_loci_unique_ids = set(all_loci_sig['unique_id'].values)
    
    # Find matches and differences
    in_both = sig_unique_ids & all_loci_unique_ids
    only_in_sig_file = sig_unique_ids - all_loci_unique_ids
    only_in_all_loci = all_loci_unique_ids - sig_unique_ids
    
    comparison_results.append({
        'Signature': sig_num,
        'In_SIG_file': len(sig_unique_ids),
        'In_all_loci': len(all_loci_unique_ids),
        'In_both': len(in_both),
        'Only_in_SIG_file': len(only_in_sig_file),
        'Only_in_all_loci': len(only_in_all_loci),
        'Match_rate': len(in_both) / max(len(sig_unique_ids), 1) * 100
    })
    
    if len(only_in_sig_file) > 0 or len(only_in_all_loci) > 0:
        print(f"\nSIG{sig_num} discrepancies:")
        if len(only_in_sig_file) > 0:
            print(f"  Only in SIG file ({len(only_in_sig_file)}):")
            for uid in list(only_in_sig_file)[:5]:  # Show first 5
                row = sig_df[sig_df['unique_id'] == uid].iloc[0]
                rsid = row.get('rsid', 'N/A')
                print(f"    {uid} (rsid: {rsid})")
        if len(only_in_all_loci) > 0:
            print(f"  Only in all_loci ({len(only_in_all_loci)}):")
            for uid in list(only_in_all_loci)[:5]:  # Show first 5
                row = all_loci_sig[all_loci_sig['unique_id'] == uid].iloc[0]
                rsid = row.get('rsid', 'N/A')
                print(f"    {uid} (rsid: {rsid})")

# Create summary table
comparison_df = pd.DataFrame(comparison_results)
print("\n" + "="*80)
print("COMPARISON SUMMARY")
print("="*80)
print(comparison_df.to_string(index=False))

# Calculate total unique loci across all signatures
all_sig_unique_ids = set()
for sig_df in sig_files_data.values():
    all_sig_unique_ids.update(sig_df['unique_id'].values)

all_loci_unique_ids = set(all_loci_df['unique_id'].values)

print(f"\n" + "="*80)
print("OVERALL SUMMARY")
print("="*80)
print(f"Total unique loci in SIG files: {len(all_sig_unique_ids)}")
print(f"Total unique loci in all_loci_annotated.tsv: {len(all_loci_unique_ids)}")
print(f"Loci in both: {len(all_sig_unique_ids & all_loci_unique_ids)}")
print(f"Only in SIG files: {len(all_sig_unique_ids - all_loci_unique_ids)}")
print(f"Only in all_loci: {len(all_loci_unique_ids - all_sig_unique_ids)}")

# Save comparison table
output_file = Path("/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/loci_comparison_table.csv")
output_file.parent.mkdir(parents=True, exist_ok=True)
comparison_df.to_csv(output_file, index=False)
print(f"\n✓ Saved comparison table to: {output_file}")

# Create detailed table of unique loci per signature
unique_loci_table = []
for sig_num in range(21):
    if sig_num not in sig_files_data:
        continue
    
    sig_df = sig_files_data[sig_num]
    for _, row in sig_df.iterrows():
        # Get column values with fallbacks
        chr_col = '#CHR' if '#CHR' in sig_df.columns else 'CHR'
        chr_val = row.get(chr_col, 'N/A')
        pos_val = row.get('POS', 'N/A')
        uid_val = row.get('UID', 'N/A')
        rsid_val = row.get('rsid', 'N/A')
        log10p_val = row.get('LOG10P', 'N/A')
        gene_val = row.get('nearestgene', 'N/A') if 'nearestgene' in row else 'N/A'
        
        unique_loci_table.append({
            'Signature': sig_num,
            'CHR': chr_val,
            'POS': pos_val,
            'UID': uid_val,
            'rsid': rsid_val,
            'LOG10P': log10p_val,
            'nearestgene': gene_val
        })

unique_loci_df = pd.DataFrame(unique_loci_table)
unique_loci_output = output_file.parent / "unique_loci_per_signature.csv"
unique_loci_df.to_csv(unique_loci_output, index=False)
print(f"✓ Saved unique loci table to: {unique_loci_output}")
print(f"  Total loci: {len(unique_loci_df)}")

"""
Convert unique_loci_per_signature.csv to VCF format for Ensembl VEP upload
"""

import pandas as pd
from pathlib import Path

# Load your loci table
loci_df = pd.read_csv("/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/unique_loci_per_signature.csv")

# Create VCF file
vcf_lines = []
vcf_lines.append("##fileformat=VCFv4.2")
vcf_lines.append("##source=unique_loci_per_signature")
vcf_lines.append("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO")

# For each locus, create a VCF line
# Note: We'll need to handle cases where we don't have REF/ALT
for _, row in loci_df.iterrows():
    chrom = f"chr{int(row['CHR'])}"
    pos = int(row['POS'])
    rsid = row.get('rsid', '.')
    
    # Try to extract REF/ALT from UID if available
    # UID format is like "1:154400320:A:G"
    ref = '.'
    alt = '.'
    if 'UID' in row and pd.notna(row['UID']):
        uid_parts = str(row['UID']).split(':')
        if len(uid_parts) >= 4:
            ref = uid_parts[2]
            alt = uid_parts[3]
    
    # If still missing, use placeholder
    if ref == '.' or alt == '.':
        ref = 'N'
        alt = 'N'
    
    # Create VCF line
    vcf_line = f"{chrom}\t{pos}\t{rsid}\t{ref}\t{alt}\t.\t.\t."
    vcf_lines.append(vcf_line)

# Write VCF file
output_vcf = Path("/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/unique_loci_for_vep.vcf")
with open(output_vcf, 'w') as f:
    f.write('\n'.join(vcf_lines))

print(f"✓ Created VCF file: {output_vcf}")
print(f"  Total variants: {len(loci_df)}")
print(f"\nYou can now upload this file to:")
print(f"  https://www.ensembl.org/Tools/VEP")

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/annotate_loci.py

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/create_annotation_comparison.py

Figure 4: Genetic Validation (manuscript Fig 4)¶

Purpose¶

Panels Required:¶

Key Message:¶