import pandas as pd
import os
import requests

def load_data(hgnc_path):
    # Load HGNC gene list (protein-coding genes)
    if os.path.exists(hgnc_path):
        hgnc_df = pd.read_csv(hgnc_path, sep='\t', dtype=str)
    else:
        url = 'https://storage.googleapis.com/public-download-files/hgnc/tsv/tsv/locus_types/gene_with_protein_product.txt'
        resp = requests.get(url)
        resp.raise_for_status()
        with open(hgnc_path, 'wb') as f:
            f.write(resp.content)
        hgnc_df = pd.read_csv(hgnc_path, sep='\t', dtype=str)

    return hgnc_df


def build_gene_symbol_map(hgnc_df):
    """
    Map all approved, alias, and previous symbols to the current HGNC symbol.
    """
    symbol_map = {}
    for _, row in hgnc_df.iterrows():
        approved = row.get('symbol')
        if pd.isna(approved):
            continue
        symbol_map[approved] = approved
        for col in ['alias_symbol', 'prev_symbol']:
            if col in row and pd.notna(row[col]):
                for alias in row[col].split('|'):
                    symbol_map[alias] = approved
    ensembl_map = dict(zip(hgnc_df['symbol'], hgnc_df['ensembl_gene_id']))
    entrez_map = dict(zip(hgnc_df['symbol'], hgnc_df['entrez_id']))
    return symbol_map, ensembl_map, entrez_map

def get_gene_info(gene, hgnc_path):
    hgnc_df = load_data(hgnc_path)
    symbol_map, ensembl_map, entrez_map = build_gene_symbol_map(hgnc_df)
    approved_symbol = symbol_map.get(gene)
    if not approved_symbol:
        raise ValueError(f"Gene symbol '{gene}' not found in HGNC data.")
    ensembl_id = ensembl_map.get(approved_symbol)
    entrez_id = entrez_map.get(approved_symbol)
    return approved_symbol, ensembl_id, entrez_id

if __name__ == "__main__":
    hgnc_path = "/home/dbanerjee/deepro/rwe/data/hgnc/gene_with_protein_product.txt"
    approved_symbol, ensembl_id, entrez_id = get_gene_info("ANGPTL3", hgnc_path)
    print(f"Approved symbol: {approved_symbol}, Ensembl ID: {ensembl_id}, Entrez ID: {entrez_id}")

