import streamlit as st
from utils import memory_moe_mlp, memory_mlp_layer, memory_for_attention_layer


def main():
    st.title("LLM Model Memory Usage Calculator")

    st.sidebar.header("Model Parameters")
    precession = st.sidebar.number_input("precession in Byte", min_value=1, max_value=4, value=2, step=2)
    hidden_size = st.sidebar.number_input("Hidden Size", min_value=512, max_value=2 ** 16, value=4096, step=512)
    num_heads = st.sidebar.number_input("Number of Attention Heads", min_value=4, max_value=128, value=32, step=4)
    batch_size = st.sidebar.number_input("Batch Size", min_value=1, max_value=1024, value=64, step=4)
    seq_len = st.sidebar.number_input("Sequence Length", min_value=512, max_value=128000, value=2048, step=512)
    intermediate_size = st.sidebar.number_input("Intermediate Size", min_value=1024, max_value=2 ** 18, value=11008,
                                                step=128)
    layers = st.sidebar.number_input("Number of Layers", min_value=6, max_value=48, value=30, step=1)
    moe = st.sidebar.checkbox("Use Mixture of Experts (MOE)", value=False)

    # Conditional rendering for MOE parameters
    if moe:
        top_k = st.sidebar.number_input("Number of Experts to use (Top K)", min_value=1, max_value=16, value=2, step=1)
        num_experts = st.sidebar.number_input("Total Number of Experts", min_value=2, max_value=32, value=4, step=2)
    else:
        top_k = 2  # Default values if MOE is not used
        num_experts = 4

    attention_memory = memory_for_attention_layer(precession,
                                                  seq_len,
                                                  batch_size,
                                                  hidden_size,
                                                  num_heads)

    dense_mlp_memory = memory_mlp_layer(precession,
                                        seq_len,
                                        batch_size,
                                        hidden_size,
                                        intermediate_size)

    dense_model_memory = layers * (attention_memory + dense_mlp_memory) // (1024 ** 3)
    st.write("This estimation is for the training phase,  all computation done for AdamW optimizer.")
    space = st.empty()
    space.markdown('<div style="height: 20px;"></div>', unsafe_allow_html=True)

    st.markdown(
        f'<div style="background-color: #b3f0ff; padding: 30px; border-radius: 5px;">'
        f'<p style="font-weight: bold;">The memory requirement for this model is ~ {dense_model_memory} GB</p>'
        f'</div>',
        unsafe_allow_html=True
    )
    space = st.empty()
    space.markdown('<div style="height: 40px;"></div>', unsafe_allow_html=True)

    if moe:
        moe_memory = memory_moe_mlp(precession,
                                    seq_len,
                                    batch_size,
                                    hidden_size,
                                    intermediate_size,
                                    num_experts,
                                    top_k)
        moe_model = layers * (attention_memory + moe_memory) // (1024 ** 3)

        st.markdown(
            f'<div style="background-color: #99ff99; padding: 30px; border-radius: 5px;">'
            f'<p style="font-weight: bold;">The memory requirement for the MOE model is ~ {moe_model} GB</p>'
            f'</div>',
            unsafe_allow_html=True
        )

    space = st.empty()
    space.markdown('<div style="height: 40px;"></div>', unsafe_allow_html=True)

    st.markdown(
        f'<div style="background-color: #f0f0f0; padding: 30px; border-radius: 5px;">'
        f'<p style="font-weight: bold;">For more information please read this article</p>'
        f'<a href="https://medium.com/@khalil.hennara.247/llm-memory-usage-f62a007a509c">Article</a>'
        f'</div>',
        unsafe_allow_html=True
    )


if __name__ == "__main__":
    main()