import streamlit as st
import pandas as pd
from model_util import get_mode_from_hf, calculate_flops_in_hugging_space


def calculate_flops_architecture(num_params, num_tokens):
    """
    Calculate FLOPs based on model architecture.
    """
    total_flops = 6 * num_params * num_tokens
    threshold_flops = 10 ** 25
    meets_criteria = total_flops > threshold_flops

    return {
        "total_flops": total_flops,
        "meets_criteria": meets_criteria
    }


def calculate_flops_gpu(gpu_hours, power_consumption_w, flops_per_gpu_s):
    """
    Calculate FLOPs based on GPU hours and type.
    """
    threshold_flops = 10 ** 25
    total_energy_wh = gpu_hours * power_consumption_w
    total_flops = gpu_hours * flops_per_gpu_s * 3600
    meets_criteria = total_flops > threshold_flops

    return {
        "total_energy_wh": total_energy_wh,
        "total_flops": total_flops,
        "meets_criteria": meets_criteria
    }


def calculate_flops_hf(model_name, input_shape, access_token, bp_factor):
    """
    Calculate FLOPs using Hugging Face model information.
    """
    model = get_mode_from_hf(model_name=model_name, library="auto", access_token=access_token)
    data, return_print = calculate_flops_in_hugging_space(model_name=model_name, empty_model=model,
                                                          access_token=access_token, input_shape=input_shape,
                                                          bp_factor=bp_factor)#, output_unit=output_unit)
    print(f"Data: {data}")

    total_flops = data[0]['Forward+Backward FLOPs']
    threshold_flops = 10 ** 25
    meets_criteria = total_flops > threshold_flops

    return {
        "total_flops": total_flops,
        "meets_criteria": meets_criteria,
        "dataframe": data,
        "return_print": return_print
    }


st.title("FLOPs Calculator for EU AI Act High Impact Capabilities")

# Choose calculation method
calculation_method = st.selectbox("Choose Calculation Method:",
                                  ["Model Architecture", "GPU Hours and Type", "Hugging Face Model"])

if calculation_method == "Model Architecture":
    num_params = st.number_input("Number of Parameters (N):", min_value=0.0, value=float(7.0 * 10 ** 9), step=1.0)
    num_tokens = st.number_input("Number of Training Tokens (D):", min_value=0.0, value=float(1500 * 10 ** 9), step=1.0)

    if st.button("Calculate FLOPs (Model Architecture)"):
        result = calculate_flops_architecture(num_params, num_tokens)
        st.write(f"Total FLOPs: {result['total_flops']:.2e} FLOPs")
        st.write(f"Meets high impact capabilities criteria: {result['meets_criteria']}")

elif calculation_method == "GPU Hours and Type":
    # Define GPU types and their corresponding FLOPs per second
    gpu_types = {
        "Nvidia A100": {"flops_per_s": 312 * 10 ** 12, "power": 400},
        "Nvidia V100": {"flops_per_s": 125 * 10 ** 12, "power": 300},
        "Nvidia H100": {"flops_per_s": 1.25 * 10 ** 15, "power": 700},
        "Nvidia T4": {"flops_per_s": 65 * 10 ** 12, "power": 70}
    }

    gpu_type = st.selectbox("Select GPU type:", list(gpu_types.keys()))
    flops_per_gpu_s = gpu_types[gpu_type]["flops_per_s"]
    power_consumption_w = gpu_types[gpu_type]["power"]

    gpu_hours = st.number_input("Total GPU hours used for training:", min_value=0.0, value=float(7.7 * 10 ** 6),
                                step=1.0)

    if st.button("Calculate FLOPs (GPU Hours and Type)"):
        result = calculate_flops_gpu(gpu_hours, power_consumption_w, flops_per_gpu_s)
        st.write(f"Total energy consumption: {result['total_energy_wh']:.2e} Wh")
        st.write(f"Total FLOPs: {result['total_flops']:.2e} FLOPs")
        st.write(f"Meets high impact capabilities criteria: {result['meets_criteria']}")

elif calculation_method == "Hugging Face Model":
    model_name = st.text_input("Model Name:", "tiiuae/falcon-40b")
    batch_size = st.number_input("Batch Size:", min_value=1, value=1)
    max_seq_length = st.number_input("Max Sequence Length:", min_value=1, value=128)
    input_shape = (batch_size, max_seq_length)
    access_token = st.text_input("Hugging Face Access Token:", "")
    bp_factor = st.number_input("Backward Pass Factor (BP Factor):", min_value=0.0, value=1.0, step=0.1)
    #output_unit = st.selectbox("Output Unit:", ["auto", "MFLOPs", "GFLOPs", "TFLOPs", "PFLOPs"])

    if st.button("Calculate FLOPs (Hugging Face Model)"):
        result = calculate_flops_hf(model_name, input_shape, access_token, bp_factor) #output_unit)
        st.write(f"Total FLOPs: {result['total_flops']:.2e} FLOPs")
        st.write(f"Meets high impact capabilities criteria: {result['meets_criteria']}")
        st.write("Detailed FLOPs Data:")
        st.dataframe(result["dataframe"])
        st.text(result["return_print"])