pages/1_data_explorer.py

"""
Feature 1: Data Explorer - Pivot data view with interactive charts
Now with Natural Language query support powered by Claude
"""
import streamlit as st
import pandas as pd
from core.db_connector import get_db_connector
from core.cached_queries import cached_get_pivot_data
from analysis.nl2sql import NLQueryParser
from ui.components import page_header, tag_multiselect, date_range_picker, no_data_message
from ui.plotly_charts import create_timeseries_chart
from core.timezone import to_utc, convert_df_timestamps_to_eastern

page_header(
    "Data Explorer",
    "View sensor data in wide format with tags as columns. Use natural language or the sidebar controls to query."
)

db = get_db_connector()
nl_parser = NLQueryParser()

# ── Natural Language Query Section ──
st.subheader("Ask in Natural Language")

nl_col1, nl_col2 = st.columns([5, 1])
with nl_col1:
    nl_query = st.text_input(
        "Describe what data you want",
        placeholder='e.g. "show me pressure and temperature from Sept 25 between 10am and 2pm"',
        key="nl_query_input",
        label_visibility="collapsed",
    )
with nl_col2:
    nl_btn = st.button("Query", type="primary", use_container_width=True, key="nl_query_btn")

# Process NL query
nl_parsed = None
if nl_btn and nl_query:
    if not nl_parser.api_available:
        st.error("Claude API key not configured. Add ANTHROPIC_API_KEY to your .env file for natural language queries.")
    else:
        with st.spinner("Interpreting your query..."):
            all_sensors = db.get_all_tag_names()
            nl_parsed = nl_parser.parse(nl_query, all_sensors)

        if nl_parsed and 'error' in nl_parsed:
            st.error(f"Could not parse query: {nl_parsed['error']}")
            nl_parsed = None
        elif nl_parsed:
            # Store parsed result in session state for data fetch
            st.session_state['nl_parsed'] = nl_parsed
        else:
            st.error("Could not understand the query. Try being more specific about sensors and time range.")

# Check for stored NL result (persists across reruns)
if 'nl_parsed' in st.session_state and st.session_state['nl_parsed']:
    nl_parsed = st.session_state['nl_parsed']

# Show parsed interpretation
if nl_parsed and 'error' not in nl_parsed:
    with st.container():
        st.success(f"**Parsed:** {nl_parsed.get('explanation', '')}")
        pcol1, pcol2, pcol3 = st.columns(3)
        pcol1.markdown(f"**Sensors:** {', '.join(nl_parsed['sensors'])}")
        pcol2.markdown(f"**From:** {nl_parsed['start_time'].strftime('%b %d, %Y %H:%M')}")
        pcol3.markdown(f"**To:** {nl_parsed['end_time'].strftime('%b %d, %Y %H:%M')}")

    col_fetch, col_clear = st.columns([1, 1])
    with col_fetch:
        nl_fetch_btn = st.button("Fetch Data", type="primary", use_container_width=True, key="nl_fetch_btn")
    with col_clear:
        nl_clear_btn = st.button("Clear", use_container_width=True, key="nl_clear_btn")

    if nl_clear_btn:
        del st.session_state['nl_parsed']
        st.rerun()

    if nl_fetch_btn:
        selected_tags = nl_parsed['sensors']
        # NL parser returns naive datetimes — treat as Eastern, convert to UTC for queries
        start_dt = to_utc(nl_parsed['start_time'])
        end_dt = to_utc(nl_parsed['end_time'])

        duration_hours = (end_dt - start_dt).total_seconds() / 3600
        table_used = db._select_table(start_dt, end_dt)
        table_label = {
            'procdatafloattable': 'Raw (10Hz)',
            'procdatafloattable_utc_1sec': '1-second aggregates',
            'procdatafloattable_utc_15sec': '15-second aggregates',
        }.get(table_used, table_used)

        with st.spinner(f"Querying {len(selected_tags)} sensors ({duration_hours:.1f} hours) from {table_label}..."):
            df_pivot = cached_get_pivot_data(tuple(selected_tags), start_dt, end_dt)

        if df_pivot.empty:
            no_data_message()
        else:
            # Convert timestamps to Eastern for display
            df_pivot = convert_df_timestamps_to_eastern(df_pivot)

            # Summary metrics
            cols = st.columns(4)
            cols[0].metric("Rows", f"{len(df_pivot):,}")
            cols[1].metric("Sensors", len(selected_tags))
            cols[2].metric("Duration", f"{duration_hours:.1f} hrs")
            cols[3].metric("Table", table_label)

            st.divider()

            # Interactive chart
            tags_in_data = [t for t in selected_tags if t in df_pivot.columns]
            if tags_in_data:
                chart_title = f"{', '.join(tags_in_data[:3])}{'...' if len(tags_in_data) > 3 else ''}"
                fig = create_timeseries_chart(df_pivot, tags_in_data, title=chart_title)
                st.plotly_chart(fig, use_container_width=True)

            st.divider()

            # Summary statistics
            with st.expander("Summary Statistics", expanded=True):
                stats_data = []
                for tag in tags_in_data:
                    col_data = df_pivot[tag].dropna()
                    if len(col_data) > 0:
                        stats_data.append({
                            'Sensor': tag,
                            'Count': len(col_data),
                            'Min': f"{col_data.min():.4f}",
                            'Max': f"{col_data.max():.4f}",
                            'Mean': f"{col_data.mean():.4f}",
                            'Std': f"{col_data.std():.4f}",
                        })
                if stats_data:
                    st.dataframe(pd.DataFrame(stats_data), use_container_width=True, hide_index=True)

            # Data table
            with st.expander("Data Table", expanded=False):
                st.dataframe(df_pivot, use_container_width=True, height=400)

            # CSV download
            csv = df_pivot.to_csv(index=False)
            st.download_button(
                label="Download CSV",
                data=csv,
                file_name=f"csh2_data_{start_dt.strftime('%Y%m%d_%H%M')}_{end_dt.strftime('%Y%m%d_%H%M')}.csv",
                mime="text/csv",
            )

    st.divider()

st.caption("Or use the sidebar controls for manual selection:")

# ── Sidebar controls (manual mode) ──
with st.sidebar:
    st.subheader("Manual Query")
    selected_tags = tag_multiselect(db, key="explorer")
    st.divider()
    start_dt, end_dt = date_range_picker(key="explorer")

    query_btn = st.button("Fetch Data", type="primary", use_container_width=True, key="manual_fetch_btn")

# Main content (manual mode)
if query_btn and selected_tags:
    # Manual date picker returns naive datetimes — treat as Eastern, convert to UTC for queries
    start_dt_utc = to_utc(start_dt)
    end_dt_utc = to_utc(end_dt)

    if start_dt >= end_dt:
        st.error("Start time must be before end time.")
    else:
        duration_hours = (end_dt_utc - start_dt_utc).total_seconds() / 3600
        table_used = db._select_table(start_dt_utc, end_dt_utc)
        table_label = {
            'procdatafloattable': 'Raw (10Hz)',
            'procdatafloattable_utc_1sec': '1-second aggregates',
            'procdatafloattable_utc_15sec': '15-second aggregates',
        }.get(table_used, table_used)

        with st.spinner(f"Querying {len(selected_tags)} sensors ({duration_hours:.1f} hours) from {table_label}..."):
            df_pivot = cached_get_pivot_data(tuple(selected_tags), start_dt_utc, end_dt_utc)

        if df_pivot.empty:
            no_data_message()
        else:
            # Convert timestamps to Eastern for display
            df_pivot = convert_df_timestamps_to_eastern(df_pivot)
            # Summary metrics
            cols = st.columns(4)
            cols[0].metric("Rows", f"{len(df_pivot):,}")
            cols[1].metric("Sensors", len(selected_tags))
            cols[2].metric("Duration", f"{duration_hours:.1f} hrs")
            cols[3].metric("Table", table_label)

            st.divider()

            # Interactive chart
            tags_in_data = [t for t in selected_tags if t in df_pivot.columns]
            if tags_in_data:
                chart_title = f"{', '.join(tags_in_data[:3])}{'...' if len(tags_in_data) > 3 else ''}"
                fig = create_timeseries_chart(df_pivot, tags_in_data, title=chart_title)
                st.plotly_chart(fig, use_container_width=True)

            st.divider()

            # Summary statistics
            with st.expander("Summary Statistics", expanded=True):
                stats_data = []
                for tag in tags_in_data:
                    col_data = df_pivot[tag].dropna()
                    if len(col_data) > 0:
                        stats_data.append({
                            'Sensor': tag,
                            'Count': len(col_data),
                            'Min': f"{col_data.min():.4f}",
                            'Max': f"{col_data.max():.4f}",
                            'Mean': f"{col_data.mean():.4f}",
                            'Std': f"{col_data.std():.4f}",
                        })
                if stats_data:
                    st.dataframe(pd.DataFrame(stats_data), use_container_width=True, hide_index=True)

            # Data table
            with st.expander("Data Table", expanded=False):
                st.dataframe(df_pivot, use_container_width=True, height=400)

            # CSV download
            csv = df_pivot.to_csv(index=False)
            st.download_button(
                label="Download CSV",
                data=csv,
                file_name=f"csh2_data_{start_dt.strftime('%Y%m%d_%H%M')}_{end_dt.strftime('%Y%m%d_%H%M')}.csv",
                mime="text/csv",
            )

elif query_btn and not selected_tags:
    st.warning("Please select at least one sensor from the sidebar.")