utils/convert_to_excel.py

import io
import time

import pandas as pd
import streamlit as st

# @st.cache_data
# def convert_dfs(dfs: list[pd.DataFrame], sheet_names: list[str]) -> bytes:
#     # IMPORTANT: Cache the conversion to prevent computation on every rerun

#     # Create a BytesIO object
#     bytes_io = io.BytesIO()

#     # Write the dataframes to the BytesIO object
#     with pd.ExcelWriter(bytes_io, engine="xlsxwriter") as writer:
#         for df, sheet_name in zip(dfs, sheet_names):
#             df.to_excel(writer, sheet_name=sheet_name, index=True)

#     # Get the bytes data
#     bytes_data = bytes_io.getvalue()

#     # Close the BytesIO object
#     bytes_io.close()

#     return bytes_data


def get_formats(workbook):
    return {
        "green": workbook.add_format(
            {"bg_color": "#37CC73", "bold": True, "border": 1}
        ),
        "blue": workbook.add_format({"bg_color": "#1A64FF", "bold": True, "border": 1}),
        "blue_light": workbook.add_format(
            {"bg_color": "#00B0F0", "bold": True, "border": 1}
        ),
        "beurre": workbook.add_format(
            {"bg_color": "#FFE699", "bold": True, "border": 1}
        ),
        "orange": workbook.add_format(
            {"bg_color": "#F47F31", "bold": True, "border": 1}
        ),
        "purple5": workbook.add_format(
            {"bg_color": "#E03DCD", "bold": True, "border": 1}
        ),
        "purple6": workbook.add_format(
            {"bg_color": "#AE83F8", "bold": True, "border": 1}
        ),
        "gray": workbook.add_format({"bg_color": "#D9D9D9", "bold": True, "border": 1}),
        "red": workbook.add_format({"bg_color": "#FF0000", "bold": True, "border": 1}),
    }


def get_format_map_by_format_type(formats: dict, format_type: str) -> dict:
    if format_type == "GSM_Analysis":
        return {
            # "name": formats["blue"],
            "amrSegLoadDepTchRateLower": formats["beurre"],
            "amrSegLoadDepTchRateUpper": formats["beurre"],
            "dedicatedGPRScapacity": formats["beurre"],
            "defaultGPRScapacity": formats["beurre"],
            "number_trx_per_cell": formats["blue_light"],
            "number_trx_per_bcf": formats["blue_light"],
            "number_tch_per_cell": formats["blue"],
            "number_sd_per_cell": formats["blue"],
            "number_bcch_per_cell": formats["blue"],
            "number_ccch_per_cell": formats["blue"],
            "number_cbc_per_cell": formats["blue"],
            "number_total_channels_per_cell": formats["blue"],
            "number_signals_per_cell": formats["blue"],
            "hf_rate_coef": formats["purple5"],
            "GPRS": formats["purple5"],
            "TCH Actual HR%": formats["green"],
            "Offered Traffic BH": formats["green"],
            "Max_Traffic BH": formats["green"],
            "Avg_Traffic BH": formats["green"],
            "Max_tch_call_blocking BH": formats["green"],
            "Avg_tch_call_blocking BH": formats["green"],
            "number_of_days_with_tch_blocking_exceeded": formats["green"],
            "Max_sdcch_real_blocking BH": formats["green"],
            "Avg_sdcch_real_blocking BH": formats["green"],
            "number_of_days_with_sdcch_blocking_exceeded": formats["green"],
            "TCH UTILIZATION (@Max Traffic)": formats["orange"],
            "Target FR CHs": formats["purple6"],
            "Target HR CHs": formats["purple6"],
            "Target TCHs": formats["purple6"],
            "Target TRXs": formats["purple6"],
            "Numberof required TRXs": formats["purple6"],
        }
    elif format_type == "database":
        return {
            "code": formats["blue"],
            "Azimut": formats["green"],
            "Longitude": formats["green"],
            "Latitude": formats["green"],
            "Hauteur": formats["green"],
            "number_trx_per_cell": formats["blue_light"],
            "number_trx_per_bcf": formats["blue_light"],
            "number_trx_per_site": formats["blue_light"],
        }
    # elif format_type == "LTE":
    #     return {
    #         "DL PRB Utilization": formats["orange"],
    #         "UL PRB Utilization": formats["orange"],
    #         "RSRP": formats["blue_light"],
    #         "RSRQ": formats["blue_light"],
    #         "Throughput (Mbps)": formats["green"],
    #     }
    else:
        return {}  # No formatting if format_type not matched


def _apply_custom_formatting(
    writer, df: pd.DataFrame, sheet_name: str, format_type: str
):
    workbook = writer.book
    worksheet = writer.sheets[sheet_name]

    formats = get_formats(workbook)
    format_map = get_format_map_by_format_type(formats, format_type)

    for col_idx, col_name in enumerate(df.columns):
        fmt = format_map.get(col_name)
        if fmt:
            worksheet.write(0, col_idx + 1, col_name, fmt)


def _write_to_excel(
    dfs: list[pd.DataFrame], sheet_names: list[str], index=True, format_type: str = None
) -> bytes:
    bytes_io = io.BytesIO()
    with pd.ExcelWriter(bytes_io, engine="xlsxwriter") as writer:
        for df, name in zip(dfs, sheet_names):
            # df.index.name = "index"
            df.to_excel(writer, sheet_name=name, index=index)
            if format_type:
                _apply_custom_formatting(writer, df, name, format_type)
    return bytes_io.getvalue()


@st.cache_data
def convert_dfs(dfs: list[pd.DataFrame], sheet_names: list[str]) -> bytes:
    return _write_to_excel(dfs, sheet_names, index=True)


@st.cache_data
def convert_gsm_dfs(dfs, sheet_names) -> bytes:
    return _write_to_excel(dfs, sheet_names, index=True, format_type="GSM_Analysis")


@st.cache_data
def convert_database_dfs(dfs, sheet_names) -> bytes:
    return _write_to_excel(dfs, sheet_names, index=True, format_type="database")


def save_dataframe(df: pd.DataFrame, sheet_name: str):
    """
    Save the dataframe to a csv file.

    Args:
        df (pd.DataFrame): The dataframe to save.
        sheet_name (str): The name of the sheet.
    """
    df.to_csv(f"data2/{sheet_name}_{time.time()}.csv", index=False, encoding="latin1")