import numpy as np
import pandas as pd
import plotly.express as px
import streamlit as st


def saturated_mutations(dna_seq):
    dna_seq = dna_seq.upper().replace("U", "T")
    dna_seq = "".join(c for c in dna_seq if c in "ATCG")
    n = len(dna_seq)
    bases = np.array(list("ATCG"))
    seq_array = np.array(list(dna_seq))
    replace_matrix = np.array([bases[bases != base] for base in seq_array])
    result = np.tile(seq_array, (3 * n, 1))
    rows = np.repeat(np.arange(n), 3)
    result[np.arange(3 * n), rows] = replace_matrix.flatten()
    result = np.vstack([seq_array, result])
    return ["".join(row) for row in result]


if __name__ == "__main__":
    model = st.session_state["model"]
    input_format = st.session_state["input_format"]
    predict = st.session_state["predict"]
    mrna_env = st.session_state["mrna_env"]
    UTR3_GLUC = st.session_state["UTR3_GLUC"]
    CDS_GLUC = st.session_state["CDS_GLUC"]
    st.set_page_config(layout="wide")
    st.title("UTR-CODE")
    with st.form("single_input_form"):
        utr5_input = st.text_area(
            "5' UTR:",
            key="utr5_input",
            max_chars=1000,
            help="Only 'A', 'C', 'G', 'T' and 'U' is allowed, other characters will be ignored",
        )
        with st.expander("Optional Input"):
            optional_input_col1, optional_input_col2 = st.columns(2)
            with optional_input_col1:
                cds_input = st.text_area(
                    "CDS",
                    value=CDS_GLUC,
                    key="cds_input",
                    max_chars=4500,
                    help="GLUC's CDS is used by default.",
                )
            with optional_input_col2:
                utr3_input = st.text_area(
                    "3' UTR:",
                    value=UTR3_GLUC,
                    key="utr3_input",
                    max_chars=4000,
                    help="GLUC's UTR3 is used by default.",
                )
        mutation_rounds = st.slider(
            "ISM Rounds",
            min_value=1,
            max_value=10,
            value=5,
            step=1,
            label_visibility="visible",
        )
        submitted = st.form_submit_button("Submit", use_container_width=True)
    if submitted:
        if not "".join([c for c in utr5_input.upper() if c in "ATCGU"]):
            st.warning("Please input 5' UTR")
        elif any(c not in "ATCGU\n\t, " for c in f"{utr5_input.upper()}{cds_input.upper()}{utr3_input.upper()}"):
            st.warning("Only 'A', 'C', 'G', 'T' and 'U' is allowed.")
        else:
            utr5_optimize_use = "".join(
                [c for c in utr5_input.upper().replace("U", "T") if c in "ATCG"]
            )
            results = dict()
            utr, cds = input_format(utr5_optimize_use, cds_input, utr3_input)
            results[0] = {
                "rpf": (np.expm1(predict(model, utr, cds, mrna_env)[0]) / 5) / 108.7144 ,
                "seq": utr5_optimize_use,
                "base_change": "Origin",
            }
            cds_seq = [
                cds_input,
            ] * (3 * len(utr5_input) + 1)
            utr3_seq = [
                utr3_input,
            ] * (3 * len(utr5_input) + 1)
            progress_bar = st.empty()
            progress_bar = st.progress(0, f"Progress: 0%")
            progress_text = st.empty()
            chart_placeholder = st.empty()
            for i in range(1, mutation_rounds + 1):
                used_seq = saturated_mutations(utr5_optimize_use)
                utr, cds = input_format(used_seq, cds_seq, utr3_seq)
                result = predict(model, utr, cds, mrna_env)
                best_idx = np.argmax(result)
                max_val = result[best_idx]
                base_change = next(
                    (
                        f"{i+1}:{c1}->{c2}"
                        for i, (c1, c2) in enumerate(
                            zip(utr5_optimize_use, used_seq[best_idx])
                        )
                        if c1 != c2
                    ),
                    "No change",
                )
                utr5_optimize_use = used_seq[best_idx]
                results[i] = {
                    "rpf": (np.expm1(max_val) / 5) / 108.7144,
                    "seq": utr5_optimize_use,
                    "base_change": base_change,
                }
                progress_bar.progress(
                    i / mutation_rounds,
                    f"Progress: {i / mutation_rounds * 100:.0f}%",
                )
                df = pd.DataFrame(
                    [
                        {
                            "Round": k,
                            "Predicted Value": float(v["rpf"]),
                            "UTR5 Sequence": v["seq"],
                            "Base Change": v["base_change"],
                        }
                        for k, v in results.items()
                    ]
                )
                fig = px.line(
                    df,
                    x="Round",
                    y="Predicted Value",
                    text="Base Change",
                    markers=True,
                )
                fig.update_layout(
                    template="plotly_white",
                    title="UTR5 Optimization Progress",
                    title_font_family="Arial, sans-serif",
                    title_font_color="black",
                    xaxis_title="Mutation Round",
                    yaxis_title="Predicted Value (RPF)",
                    xaxis_title_font=dict(color="black"),
                    yaxis_title_font=dict(color="black"),
                    font_family="Arial, sans-serif",
                    font_color="black",
                    margin=dict(l=50, r=50, t=80, b=50),
                )
                fig.update_xaxes(
                    showgrid=False,
                    showline=True,
                    linecolor="black",
                    tickcolor="black",
                    ticks="outside",
                    tickfont_color="black",
                )
                fig.update_yaxes(
                    showgrid=False,
                    showline=True,
                    linecolor="black",
                    tickcolor="black",
                    ticks="outside",
                    tickfont_color="black",
                    tickformat=".2f",
                )
                fig.update_traces(
                    textposition="top center",
                    marker=dict(color="#cf3326", line=dict(color="#cf3326")),
                    line=dict(color="#cf3326"),
                    textfont_color="black",
                )
                fig.update_xaxes(dtick=1)
                chart_placeholder.plotly_chart(fig, use_container_width=True)
            progress_bar.empty()
            st.download_button(
                "Download Result",
                data=df.to_csv(index=False),
                file_name="results.csv",
                mime="text/csv",
                use_container_width=True,
                icon="📥",
            )