import pandas as pd

# Load predictions
predictions = pd.read_csv(
    'data/genomes/GRCh38/g27/trifid_predictions.tsv.gz', 
    compression='gzip', 
    sep='\t'
)

# Filter for FGFR1
gene_name = 'FGFR1'
fgfr1_data = predictions.loc[
    predictions['gene_name'] == gene_name
][
    ['transcript_id', 'gene_name', 'trifid_score', 'norm_trifid_score', 
     'appris', 'length', 'sequence']
]

print(fgfr1_data)

from trifid.models.interpret import TreeInterpretation
import pickle

# Load model
model = pickle.load(open('models/selected_model.pkl', 'rb'))

# Load training data
df_training = pd.read_csv(
    'data/model/training_set_final.g27.tsv.gz', 
    sep='\t', 
    compression='gzip'
)

# Create interpretation object
interpretation = TreeInterpretation(
    model=model,
    df=df_training,
    features_col=training_features,
    target_col='label',
    random_state=123,
    test_size=0.25
)

# Explain specific isoform
explanation = interpretation.local_explanation(
    df_predictions, 
    sample='ENST00000356207'
)
print(explanation.head(10))

def analyze_gene(gene_name, predictions_file):
    """
    Analyze TRIFID predictions for a specific gene.
    
    Args:
        gene_name: Gene symbol (e.g., 'FGFR1')
        predictions_file: Path to TRIFID predictions
    
    Returns:
        DataFrame with isoform analysis
    """
    df = pd.read_csv(predictions_file, sep='\t', compression='gzip')
    
    gene_data = df[
        df['gene_name'] == gene_name
    ].sort_values('trifid_score', ascending=False)
    
    print(f"\nAnalysis for {gene_name}:")
    print(f"Total isoforms: {len(gene_data)}")
    print(f"Functional (score >= 0.5): {(gene_data['trifid_score'] >= 0.5).sum()}")
    print(f"\nTop isoform: {gene_data.iloc[0]['transcript_id']}")
    print(f"Score: {gene_data.iloc[0]['trifid_score']:.2f}")
    
    return gene_data

# Run analysis
fgfr1_analysis = analyze_gene('FGFR1', 'data/genomes/GRCh38/g27/trifid_predictions.tsv.gz')