#!/usr/bin/env python3
"""
Map probes to gene symbols using GEO SOFT format annotations
"""

import os
import re
import gzip
import urllib.request
import pandas as pd
import numpy as np

# Paths
BASE_DIR = "/app/sandbox/session_20251217_123457_77eda5efa279"
WORKFLOW_DIR = f"{BASE_DIR}/workflow"
DATA_DIR = f"{BASE_DIR}/data"

print("=" * 80)
print("MAPPING PROBES TO GENES VIA GEO SOFT FILE")
print("=" * 80)

# Download GPL SOFT file
platform_id = "21185"
soft_url = f"https://ftp.ncbi.nlm.nih.gov/geo/platforms/GPL21nnn/GPL{platform_id}/soft/GPL{platform_id}_family.soft.gz"

print(f"\n[1/3] Downloading GPL{platform_id} SOFT file...")
print(f"  → URL: {soft_url}")

try:
    response = urllib.request.urlopen(soft_url, timeout=120)
    compressed_data = response.read()
    print(f"  → Downloaded {len(compressed_data)} bytes")

    # Decompress
    decompressed_data = gzip.decompress(compressed_data)
    soft_file = f"{DATA_DIR}/GPL{platform_id}.soft"

    with open(soft_file, 'wb') as f:
        f.write(decompressed_data)

    print(f"  → Saved to: {soft_file}")

    # Parse SOFT file
    print(f"\n[2/3] Parsing SOFT file for probe-to-gene mappings...")

    with open(soft_file, 'r', encoding='latin-1') as f:
        lines = f.readlines()

    # Find table section
    table_start = None
    for i, line in enumerate(lines):
        if line.startswith('!platform_table_begin'):
            table_start = i + 1
            break

    if table_start is None:
        raise ValueError("Could not find table start in SOFT file")

    # Find table end
    table_end = None
    for i in range(table_start, len(lines)):
        if lines[i].startswith('!platform_table_end'):
            table_end = i
            break

    # Extract table
    table_lines = lines[table_start:table_end]

    # Parse header
    header = table_lines[0].strip().split('\t')
    print(f"  → Columns: {header[:10]}")

    # Find gene symbol column
    gene_col_idx = None
    for idx, col in enumerate(header):
        if col.lower() in ['gene_symbol', 'gene symbol', 'symbol', 'gene name']:
            gene_col_idx = idx
            print(f"  → Found gene symbol column: '{col}' at index {idx}")
            break

    if gene_col_idx is None:
        # Try to find by pattern
        for idx, col in enumerate(header):
            if 'gene' in col.lower() and 'symbol' in col.lower():
                gene_col_idx = idx
                print(f"  → Found gene symbol column: '{col}' at index {idx}")
                break

    if gene_col_idx is None:
        print(f"  ! Warning: Could not find gene symbol column")
        print(f"  ! Available columns: {header}")
        raise ValueError("No gene symbol column found")

    # Parse mappings
    probe_to_gene = {}
    for i, line in enumerate(table_lines[1:]):
        if i % 5000 == 0:
            print(f"    Parsing row {i}/{len(table_lines)-1}...")

        parts = line.strip().split('\t')
        if len(parts) > gene_col_idx:
            probe_id = parts[0]  # First column is typically ID
            gene_symbol = parts[gene_col_idx]

            # Clean gene symbol
            if gene_symbol and gene_symbol not in ['', '---', 'NA', 'na', 'null']:
                # Take first gene if multiple (separated by /// or ,)
                gene_symbol = re.split(r'[,/]+', gene_symbol)[0].strip()
                probe_to_gene[probe_id] = gene_symbol

    print(f"  → Mapped {len(probe_to_gene)} probes to gene symbols")

    # Apply to expression matrix
    print(f"\n[3/3] Applying mappings to expression matrix...")

    expr_df = pd.read_csv(f"{WORKFLOW_DIR}/expression_matrix.csv", index_col=0)
    print(f"  → Loaded expression matrix: {expr_df.shape}")

    # Clean probe IDs (remove quotes)
    expr_df.index = [idx.strip('"') for idx in expr_df.index]

    # Map to genes
    gene_symbols = []
    mapped_count = 0

    for probe in expr_df.index:
        if probe in probe_to_gene:
            gene_symbols.append(probe_to_gene[probe])
            mapped_count += 1
        else:
            # Keep probe ID if no mapping
            gene_symbols.append(probe)

    print(f"  → Successfully mapped {mapped_count}/{len(expr_df)} probes ({mapped_count/len(expr_df):.1%})")

    # Create gene-level expression
    expr_df['Gene_Symbol'] = gene_symbols

    # Handle duplicates by averaging
    print("  → Aggregating duplicate genes (averaging)...")
    expr_df_genes = expr_df.groupby('Gene_Symbol').mean()

    print(f"  → Final matrix: {expr_df_genes.shape[0]} genes x {expr_df_genes.shape[1]} samples")

    # Save
    output_file = f"{WORKFLOW_DIR}/expression_matrix.csv"
    expr_df_genes.to_csv(output_file)
    print(f"  → Saved to: {output_file}")

    # Show some examples
    if mapped_count > 0:
        gene_list = [g for g in expr_df_genes.index if g in probe_to_gene.values()][:15]
        print(f"  → Example genes: {', '.join(gene_list)}")

    print("\n" + "=" * 80)
    print("✓ SUCCESS: Probe-to-gene mapping completed")
    print("=" * 80)

except Exception as e:
    print(f"\n✗ ERROR: {e}")
    print("  → Keeping existing expression matrix with probe IDs")
    print("  → Analysis can still proceed with probe identifiers")
    print("=" * 80)