import streamlit as st
from predict_chromosome import predict_and_write


st.title("DeepLoop")
st.write("Looking for some example data?")
st.write("https://huggingface.co/datasets/funlab/HiCorr_test_data")
st.write("Use hg19, HindIII, chr11, CPGZ and LoopDenoise for the Demo data")


#########
# INPUT #
#########
HiCorr_data_hf_repo = st.text_input(
    "HiCorr Data 🤗 Dataset", value="funlab/HiCorr_test_data"
)
training_set = st.selectbox("Select Training Set", ["CPGZ", "H9"], index=0)
depth = st.selectbox("Select Depth", ["LoopDenoise", "50M", "101K"], index=0)
# TODO Throw a warning that h9 only has LoopDenoise and 100M
chromosome = st.selectbox("Select Chromosome", ["chr11"])
prefix = "results"
genome = st.selectbox(
    "Reference Genome",
    [
        "hg19",
        "hg38",
        "mm10",
    ],
)
digestion_enzyme = st.selectbox(
    "Digestion Enzyme",
    [
        "HindIII",
        "Arima",
        # FIXME "microC_5kb_bin"
        "DPNII",
    ],
    index=0,
)
# TODO Add the other toggles
# small_matrix_size
# step_size
# max_dist
# dummy
# val_cols
# keep_zeros

# Load the model from hugging face
from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras(f"funlab/DeepLoop-{training_set}-{depth}")

from huggingface_hub import snapshot_download

HiCorr_data = snapshot_download(HiCorr_data_hf_repo, repo_type="dataset")

anchors = snapshot_download(
    repo_id=f"funlab/{genome}_{digestion_enzyme}_anchor_bed", repo_type="dataset"
)

import os

os.makedirs(prefix, exist_ok=True)

if st.button("Run Deeploop", type="primary"):
    denoised_anchor_to_anchor = None
    with st.spinner("Running the model"):
        denoised_anchor_to_anchor = predict_and_write(
            model,
            full_matrix_dir=HiCorr_data + f"/anchor_2_anchor.loop.{chromosome}",
            input_name=HiCorr_data + f"/anchor_2_anchor.loop.{chromosome}.p_val",
            outdir=prefix,
            anchor_dir=anchors,
            chromosome=chromosome,
            small_matrix_size=128,
            step_size=128,
            dummy=5,
            # max_dist,
            val_cols=["obs", "exp", "pval"],
            # keep_zeros,
        )

    st.success("Done!")

    # Print and list the files
    st.download_button(
        label="Download Denoised Results",
        data=os.path.join(prefix, chromosome + ".denoised.anchor.to.anchor"),
        file_name=chromosome + ".denoised.anchor.to.anchor",
        mime="text/csv",
    )
    st.dataframe(denoised_anchor_to_anchor)
    st.write(os.listdir(prefix))