#!/usr/bin/env python3
"""
Offline network enrichment and interaction analysis using local STRING-DB files.
Replicates the NETWORK_ENRICHMENT Nextflow process without internet connection.
"""

import pandas as pd
import argparse
from pathlib import Path
from scipy import stats
import numpy as np


def load_string_data(string_db_dir):
    """Load all required STRING-DB files."""
    string_db_dir = Path(string_db_dir)
    
    print("Loading STRING-DB files...")
    
    # Load alias mapping (ENST -> ENSP)
    aliases = pd.read_csv(
        string_db_dir / "9606.protein.aliases.enst.v12.0.tsv",
        sep='\t',
        comment='#',
        names=['string_protein_id', 'alias', 'source']
    )
    enst_to_ensp = aliases[aliases['source'] == 'Ensembl_transcript'].set_index('alias')['string_protein_id'].to_dict()
    
    # Load protein info (for preferred names)
    info = pd.read_csv(
        string_db_dir / "9606.protein.info.v12.0.txt",
        sep='\t',
        comment='#',
        names=['string_protein_id', 'preferred_name', 'protein_size', 'annotation']
    )
    ensp_to_name = info.set_index('string_protein_id')['preferred_name'].to_dict()
    
    # Load protein links
    # Load protein links - USE DETAILED FILE
    links = pd.read_csv(
        string_db_dir / "9606.protein.links.detailed.v12.0.txt",
        sep=' '
    )
    
    # Load enrichment terms
    enrichment = pd.read_csv(
        string_db_dir / "9606.protein.enrichment.terms.v12.0.txt",
        sep='\t',
        comment='#',
        names=['string_protein_id', 'category', 'term', 'description']
    )
    
    print(f"Loaded {len(enst_to_ensp)} ENST->ENSP mappings")
    print(f"Loaded {len(links)} protein-protein interactions")
    print(f"Loaded {len(enrichment)} enrichment annotations")
    
    return enst_to_ensp, ensp_to_name, links, enrichment


def perform_enrichment_analysis(query_proteins, enrichment_df, ensp_to_name, background_size=20000):
    """
    Perform enrichment analysis similar to STRING-DB API.
    Uses Fisher's exact test for each term.
    """
    results = []
    
    query_set = set(query_proteins)
    n_query = len(query_set)
    
    # Group by category and term
    for (category, term), group in enrichment_df.groupby(['category', 'term']):
        term_proteins = set(group['string_protein_id'])
        
        # Intersection: proteins in both query and term
        intersection = query_set & term_proteins
        n_intersection = len(intersection)
        
        if n_intersection == 0:
            continue
        
        # Background: total proteins with this term
        n_term_background = len(term_proteins)
        
        # Fisher's exact test
        # Contingency table:
        #                  In term    Not in term
        # In query           a            b
        # Not in query       c            d
        
        a = n_intersection
        b = n_query - a
        c = n_term_background - a
        d = background_size - n_term_background - b
        
        # One-sided test for enrichment
        oddsratio, pvalue = stats.fisher_exact([[a, b], [c, d]], alternative='greater')
        
        # Get description (take first)
        description = group['description'].iloc[0]
        
        # Get gene names
        intersection_list = list(intersection)
        input_genes = ','.join(intersection_list)
        preferred_names = ','.join([ensp_to_name.get(p, p.split('.')[-1]) for p in intersection_list])
        
        results.append({
            'category': category,
            'term': term,
            'number_of_genes': n_intersection,
            'number_of_genes_in_background': n_term_background,
            'ncbiTaxonId': 9606,
            'inputGenes': input_genes,
            'preferredNames': preferred_names,
            'p_value': pvalue,
            'description': description
        })
    
    if not results:
        return pd.DataFrame()
    
    # Create DataFrame and calculate FDR
    results_df = pd.DataFrame(results)
    results_df = results_df.sort_values('p_value')
    
    # Benjamini-Hochberg FDR correction
    n_tests = len(results_df)
    results_df['fdr'] = results_df['p_value'] * n_tests / (np.arange(1, n_tests + 1))
    results_df['fdr'] = results_df['fdr'].clip(upper=1.0)
    
    # Sort by p-value
    results_df = results_df.sort_values('p_value').reset_index(drop=True)
    
    return results_df


def get_network_interactions(query_proteins, links_df, ensp_to_name):
    """
    Extract network interactions for query proteins.
    Converts combined_score to normalized scores similar to STRING API.
    """
    query_set = set(query_proteins)
    
    # Filter links where both proteins are in query set
    interactions = links_df[
        links_df['protein1'].isin(query_set) & 
        links_df['protein2'].isin(query_set)
    ].copy()
    
    if len(interactions) == 0:
        return pd.DataFrame()
    
    interactions['score'] = interactions['combined_score'] / 1000.0
    interactions['nscore'] = interactions['neighborhood'] / 1000.0
    interactions['fscore'] = interactions['fusion'] / 1000.0
    interactions['pscore'] = interactions['cooccurence'] / 1000.0  # Note: typo in STRING file
    interactions['ascore'] = interactions['coexpression'] / 1000.0
    interactions['escore'] = interactions['experimental'] / 1000.0
    interactions['dscore'] = interactions['database'] / 1000.0
    interactions['tscore'] = interactions['textmining'] / 1000.0
    
    # Add preferred names
    interactions['preferredName_A'] = interactions['protein1'].map(
        lambda x: ensp_to_name.get(x, x.split('.')[-1])
    )
    interactions['preferredName_B'] = interactions['protein2'].map(
        lambda x: ensp_to_name.get(x, x.split('.')[-1])
    )
    
    # Rename columns to match STRING API output
    interactions = interactions.rename(columns={
        'protein1': 'stringId_A',
        'protein2': 'stringId_B'
    })
    
    interactions['ncbiTaxonId'] = 9606
    
    # Select and order columns to match original output
    column_order = [
        'stringId_A', 'stringId_B', 'preferredName_A', 'preferredName_B',
        'ncbiTaxonId', 'score', 'nscore', 'fscore', 'pscore', 
        'ascore', 'escore', 'dscore', 'tscore'
    ]
    
    interactions = interactions[column_order].reset_index(drop=True)
    
    return interactions


def main():
    parser = argparse.ArgumentParser(
        description='Offline network enrichment analysis using local STRING-DB files'
    )
    parser.add_argument(
        'interactions_file',
        type=str,
        help='Input TSV file with significant interactions (conplex output)'
    )
    parser.add_argument(
        '--string-db-dir',
        type=str,
        default='/data/bugra/digital_trials/app_network',
        help='Directory containing STRING-DB files'
    )
    parser.add_argument(
        '--threshold',
        type=float,
        default=0.0,
        help='ConPlex score threshold for protein network'
    )
    parser.add_argument(
        '--max-proteins',
        type=int,
        default=450,
        help='Maximum number of proteins to analyze'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='.',
        help='Output directory for results'
    )
    
    args = parser.parse_args()
    
    # Load STRING-DB data
    enst_to_ensp, ensp_to_name, links, enrichment = load_string_data(args.string_db_dir)
    
    # Load interaction data
    print(f"\nLoading interaction data from {args.interactions_file}...")
    interaction_data = pd.read_csv(args.interactions_file, sep='\t')
    interaction_data = interaction_data.sort_values('conplex_score', ascending=False)
    
    # Remove duplicates and limit
    interaction_data = interaction_data.drop_duplicates('transcipt')
    if interaction_data.shape[0] > args.max_proteins:
        print(f"Limiting to top {args.max_proteins} proteins")
        interaction_data = interaction_data.iloc[:args.max_proteins]
    
    # Filter by threshold
    filtered_data = interaction_data[interaction_data['conplex_score'] > args.threshold]
    print(f"Found {len(filtered_data)} proteins above threshold {args.threshold}")
    
    # Convert ENST to ENSP
    enst_list = filtered_data['transcipt'].tolist()
    ensp_list = [enst_to_ensp.get(enst) for enst in enst_list]
    ensp_list = [e for e in ensp_list if e is not None]
    
    print(f"Mapped {len(ensp_list)} ENST IDs to ENSP IDs")
    
    if len(ensp_list) == 0:
        print("ERROR: No valid ENSP mappings found!")
        return
    
    # Perform enrichment analysis
    print("\nPerforming enrichment analysis...")
    enrichment_results = perform_enrichment_analysis(ensp_list, enrichment, ensp_to_name)
    print(f"Found {len(enrichment_results)} enriched terms")
    
    # Get network interactions
    print("\nExtracting network interactions...")
    network_interactions = get_network_interactions(ensp_list, links, ensp_to_name)
    print(f"Found {len(network_interactions)} protein-protein interactions")
    
    # Generate output filename
    input_path = Path(args.interactions_file)
    output_name = input_path.stem.replace('_significant_interactions', '')
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Save results
    enrichment_file = output_dir / f"{output_name}_network_enrichment.tsv"
    interactions_file = output_dir / f"{output_name}_network_interactions.tsv"
    
    enrichment_results.to_csv(enrichment_file, sep='\t')
    network_interactions.to_csv(interactions_file, sep='\t')
    
    print(f"\nResults saved:")
    print(f"  Enrichment: {enrichment_file}")
    print(f"  Interactions: {interactions_file}")


if __name__ == '__main__':
    main()