import streamlit as st
import plotly.graph_objects as go
import pandas as pd
from datetime import datetime, time
from datasets import load_dataset
import data

st.set_page_config(layout="wide")

st.write("# Time Window Recommendation")

token = data.token

@st.cache_data(show_spinner=False)
def load_data():
    dataset = load_dataset('skt-asap/labeled-dataset', data_files="Labeled.csv", token=token)
    df = dataset['train'].to_pandas()
    df.fillna(0, inplace=True)

    df['timestamp'] = pd.to_datetime(df['timestamp'])

    return df

df = load_data()

unique_cells = df['enbid_pci'].unique().tolist()
selected_enbid_pci = st.selectbox('셀 ID:', unique_cells)

available_dates = df['timestamp'].dt.date.unique()
selected_date = st.date_input('날짜 선택:', min_value=min(available_dates), max_value=max(available_dates), value=min(available_dates))

filtered_df = df[(df['timestamp'].dt.date == selected_date) & (df['enbid_pci'] == selected_enbid_pci)]

timeline_data = []
status_columns = ['Status_2100', 'Status_2600_10', 'Status_2600_20']

for col in status_columns:
    daily_status = filtered_df[['timestamp', col]].set_index('timestamp').resample('15T').first()
    daily_status['start_time'] = daily_status.index
    daily_status['end_time'] = daily_status['start_time'] + pd.Timedelta(minutes=15)

    for _, row in daily_status.iterrows():
        if pd.notna(row[col]):  # NaN 값 처리
            timeline_data.append(dict(
                Task=col,
                Start=row['start_time'],
                Finish=row['end_time'],
                Resource='ON' if row[col] == 1 else 'OFF'
            ))

timeline_df = pd.DataFrame(timeline_data)

fig = go.Figure()

for status in status_columns:
    status_data = timeline_df[timeline_df['Task'] == status]
    if not status_data.empty:
        fig.add_trace(go.Bar(
            x=[(d['Finish'] - d['Start']).total_seconds() / 3600 for d in status_data.to_dict('records')],
            y=[status] * len(status_data),
            orientation='h',
            name=status,
            text=status_data['Resource'],
            textposition='inside',
            insidetextanchor='middle',
            marker=dict(color=['#FF8181' if r == 'OFF' else '#7AA1EE' for r in status_data['Resource']])
        ))

fig.update_layout(
    title=f'일일 ON/OFF 타임라인 - {selected_enbid_pci} ({selected_date})',
    xaxis_title='시간',
    yaxis_title='주파수 대역',
    height=400,
    barmode='stack',
    xaxis=dict(
        tickmode='array',
        tickvals=[i for i in range(0, 25, 3)],
        ticktext=[f'{i:02d}:00' for i in range(0, 25, 3)]
    ),
    legend_title_text='주파수 대역',
    margin=dict(l=0, r=0, t=50, b=0)
)

fig.update_traces(textfont_size=10)

# Streamlit 앱에 차트 추가
st.plotly_chart(fig, use_container_width=True)