import numpy as np
import streamlit as st

st.title("GCP Resource Alloctor")

st.subheader("Configuration")

# GPU Type
gpu_type = st.selectbox(
    'GPU Type',
    (  
        'H100 80GB', 
        'A100 80GB',
        'A100 40GB',
        'V100 16GB',
        'P100 16GB',
        'L4 24GB',
        'T4 16GB',
        'P4 8GB',
    )
)

# Number of GPUs
gpu_number_mapping = {
    'H100 80GB': [8],
    'A100 80GB': [1,2,4,8],
    'A100 40GB': [1,2,4,8,16],
    'V100 16GB': [1,2,4,8],
    'P100 16GB': [1,2,4],
    'L4 24GB': [1,2,4,8],
    'T4 16GB': [1,2,4],
    'P4 8GB': [1,2,4],
}
gpu_number = st.selectbox('Number of GPUs', gpu_number_mapping[gpu_type])

# Instance Type
gpu_type_mapping = {
    'H100 80GB': ["A3"],
    'A100 80GB': ["A2"],
    'A100 40GB': ["A2"],
    'V100 16GB': ["N1", "CUSTOM"],
    'P100 16GB': ["N1", "CUSTOM"],
    'L4 24GB': ["G2", "CUSTOM"],
    'T4 16GB': ["N1", "CUSTOM"],
    'P4 8GB': ["N1", "CUSTOM"],
}
instance_type = st.selectbox('Instance Type', gpu_type_mapping[gpu_type])

# CPU Cores
cpu_cores_mapping = {
    "A3": [208],
    "A2": [12*gpu_number],
    "G2": [12*gpu_number] if gpu_number > 1 else [4,8,12,16,32],
    "N1": [1,2,4,8,16,32,96],
    "CUSTOM": [1] + [i for i in range(2, 96+1, 2)]
}
if gpu_type != "CUSTOM":
    cpu_cores = st.selectbox('Cores (vCPU)', cpu_cores_mapping[instance_type])
else:
    cpu_cores = st.select_slider('Cores (vCPU)', cpu_cores_mapping[instance_type])

# Memory Size
memory_size_mapping = {
    "A3": [1872],
    "A2": [170*gpu_number],
    "G2": [4*cpu_cores] if gpu_number > 1 else [48*gpu_number],
    "N1": [cpu_cores*3.75],
    "CUSTOM": [i for i in np.arange(cpu_cores, cpu_cores*6.5+1, 1)]
}
if gpu_type != "CUSTOM":
    memory_size = st.selectbox('Memory (GB)', memory_size_mapping[instance_type])
else:
    memory_size = st.select_slider('Memory (GB)', memory_size_mapping[instance_type])

# Balanced Disk
balanced_disk_size = st.select_slider('Balanced Disk (GB)', [i for i in range(10, 65536, 10)])

# SSD Disk
ssd_disk_size = st.select_slider('SSD Disk (GB)', [i * 375 for i in [0,1,2,3,4,5,6,7,8,16,24]])

# Hours
hours = st.select_slider('Duration (Hours)', [i for i in range(1, 168+1)])

# Pricing Estimate
serivces_mapping = {
    "Core": {
        "A3": 0.029917642,
        "A2": 0.017880447,
        "G2": 0.016626389,
        "N1": 0.007834495,
        "CUSTOM": 0.00782101,
    },
    "RAM": {
        "A3": 0.002605197,
        "A2": 0.002396196,
        "G2": 0.00194851,
        "N1": 0.001049094,
        "CUSTOM": 0.001047746,
    },
    "GPU": {
        'H100 80GB': 12.112232328,
        'A100 80GB': 2.61383548,
        'A100 40GB': 1.67288707,
        'V100 16GB': 0.997853,
        'P100 16GB': 0.5798335,
        'L4 24GB': 0.279501996,
        'T4 16GB': 0.1483295,
        'P4 8GB': 0.29800745,
    },
    "PD": 0.1483295 / 30 / 24,
    "SSD": 0.108550225 / 30 / 24,
}

core_price = serivces_mapping['Core'][instance_type] * cpu_cores
memory_price = serivces_mapping['RAM'][instance_type] * memory_size
gpu_price = serivces_mapping['GPU'][gpu_type] * gpu_number
balanced_disk_price = serivces_mapping['PD'] * balanced_disk_size
ssd_disk_price = serivces_mapping['SSD'] * ssd_disk_size
duration_total_price = core_price + memory_price + gpu_price + balanced_disk_price + ssd_disk_price
total_price = duration_total_price * hours

st.divider()

st.subheader("Hourly estimate")
st.write(f"Core: SGD :blue[{core_price}]")
st.write(f"Memory: SGD :blue[{memory_price}]")
st.write(f"GPU: SGD :blue[{gpu_price}]")
st.write(f"Balance Disk: SGD :blue[{balanced_disk_price}]")
st.write(f"SSD Disk: SGD :blue[{ssd_disk_price}]")
st.write(f"Hourly Total: SGD :blue[{duration_total_price}]")
st.write(f"Duration Total: SGD :blue[{total_price}]")