import streamlit as st
from datasets import load_dataset
import pandas as pd
from datetime import datetime, time
from streamlit_folium import folium_static as st_folium
import map_recommend
import data
import os

token = data.token

st.set_page_config(layout="wide")

st.write('# Recommendation By Time')

# 세션 상태 초기화
if 'selected_model' not in st.session_state:
    st.session_state.selected_model = None

if 'date_input' not in st.session_state:
    st.session_state.date_input = datetime(2024, 4, 1)

if 'time_input' not in st.session_state:
    st.session_state.time_input = time(0, 0)

if 'frequency_input' not in st.session_state:
    st.session_state.frequency_input = '2100'

if 'recommended_cell_states' not in st.session_state:
    st.session_state.recommended_cell_states = None

if 'uploaded_csv' not in st.session_state:
    st.session_state.uploaded_csv = None

if 'map_data' not in st.session_state:
    st.session_state.map_data = None

with st.sidebar:
    with st.spinner('데이터 로딩 중...'):
        df, df_map = data.load_data()
        st.sidebar.success('데이터 로드 완료')

def load_and_predict(model_name, df_map, timestamp_input, frequencies, progress_callback=None):
    recommended_cell_states = {freq: {} for freq in frequencies}

    if progress_callback:
        progress_callback(10)

    if model_name == 'CSV File' and st.session_state.uploaded_csv is not None:
        df = st.session_state.uploaded_csv
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        timestamp_input = pd.to_datetime(timestamp_input)
        st.info(f"**{df['timestamp'].min().strftime('%Y/%m/%d')} ~ {df['timestamp'].max().strftime('%Y/%m/%d')}** 사이의 날짜를 입력해주세요.")
        if timestamp_input not in df['timestamp'].values:
            st.error('CSV 파일에 해당 타임스탬프가 없습니다.')
            st.stop()
        else:
            df = df[df['timestamp'] == timestamp_input]
    elif model_name == 'Random':
        with st.spinner('Random 데이터셋 로딩 중...'):
            dataset = load_dataset('skt-asap/examples', data_files='Random.csv', token=token)
        df = dataset['train'].to_pandas()
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        timestamp_input = pd.to_datetime(timestamp_input)
        st.info(f"**{df['timestamp'].min().strftime('%Y/%m/%d')} ~ {df['timestamp'].max().strftime('%Y/%m/%d')}** 사이의 날짜를 입력해주세요.")
        if timestamp_input not in df['timestamp'].values:
            st.error('CSV 파일에 해당 타임스탬프가 없습니다.')
            st.stop()
        else:
            df = df[df['timestamp'] == timestamp_input]

    for freq in frequencies:
        if model_name == 'Time-based':
            for idx, row in df_map.iterrows():
                if timestamp_input.hour < 6 or timestamp_input.hour > 22:
                    recommended_cell_states[freq][row['enbid_pci']] = 'OFF'
                else:
                    recommended_cell_states[freq][row['enbid_pci']] = 'ON'
                if progress_callback:
                    progress_callback(10 + 80 * (idx + 1) // len(df_map))
        elif model_name == 'CSV File':
            if st.session_state.uploaded_csv is not None:
                for idx, row in df.iterrows():
                    if row[f'Status_{freq}'] == 0:
                        recommended_cell_states[freq][row['enbid_pci']] = 'OFF'
                    elif row[f'Status_{freq}'] == 1:
                        recommended_cell_states[freq][row['enbid_pci']] = 'ON'
                    else:
                        recommended_cell_states[freq][row['enbid_pci']] = 'UNKNOWN'
            else:
                st.error('CSV 파일을 업로드해주세요.')
                st.stop()
        elif model_name == 'Random':
            for idx, row in df.iterrows():
                if row[f'Status_{freq}'] == 0:
                    recommended_cell_states[freq][row['enbid_pci']] = 'OFF'
                elif row[f'Status_{freq}'] == 1:
                    recommended_cell_states[freq][row['enbid_pci']] = 'ON'
                else:
                    recommended_cell_states[freq][row['enbid_pci']] = 'UNKNOWN'

        else:
            st.error('올바른 모델을 선택해주세요.')
            st.stop()

    return recommended_cell_states

frequencies_to_check = ['2100', '2600_10', '2600_20']

model_options = ['CSV File', 'Random', 'Time-based']
selected_model = st.selectbox('모델:', model_options, index=model_options.index(st.session_state.selected_model) if st.session_state.selected_model else 0)

if selected_model == 'CSV File':
    uploaded_file = st.file_uploader("CSV File Upload", type="csv", label_visibility='collapsed')
    if uploaded_file is not None:
        st.session_state.uploaded_csv = pd.read_csv(uploaded_file)

with st.form('form'):
    date_input = st.date_input('날짜:', st.session_state.date_input)
    time_input = st.time_input('시간:', st.session_state.time_input)

    frequency_options = ['2100MHz', '2600MHz - 10', '2600MHz - 20']
    frequency_input = st.selectbox('주파수:', frequency_options)

    run_button = st.form_submit_button('Run')

    if run_button:
        # 입력값 업데이트
        st.session_state.date_input = date_input
        st.session_state.time_input = time_input
        st.session_state.selected_model = selected_model
        if frequency_input == '2100MHz':
            st.session_state.frequency_input = '2100'
        elif frequency_input == '2600MHz - 10':
            st.session_state.frequency_input = '2600_10'
        elif frequency_input == '2600MHz - 20':
            st.session_state.frequency_input = '2600_20'

        # 분석 수행 및 상태 저장
        progress_bar = st.sidebar.progress(0)

        timestamp_input = datetime.combine(date_input, time_input)
        st.session_state.recommended_cell_states = load_and_predict(
            selected_model,
            df_map,
            timestamp_input,
            frequencies_to_check,
            progress_callback=progress_bar.progress
        )

        progress_bar.progress(100)
        st.sidebar.success('분석 완료')

        progress_bar.empty()

        current_recommendations = st.session_state.recommended_cell_states[st.session_state.frequency_input] if st.session_state.recommended_cell_states else {}

        # 지도를 생성하고 세션 상태에 저장
        m = map_recommend.create_map(df_map, current_recommendations)
        st.session_state.map_data = m

# 지도 출력 (Run 버튼을 눌렀을 때만 재렌더링)
if st.session_state.map_data:
    folium_data = st_folium(st.session_state.map_data)

    st.write("""
    <style>
    iframe {
        max-width: 1200px;
        width: 100%;
        height: 600px;
        border-radius: 10px;
        border: 1px solid #d3d3d3;
    }
    </style>
    """, unsafe_allow_html=True)

    # Cell ID 선택 및 상태 표시
    matching_enbid_pci = '33011_221'
    unique_cells = df['enbid_pci'].unique().tolist()
    matching_enbid_pci = st.selectbox('셀 ID:', unique_cells, index=unique_cells.index(matching_enbid_pci) if matching_enbid_pci in unique_cells else 0)

    st.markdown(f"##### **Cell**: {matching_enbid_pci}")

    for freq in frequencies_to_check:
        if st.session_state.recommended_cell_states is None:
            status = 'Unknown'
        else:
            if matching_enbid_pci not in st.session_state.recommended_cell_states[freq] or st.session_state.recommended_cell_states[freq][matching_enbid_pci] is None:
                status = 'Unknown'
            else:
                status = st.session_state.recommended_cell_states[freq][matching_enbid_pci]
        if status == 'OFF':
            color = '#D32F2F'
        elif status == 'ON':
            color = '#1976D2'
        else:
            color = '#aaa'
        if freq == '2100':
            freq_check = '2100MHz'
        elif freq == '2600_10':
            freq_check = '2600MHz - 10'
        elif freq == '2600_20':
            freq_check = '2600MHz - 20'

        st.markdown(f"""
        <div style="background-color: #f0f4f8; padding: 20px; border-radius: 10px; border: 1px solid #ccc; margin-bottom: 10px;">
            <strong>{freq_check}</strong>: <span style="color: {color};">{status}</span>
        </div>
        """, unsafe_allow_html=True)