app.py

import gradio as gr
import mtdna_backend
import json
from iterate3 import data_preprocess, model, pipeline
import os
import hashlib
import threading
# Gradio UI
#stop_flag = gr.State(value=False)
class StopFlag:
    def __init__(self):
        self.value = False
global_stop_flag = StopFlag()  # Shared between run + stop

with gr.Blocks() as interface:
    gr.Markdown("# 🧬 mtDNA Location Classifier (MVP)")

    #inputMode = gr.Radio(choices=["Single Accession", "Batch Input"], value="Single Accession", label="Choose Input Mode")
    user_email = gr.Textbox(label="📧 Your email (used to track free quota)")
    usage_display = gr.Markdown("", visible=False)

    # with gr.Group() as single_input_group:
    #     single_accession = gr.Textbox(label="Enter Single Accession (e.g., KU131308)")

    # with gr.Group(visible=False) as batch_input_group:
    #     raw_text = gr.Textbox(label="🧬 Paste Accession Numbers (e.g., MF362736.1,MF362738.1,KU131308,MW291678)")
    #     resume_file = gr.File(label="🗃️ Previously saved Excel output (optional)", file_types=[".xlsx"], interactive=True)
    #     gr.HTML("""<a href="https://drive.google.com/file/d/1t-TFeIsGVu5Jh3CUZS-VE9jQWzNFCs_c/view?usp=sharing" download target="_blank">Download Example CSV Format</a>""")
    #     gr.HTML("""<a href="https://docs.google.com/spreadsheets/d/1lKqPp17EfHsshJGZRWEpcNOZlGo3F5qU/edit?usp=sharing&ouid=112390323314156876153&rtpof=true&sd=true" download target="_blank">Download Example Excel Format</a>""")
    #     file_upload = gr.File(label="📁 Or Upload CSV/Excel File", file_types=[".csv", ".xlsx"], interactive=True, elem_id="file-upload-box")
    raw_text = gr.Textbox(label="🧚 Input Accession Number(s) (single (KU131308) or comma-separated (e.g., MF362736.1,MF362738.1,KU131308,MW291678))")
    #resume_file = gr.File(label="🗃️ Previously saved Excel output (optional)", file_types=[".xlsx"], interactive=True)
    gr.HTML("""<a href="https://docs.google.com/spreadsheets/d/1lKqPp17EfHsshJGZRWEpcNOZlGo3F5qU/edit?usp=sharing" download target="_blank">Download Example Excel Format</a>""")
    file_upload = gr.File(label="📁 Or Upload CSV/Excel File", file_types=[".csv", ".xlsx"], interactive=True)

    with gr.Row():
        run_button = gr.Button("🔍 Submit and Classify")
        stop_button = gr.Button("❌ Stop Batch", visible=True)
        reset_button = gr.Button("🔄 Reset")

    status = gr.Markdown(visible=False)
    
    with gr.Group(visible=False) as results_group:
      # with gr.Accordion("Open to See the Result", open=False) as results:  
      #     with gr.Row():
      #         output_summary = gr.Markdown(elem_id="output-summary")
      #         output_flag = gr.Markdown(elem_id="output-flag")
          
      #     gr.Markdown("---")

      with gr.Accordion("Open to See the Output Table", open=False) as table_accordion:    
          output_table = gr.HTML(render=True)

      with gr.Row():
          output_type = gr.Dropdown(choices=["Excel", "JSON", "TXT"], label="Select Output Format", value="Excel")
          download_button = gr.Button("⬇️ Download Output")
      #download_file = gr.File(label="Download File Here",visible=False)
      download_file = gr.File(label="Download File Here", visible=False, interactive=True)
      progress_box = gr.Textbox(label="Live Processing Log", lines=20, interactive=False)

      gr.Markdown("---")

      gr.Markdown("### 💬 Feedback (required)")
      q1 = gr.Textbox(label="1️⃣ Was the inferred location accurate or helpful? Please explain.")
      q2 = gr.Textbox(label="2️⃣ What would improve your experience with this tool?")
      contact = gr.Textbox(label="📧 Your email or institution (optional)")
      submit_feedback = gr.Button("✅ Submit Feedback")
      feedback_status = gr.Markdown()

    # Functions
    # def toggle_input_mode(mode):
    #     if mode == "Single Accession":
    #         return gr.update(visible=True), gr.update(visible=False)
    #     else:
    #         return gr.update(visible=False), gr.update(visible=True)

    def classify_with_loading():
        return gr.update(value="⏳ Please wait... processing...",visible=True)  # Show processing message
         
    # def classify_dynamic(single_accession, file, text, resume, email, mode):
    #     if mode == "Single Accession":
    #         return classify_main(single_accession)  + (gr.update(visible=False),)
    #     else:
    #         #return summarize_batch(file, text) + (gr.update(visible=False),)  # Hide processing message
    #         return classify_mulAcc(file, text, resume) + (gr.update(visible=False),)  # Hide processing message
    # Logging helpers defined early to avoid NameError
    

    # def classify_dynamic(single_accession, file, text, resume, email, mode):
    #   if mode == "Single Accession":
    #       return classify_main(single_accession) + (gr.update(value="", visible=False),)
    #   else:
    #       return classify_mulAcc(file, text, resume, email, log_callback=real_time_logger, log_collector=log_collector)

    # for single accession
    # def classify_main(accession):
    #     #table, summary, labelAncient_Modern, explain_label = mtdna_backend.summarize_results(accession)
    #     table = mtdna_backend.summarize_results(accession)
    #     #flag_output = f"### 🏺 Ancient/Modern Flag\n**{labelAncient_Modern}**\n\n_Explanation:_ {explain_label}"
    #     return (
    #         #table,
    #         make_html_table(table),
    #         # summary,
    #         # flag_output,
    #         gr.update(visible=True),
    #         gr.update(visible=False),
    #         gr.update(visible=False)
    #     )
    
    #stop_flag = gr.State(value=False)
    #stop_flag = StopFlag()

    # def stop_batch(stop_flag):
    #   stop_flag.value = True
    #   return gr.update(value="❌ Stopping...", visible=True), stop_flag
    def stop_batch():
      global_stop_flag.value = True
      return gr.update(value="❌ Stopping...", visible=True)

    # def threaded_batch_runner(file, text, email):
    #   global_stop_flag.value = False
    #   log_lines = []

    #   def update_log(line):
    #       log_lines.append(line)
    #       yield (
    #           gr.update(visible=False),  # output_table (not yet)
    #           gr.update(visible=False),  # results_group
    #           gr.update(visible=False),  # download_file
    #           gr.update(visible=False),  # usage_display
    #           gr.update(value="⏳ Still processing...", visible=True),  # status
    #           gr.update(value="\n".join(log_lines))  # progress_box
    #       )

    #   # Start a dummy update to say "Starting..."
    #   yield from update_log("🚀 Starting batch processing...")

    #   rows, file_path, count, final_log, warning = mtdna_backend.summarize_batch(
    #       file=file,
    #       raw_text=text,
    #       resume_file=None,
    #       user_email=email,
    #       stop_flag=global_stop_flag,
    #       yield_callback=lambda line: (yield from update_log(line))
    #   )

    #   html = make_html_table(rows)
    #   file_update = gr.update(value=file_path, visible=True) if os.path.exists(file_path) else gr.update(visible=False)
    #   usage_or_warning_text = f"**{count}** samples used by this email." if email.strip() else warning

    #   yield (
    #       html,
    #       gr.update(visible=True),        # results_group
    #       file_update,                    # download_file
    #       gr.update(value=usage_or_warning_text, visible=True),
    #       gr.update(value="✅ Done", visible=True),
    #       gr.update(value=final_log)
    #   )
    
    def threaded_batch_runner(file=None, text="", email=""):
      print("📧 EMAIL RECEIVED:", email)
      import tempfile
      from mtdna_backend import (
          extract_accessions_from_input,
          summarize_results,
          save_to_excel,
          hash_user_id,
          increment_usage,
      )
      import os

      global_stop_flag.value = False  # reset stop flag

      tmp_dir = tempfile.mkdtemp()
      output_file_path = os.path.join(tmp_dir, "batch_output_live.xlsx")
      limited_acc = 50 + (10 if email.strip() else 0)

      # Step 1: Parse input
      accessions, error = extract_accessions_from_input(file, text)
      if error:
          yield (
              "",                          # output_table
              gr.update(visible=False),   # results_group
              gr.update(visible=False),   # download_file
              "",                          # usage_display
              "❌ Error",                  # status
              str(error)                        # progress_box
          )
          return

      total = len(accessions)
      if total > limited_acc:
          accessions = accessions[:limited_acc]
          warning = f"⚠️ Only processing first {limited_acc} accessions."
      else:
          warning = f"✅ All {total} accessions will be processed."

      all_rows = []
      log_lines = []

      # Step 2: Loop through accessions
      for i, acc in enumerate(accessions):
          if global_stop_flag.value:
              log_lines.append(f"🛑 Stopped at {acc} ({i+1}/{total})")
              usage_text = ""
              if email.strip():
                  # user_hash = hash_user_id(email)
                  # usage_count = increment_usage(user_hash, len(all_rows))
                  usage_count = increment_usage(email, len(all_rows))
                  usage_text = f"**{usage_count}** samples used by this email. Ten more samples are added first (you now have 60 limited accessions), then wait we will contact you via this email."
              else:
                  usage_text = f"The limited accession is 50. The user has used {len(all_rows)}, and only {50-len(all_rows)} left."    
              yield (
                  make_html_table(all_rows),
                  gr.update(visible=True),
                  gr.update(value=output_file_path, visible=True),
                  gr.update(value=usage_text, visible=True),
                  "🛑 Stopped",
                  "\n".join(log_lines)
              )
              return

          log_lines.append(f"[{i+1}/{total}] Processing {acc}")
          yield (
              make_html_table(all_rows),
              gr.update(visible=True),
              gr.update(visible=False),
              "",
              "⏳ Processing...",
              "\n".join(log_lines)
          )

          try:
              rows = summarize_results(acc)
              all_rows.extend(rows)
              save_to_excel(all_rows, "", "", output_file_path, is_resume=False)
              log_lines.append(f"✅ Processed {acc} ({i+1}/{total})")
          except Exception as e:
              log_lines.append(f"❌ Failed to process {acc}: {e}")

          yield (
              make_html_table(all_rows),
              gr.update(visible=True),
              gr.update(visible=False),
              "",
              "⏳ Processing...",
              "\n".join(log_lines)
          )

      # Final update
      usage_text = ""
      if email.strip():
          # user_hash = hash_user_id(email)
          # usage_count = increment_usage(user_hash, len(all_rows))
          usage_count = increment_usage(email, len(all_rows))
          usage_text = f"**{usage_count}** samples used by this email. Ten more samples are added first (you now have 60 limited accessions), then wait we will contact you via this email."
      else:
          usage_text = f"The limited accession is 50. The user has used {len(all_rows)}, and only {50-len(all_rows)} left."
      yield (
          make_html_table(all_rows),
          gr.update(visible=True),
          gr.update(value=output_file_path, visible=True),
          gr.update(value=usage_text, visible=True),
          "✅ Done",
          "\n".join(log_lines)
      )

    # def threaded_batch_runner(file=None, text="", email=""):
    #   global_stop_flag.value = False

    #   # Dummy test output that matches expected schema
    #   return (
    #       "<div>✅ Dummy output table</div>",   # HTML string
    #       gr.update(visible=True),             # Group visibility
    #       gr.update(visible=False),            # Download file
    #       "**0** samples used.",               # Markdown
    #       "✅ Done",                           # Status string
    #       "Processing finished."               # Progress string
    #   )

    
    # def classify_mulAcc(file, text, resume, email, log_callback=None, log_collector=None):
    #     stop_flag.value = False
    #     return threaded_batch_runner(file, text, resume, email, status, stop_flag, log_callback=log_callback, log_collector=log_collector)
    

    def make_html_table(rows):
      html = """

      <div style='overflow-x: auto; padding: 10px;'>

          <div style='max-height: 400px; overflow-y: auto; border: 1px solid #444; border-radius: 8px;'>

              <table style='width:100%; border-collapse: collapse; table-layout: auto; font-size: 14px; color: #f1f1f1; background-color: #1e1e1e;'>

                  <thead style='position: sticky; top: 0; background-color: #2c2c2c; z-index: 1;'>

                      <tr>

      """
      headers = ["Sample ID", "Predicted Country", "Country Explanation", "Predicted Sample Type", "Sample Type Explanation", "Sources", "Time cost"]
      html += "".join(
          f"<th style='padding: 10px; border: 1px solid #555; text-align: left; white-space: nowrap;'>{h}</th>"
          for h in headers
      )
      html += "</tr></thead><tbody>"

      for row in rows:
          html += "<tr>"
          for i, col in enumerate(row):
              header = headers[i]
              style = "padding: 10px; border: 1px solid #555; vertical-align: top;"

              # For specific columns like Haplogroup, force nowrap
              if header in ["Country Explanation", "Sample Type Explanation"]:
                style += " max-width: 400px; word-wrap: break-word; white-space: normal;"
              elif header in ["Sample ID", "Predicted Country", "Predicted Sample Type", "Time cost"]:
                  style += " white-space: nowrap; text-overflow: ellipsis; max-width: 200px; overflow: hidden;"

              # if header == "Sources" and isinstance(col, str) and col.strip().lower().startswith("http"):
              #     col = f"<a href='{col}' target='_blank' style='color: #4ea1f3; text-decoration: underline;'>{col}</a>"

              #html += f"<td style='{style}'>{col}</td>"
              if header == "Sources" and isinstance(col, str):
                  links = [f"<a href='{url.strip()}' target='_blank' style='color: #4ea1f3; text-decoration: underline;'>{url.strip()}</a>" for url in col.strip().split("\n") if url.strip()]
                  col = "- "+"<br>- ".join(links)
              elif isinstance(col, str):
                  # lines = []
                  # for line in col.split("\n"):
                  #     line = line.strip()
                  #     if not line:
                  #         continue
                  #     if line.lower().startswith("rag_llm-"):
                  #         content = line[len("rag_llm-"):].strip()
                  #         line = f"{content} (Method: RAG_LLM)"
                  #     lines.append(f"- {line}")
                  col = col.replace("\n", "<br>")
                  #col = col.replace("\t", "&nbsp;&nbsp;&nbsp;&nbsp;")
                  #col = "<br>".join(lines)

              html += f"<td style='{style}'>{col}</td>"
          html += "</tr>"

      html += "</tbody></table></div></div>"
      return html
  

    # def reset_fields():
    #     global_stop_flag.value = False  # 💡 Add this to reset the flag
    #     return (
    #         #gr.update(value=""),  # single_accession
    #         gr.update(value=""),  # raw_text
    #         gr.update(value=None), # file_upload
    #         #gr.update(value=None),  # resume_file
    #         #gr.update(value="Single Accession"), # inputMode
    #         gr.update(value=[], visible=True), # output_table
    #         # gr.update(value="", visible=True), # output_summary
    #         # gr.update(value="", visible=True), # output_flag
    #         gr.update(visible=False), # status
    #         gr.update(visible=False),  # results_group
    #         gr.update(value="", visible=False),  # usage_display
    #         gr.update(value="", visible=False),  # progress_box
    #     )
    def reset_fields():
      global_stop_flag.value = False  # Reset the stop flag

      return (
          gr.update(value=""),          # raw_text
          gr.update(value=None),        # file_upload
          gr.update(value=[], visible=True),   # output_table
          gr.update(value="", visible=True),   # status — reset and make visible again
          gr.update(visible=False),     # results_group
          gr.update(value="", visible=True),   # usage_display — reset and make visible again
          gr.update(value="", visible=True),   # progress_box — reset AND visible!
      )
    #inputMode.change(fn=toggle_input_mode, inputs=inputMode, outputs=[single_input_group, batch_input_group])
    #run_button.click(fn=classify_with_loading, inputs=[], outputs=[status])
    # run_button.click(
    #     fn=classify_dynamic,
    #     inputs=[single_accession, file_upload, raw_text, resume_file,user_email,inputMode],
    #     outputs=[output_table, 
    #     #output_summary, output_flag, 
    #     results_group, download_file, usage_display,status, progress_box]
    # )
    
    # run_button.click(
    #     fn=threaded_batch_runner,
    #     #inputs=[file_upload, raw_text, resume_file, user_email],
    #     inputs=[file_upload, raw_text, user_email],
    #     outputs=[output_table, results_group, download_file, usage_display, status, progress_box]
    # )
#     run_button.click(
#     fn=threaded_batch_runner,
#     inputs=[file_upload, raw_text, user_email],
#     outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
#     every=0.5  # <-- this tells Gradio to expect streaming
# )
    # output_table = gr.HTML()
    # results_group = gr.Group(visible=False)
    # download_file = gr.File(visible=False)
    # usage_display = gr.Markdown(visible=False)
    # status = gr.Markdown(visible=False)
    # progress_box = gr.Textbox(visible=False)

#     run_button.click(
#     fn=threaded_batch_runner,
#     inputs=[file_upload, raw_text, user_email],
#     outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
#     every=0.5,  # streaming enabled
#     show_progress="full"
# )
    print("🎯 DEBUG COMPONENT TYPES")
    print(type(output_table))
    print(type(results_group))
    print(type(download_file))
    print(type(usage_display))
    print(type(status))
    print(type(progress_box))
    

    # interface.stream(
    #     fn=threaded_batch_runner,
    #     inputs=[file_upload, raw_text, user_email],
    #     outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
    #     trigger=run_button,
    #     every=0.5,
    #     show_progress="full",
    # )
    interface.queue()  # No arguments here!

    run_button.click(
        fn=threaded_batch_runner,
        inputs=[file_upload, raw_text, user_email],
        outputs=[output_table, results_group, download_file, usage_display, status, progress_box],
        concurrency_limit=1,  # ✅ correct in Gradio 5.x
        queue=True,            # ✅ ensure the queue is used
        #every=0.5
    )
    



    stop_button.click(fn=stop_batch, inputs=[], outputs=[status])

    # reset_button.click(
    #     #fn=reset_fields,
    #     fn=lambda: (
    #         gr.update(value=""), gr.update(value=""), gr.update(value=None), gr.update(value=None), gr.update(value="Single Accession"),
    #         gr.update(value=[], visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(value="", visible=False), gr.update(value="", visible=False)
    #     ),
    #     inputs=[],
    #     outputs=[
    #         single_accession, raw_text, file_upload, resume_file,inputMode,
    #         output_table,# output_summary, output_flag,
    #         status, results_group, usage_display, progress_box
    #     ]
    # )
    #stop_button.click(fn=lambda sf: (gr.update(value="❌ Stopping...", visible=True), setattr(sf, "value", True) or sf), inputs=[gr.State(stop_flag)], outputs=[status, gr.State(stop_flag)])

    reset_button.click(
        fn=reset_fields,
        inputs=[],
        #outputs=[raw_text, file_upload, resume_file, output_table, status, results_group, usage_display, progress_box]
        outputs=[raw_text, file_upload, output_table, status, results_group, usage_display, progress_box]
    )  

    download_button.click(
      fn=mtdna_backend.save_batch_output, 
      #inputs=[output_table, output_summary, output_flag, output_type], 
      inputs=[output_table, output_type], 
      outputs=[download_file])

    # submit_feedback.click(
    #     fn=mtdna_backend.store_feedback_to_google_sheets, 
    #     inputs=[single_accession, q1, q2, contact], outputs=feedback_status
    # )
    submit_feedback.click(
        fn=mtdna_backend.store_feedback_to_google_sheets,
        inputs=[raw_text, q1, q2, contact],
        outputs=[feedback_status]
    )
    #     # Custom CSS styles
    # gr.HTML("""
    # <style>
    #   /* Ensures both sections are equally spaced with the same background size */
    #   #output-summary, #output-flag {
    #       background-color: #f0f4f8; /* Light Grey for both */
    #       padding: 20px;
    #       border-radius: 10px;
    #       margin-top: 10px;
    #       width: 100%; /* Ensure full width */
    #       min-height: 150px; /* Ensures both have a minimum height */
    #       box-sizing: border-box; /* Prevents padding from increasing size */
    #       display: flex;
    #       flex-direction: column;
    #       justify-content: space-between;
    #   }
      
    #   /* Specific background colors */
    #   #output-summary {
    #       background-color: #434a4b; 
    #   }

    #   #output-flag {
    #       background-color: #141616; 
    #   }

    #   /* Ensuring they are in a row and evenly spaced */
    #   .gradio-row {
    #       display: flex;
    #       justify-content: space-between;
    #       width: 100%;
    #   }
    # </style>
    # """)


interface.launch(share=True,debug=True)