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

def get_tissue_distribution(file_name):
    # Convert to Path object
    file_path = Path(file_name)
    
    # Load interaction data — only the 'transcipt' column is referenced downstream
    # (the bulk of the file is the 'smile' column we never use). Skipping it cuts
    # memory ~30x on metabolite-heavy drugs (11 GB file -> ~400 MB resident).
    significant_interactions = pd.read_csv(file_path, sep='\t', usecols=['transcipt'])
    
    # Check if data exists
    if len(significant_interactions) == 0:
        print(f"No significant interactions found in {file_name}. Creating empty output.")
        tissue_per_prot_symbol_pd = pd.DataFrame()
        save_name = str(file_path.parent / (file_path.stem.replace('_significant_interactions', '_tissue_distribution') + '.tsv'))
        tissue_per_prot_symbol_pd.to_csv(save_name, sep='\t')
        print(f"Saved empty tissue distribution to: {save_name}")
        return
    
    # Load protein expression per tissue
    HPA = pd.read_csv('/home/omic/HPA_normal_ihc_data.tsv', sep='\t')
    # Load enst symbol mapping data
    MANE_all_transcipts = pd.read_csv('/home/omic/MANE_all_transcipts.csv')

    # Filter out tissues with low confidence
    HPA = HPA[(HPA['Level'] == 'Medium') | (HPA['Level'] == 'High')]


    # Get protein symbols from significant interactions
    prot_enst = [i.split('_')[0] for i in significant_interactions['transcipt'].unique()]
    
    # Safely map to symbols, skipping missing ones
    prot_symbol = []
    for i in prot_enst:
        matches = MANE_all_transcipts[MANE_all_transcipts['transcipt'] == i]
        if not matches.empty:
            prot_symbol.append(matches['symbol'].iloc[0])
        else:
            print(f"Warning: Transcript {i} not found in MANE mapping. Skipping.")
            
    prot_symbol = pd.unique(prot_symbol)

    # Get unique tissues
    unique_tissues = HPA['Tissue'].unique()
    
    # Get protein expression per tissue
    tissue_per_prot_symbol = [list(HPA[HPA['Gene name'] == i]['Tissue'].unique()) for i in prot_symbol]
    # Transform list to array
    tissue_per_prot_symbol_array = np.array([[tissue in gene_tissues for tissue in unique_tissues] for gene_tissues in tissue_per_prot_symbol])
    
    tissue_per_prot_symbol_pd = pd.DataFrame(tissue_per_prot_symbol_array, columns=unique_tissues, index=prot_symbol)
    
    # Save
    save_name = str(file_path.parent / (file_path.stem.replace('_significant_interactions', '_tissue_distribution') + '.tsv'))
    tissue_per_prot_symbol_pd.to_csv(save_name, sep='\t')
    print(f"Saved tissue distribution to: {save_name}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Get tissue distribution from *significant_interactions.tsv.")
    parser.add_argument("--file_name", required=True, help="file_name")
    args = parser.parse_args()
    
    get_tissue_distribution(args.file_name)