import gradio as gr
import pandas as pd

quants = (
    pd.read_csv("quants.csv")
    .applymap(str)
    .groupby("quant")["bpw"]
    .apply(float)
    .to_dict()
)
models = (
    pd.read_csv("models.csv")
    .applymap(str)
    .groupby("model")["params"]
    .apply(float)
    .to_dict()
)


def context_sizes(model):
    return pd.read_csv(
        "context_sizes/" + model.replace("/", "_") + ".csv",
        header=None,
        names=["context", "size"],
    )


def linear_regression(xs, ys) -> tuple[float, float]:
    sum_y = ys.sum()
    sum_x = sum(xs)
    sum_xy = sum([x * y for x, y in zip(xs, ys)])
    sum_x2 = sum([x**2 for x in xs])
    n = len(xs)

    a = (sum_y * sum_x2 - sum_x * sum_xy) / (n * sum_x2 - sum_x**2)
    b = (n * sum_xy - sum_x * sum_y) / (n * sum_x2 - sum_x**2)
    return a, b


def calc_model_size(parameters: float, quant: float) -> float:
    return round(parameters * quant / 8, 2)


def calc_context_size(context, model) -> float:
    sizes = context_sizes(model)
    a, b = linear_regression(sizes["context"], sizes["size"])
    return round((a + b * context) / 1024, 2)


def calc(model_base, context, quant_size):
    model_params = models[model_base]
    quant_bpw = quants[quant_size]

    model_size = calc_model_size(model_params, quant_bpw)
    context_size = calc_context_size(context, model_base)

    return model_size, context_size, model_size + context_size


title = "GGUF VRAM Calculator"

with gr.Blocks(title=title, theme=gr.themes.Monochrome()) as app:
    default_model = "Mistral 7B"
    default_quant = "Q4_K_S"
    default_context = 8192
    default_model_size = calc_model_size(models[default_model], quants[default_quant])
    default_context_size = calc_context_size(default_context, default_model)

    gr.Markdown(f"# {app.title}")
    model = gr.Dropdown(
        list(models.keys()), value=default_model, label="Select Model Base"
    )
    context = gr.Number(minimum=1, value=default_context, label="Context Size (Tokens)")
    quant = gr.Dropdown(
        list(quants.keys()), value=default_quant, label="Select Quant Size"
    )
    btn = gr.Button(value="Submit", variant="primary")
    btn.click(
        calc,
        inputs=[
            model,
            context,
            quant,
        ],
        outputs=[
            gr.Number(
                label="Model Size (GB)",
                value=default_model_size,
            ),
            gr.Number(
                label="Context Size (GB)",
                value=default_context_size,
            ),
            gr.Number(
                label="Total Size (GB)",
                value=default_model_size + default_context_size,
            ),
        ],
    )

    app.launch()