import sys
import gradio as gr
import os
import hashlib
from pathlib import Path
import pandas as pd
from io import StringIO
from usalign_runner import USalignRunner
from utils import (calculate_md5, get_TM_mat_from_df, get_cluster_z_from_df, get_newick_str_from_Z,
                  build_graph_from_mat_df, fill_community_to_graph, get_graph_fig, run_community_analysis, run_usalign, save_pdb_files)
import matplotlib.pyplot as plt
os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# This is a Temp Title")
    
    with gr.Row():
        file_input = gr.File(
            label="Upload PDB Files",
            file_count="multiple",
            file_types=[".pdb"],
            height=200,
            # lines=5,           # 默认显示行数
            # max_lines=10,      # 最大可见行数（超过后自动滚动）
            # container=True  
        )

    output = gr.Textbox(label="Upload Results",
                lines=5,           # 默认显示行数
                max_lines=5,      # 最大可见行数（超过后自动滚动）
                container=True  )
    
    threshold = gr.Slider(minimum=0, maximum=1, value=0.75, label="Threshold")
    
    with gr.Row():
       
        submit_btn = gr.Button("Upload Files")
        run_usalign_btn = gr.Button("Run USalign")
        community_btn = gr.Button("Run Community")
    
    md5_hash = gr.State("")
    with gr.Tab("USalign Results"):
        results_df = gr.DataFrame(label="USalign Results",height=400,wrap=True,)
    with gr.Tab("TM Matrix"):    
        # Add new output components for community analysis with height limits
        tm_matrix_output = gr.DataFrame(label="TM Matrix",height=400,wrap=True,show_label=True)
    with gr.Tab("Newick Tree"):
        newick_output = gr.Textbox(label="Newick Tree",
            lines=5,           # 默认显示行数
            max_lines=10,      # 最大可见行数（超过后自动滚动）
            container=True  )
    with gr.Tab("Structure Similarity Network"):
        network_plot = gr.Plot(label="Structure Similarity Network")
    
    # Combine download buttons into a single row
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Download Results")
            download_tm = gr.File(label="Download Files")

    submit_btn.click(
        fn=save_pdb_files,
        inputs=[file_input],
        outputs=output
    )   

    def update_md5_hash(files):
        if files:
            return calculate_md5(files)
        return ""
    
    file_input.change(
        fn=update_md5_hash,
        inputs=[file_input],
        outputs=[md5_hash]
    )
    
    run_usalign_btn.click(
        fn=run_usalign,
        inputs=[md5_hash],
        outputs=[results_df]
    )
    
    def process_community_analysis(results_df, md5_hash,threshold):
        if results_df.empty:
            return None, None, None, None, None, None
        
        results = run_community_analysis(results_df, "./data", md5_hash,threshold)
        
        if "Error" in results:
            return None, None, None, None, None, None
            
        # Prepare download files

        
        return (
            results["tm_matrix"],
            results["newick_str"],
            results["network_fig"],
            results["files"]
        )
    
    community_btn.click(
        fn=process_community_analysis,
        inputs=[results_df, md5_hash,threshold],
        outputs=[
            tm_matrix_output,
            newick_output,
            network_plot,
            download_tm,
        ]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0",server_port=7869)