import plotly.graph_objects as go
import streamlit


@streamlit.cache_data
def plot_sankey(data, output_file='sankey_diagram_oct.html', i=0, do='до'):
    labels = []
    sources = []
    targets = []
    values = []

    # Сбор данных для диаграммы
    for event, next_events in data.items():
        # Добавляем основное событие в метки
        if event not in labels:
            labels.append(event)

        current_event_index = labels.index(event)

        for next_event, (user_count, percentage) in next_events.items():
            if next_event not in labels:
                labels.append(next_event)

            next_event_index = labels.index(next_event)
            sources.append(current_event_index)
            targets.append(next_event_index)
            values.append(user_count)

    # Создание диаграммы
    fig = go.Figure(data=[go.Sankey(
        valueformat= ".0f",
        valuesuffix="TWh",

        # font = dict(family='Arial',
        #             size=14,
        #             color='black'),
        node=dict(
            # family='Arial',
            pad=15,
            thickness=20,
            line=dict(color="black", width=0.1),
            label=labels,
            # font = dict(size=15, color="black")
        ),
        link=dict(
            # family = 'Arial',
            source=sources,
            target=targets,
            value=values,
            label=[f"{user_count} ({percentage})" for user_count, percentage in
                   zip(values, [data[labels[sources[i]]][labels[targets[i]]][1] for i in range(len(sources))])],
        )
    )])

    return fig.update_layout(title_text=f"Воронка сервиса Финансовое здоровье за октябрь 2024, часть {i} ({do} ГЭС ФЗ)",
                      font_size=10,
                             font= dict(size=10, color='white'),
                             plot_bgcolor='white',
                             paper_bgcolor='white')

    #     Сохранение диаграммы в HTML файл
    # fig.write_html(f'test_sanki_oct_part{i}.html')
    #     print(f"Диаграмма успешно сохранена в файл: test_sanki.html")
    # fig.show()