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


def memory_to_readable(memory_values):
    """
    Convert a list of memory values in bytes to a more readable format uniformly (KB, MB, GB, etc.)
    """
    max_memory = max(memory_values)
    if max_memory < 1024:
        unit = 'bytes'
        scale = 1
    elif max_memory < 1024 ** 2:
        unit = 'KB'
        scale = 1024
    elif max_memory < 1024 ** 3:
        unit = 'MB'
        scale = 1024 ** 2
    elif max_memory < 1024 ** 4:
        unit = 'GB'
        scale = 1024 ** 3
    else:
        unit = 'TB'
        scale = 1024 ** 4

    readable_memory_values = [m / scale for m in memory_values]
    return readable_memory_values, unit


def simplex_transport_memory_requirements(num_locations, num_products):
    SIZE_OF_DOUBLE = 8  # Memory size for a double precision float
    SIZE_OF_INT = 4  # Memory size for an integer


    total_memory = 0

    # Transport item p from i -->j, cannot do both, i --> j and j --> i for the same product
    num_variables = num_products * num_locations * (num_locations - 1) / 2

    # Check that transport for product p can go only in one direction
    # 0 <= t_i,j,p + t_j,i,p < 2
    total_memory += num_variables * 4 * SIZE_OF_INT

    # min_stock <= t_i,j,p
    # t_i,j,p <= max_stock
    total_memory += num_variables * 2 * 2 * SIZE_OF_INT

    total_memory += num_variables * SIZE_OF_DOUBLE # Cost function

    return total_memory


def plot_memory_vs_locations(num_products, max_locations):
    locations_range = range(2, max_locations + 1)  # Need at least two locations
    memory_values = [simplex_transport_memory_requirements(loc, num_products) for loc in locations_range]
    readable_memory_values, unit = memory_to_readable(memory_values)
    fig = px.line(x=locations_range, y=readable_memory_values,
                  labels={'x': 'Number of Locations', 'y': f'Memory ({unit})'},
                  title=f"Memory Footprint vs. Number of Locations (Products = {num_products})")
    return fig


def plot_memory_vs_products(num_locations, max_products):
    products_range = range(1, max_products + 1)
    memory_values = [simplex_transport_memory_requirements(num_locations, prod) for prod in products_range]
    readable_memory_values, unit = memory_to_readable(memory_values)
    fig = px.line(x=products_range, y=readable_memory_values,
                  labels={'x': 'Number of Products', 'y': f'Memory ({unit})'},
                  title=f"Memory Footprint vs. Number of Products (Locations = {num_locations})")
    return fig


# Streamlit application layout
st.title('Transport Network Memory Footprint Calculator')

st.sidebar.header('Configuration')
num_locations = st.sidebar.number_input('Number of Locations', min_value=2, value=3)
num_products = st.sidebar.number_input('Number of Products', min_value=1, value=5)
max_locations = st.sidebar.number_input('Max Locations for Plot', min_value=2, value=10)
max_products = st.sidebar.number_input('Max Products for Plot', min_value=1, value=10)

memory_required = simplex_transport_memory_requirements(num_locations, num_products)
readable_memory, unit = memory_to_readable([memory_required])
st.write(
    f"Estimated memory requirement for {num_locations} locations and {num_products} products: {readable_memory[0]:.2f} {unit}")

st.plotly_chart(plot_memory_vs_locations(num_products, max_locations))
st.plotly_chart(plot_memory_vs_products(num_locations, max_products))