pages/5_cycle_comparison.py

"""
Feature 5: LLM-Driven Cycle Comparison
Side-by-side comparison of two testing cycles with AI narrative.
Supports cross-month comparison — cycles loaded directly from cycle_periods DB table.
"""
import streamlit as st
import pandas as pd
from core.db_connector import get_db_connector
from core import config
from analysis.basic_analysis import DataAnalyzer
from analysis.llm_analyzer import LLMCycleAnalyzer
from ui.components import page_header, cycle_selector, metric_row, analysis_card, follow_up_section
from ui.plotly_charts import create_comparison_chart
from core.timezone import convert_df_timestamps_to_eastern

page_header(
    "Cycle Comparison",
    "Compare two testing cycles side-by-side with AI-powered analysis of differences."
)

db = get_db_connector()
analyzer = DataAnalyzer()
llm = LLMCycleAnalyzer()


# ── Load all cycles directly from DB ────────────────────────────

@st.cache_data(ttl=600)
def load_all_cycles():
    return db.get_cycle_periods()


@st.cache_data(ttl=3600)
def get_available_months():
    return db.get_available_months()


@st.cache_data(ttl=3600)
def get_cycle_months():
    return db.get_cycle_months()


all_cycles = load_all_cycles()
sensor_months = get_available_months()
cycle_months = get_cycle_months()
available_months = sorted(set(sensor_months) | set(cycle_months))

if len(all_cycles) < 2:
    st.warning("Need at least 2 cycles in the database for comparison.")
    st.stop()


# ── Cycle Selection with Independent Month Filters ──────────────

@st.cache_data(ttl=600)
def load_cycle_data(start_time, end_time):
    """Load and analyze a cycle's data"""
    detail_tags = [t for t in config.CYCLE_DETAIL_TAGS if db.get_tag_index(t) is not None]
    df_pivot = db.get_pivot_data(detail_tags, start_time, end_time)
    if df_pivot.empty:
        return None, None
    # Convert timestamps to Eastern for display
    df_pivot = convert_df_timestamps_to_eastern(df_pivot)
    stats = analyzer.analyze_cycle(df_pivot)
    return df_pivot, stats


month_options = ["All Months"] + available_months

col_a, col_b = st.columns(2)

with col_a:
    st.subheader("Cycle A")
    month_a = st.selectbox("Filter by month", month_options, key="comp_month_a")
    if month_a == "All Months":
        filtered_a = all_cycles
    else:
        filtered_a = [c for c in all_cycles if (c.get('start_time_et') or c['start_time']).strftime('%Y-%m') == month_a]
    cycle_a = cycle_selector(filtered_a, key_prefix="comp_a")

with col_b:
    st.subheader("Cycle B")
    month_b = st.selectbox("Filter by month", month_options, key="comp_month_b")
    if month_b == "All Months":
        filtered_b = all_cycles
    else:
        filtered_b = [c for c in all_cycles if (c.get('start_time_et') or c['start_time']).strftime('%Y-%m') == month_b]
    cycle_b = cycle_selector(filtered_b, key_prefix="comp_b")

if cycle_a is None or cycle_b is None:
    st.stop()

if cycle_a['cycle_id'] == cycle_b['cycle_id']:
    st.warning("Please select two different cycles to compare.")
    st.stop()

# Load data for both cycles
with st.spinner("Loading cycle data..."):
    df_a, stats_a = load_cycle_data(cycle_a['start_time'], cycle_a['end_time'])
    df_b, stats_b = load_cycle_data(cycle_b['start_time'], cycle_b['end_time'])

if df_a is None or df_b is None:
    st.error("Could not load data for one or both cycles.")
    st.stop()

st.divider()

# Side-by-side metrics
st.subheader("Key Metrics Comparison")

metrics_table = []


def add_metric(label, key_path_a, key_path_b, fmt=".1f", unit=""):
    """Helper to extract nested dict values for comparison"""
    val_a = stats_a
    val_b = stats_b
    for key in key_path_a:
        val_a = val_a.get(key, {}) if isinstance(val_a, dict) else None
    for key in key_path_b:
        val_b = val_b.get(key, {}) if isinstance(val_b, dict) else None

    # Ensure final values are numeric, not residual dicts from missing keys
    if not isinstance(val_a, (int, float)):
        val_a = None
    if not isinstance(val_b, (int, float)):
        val_b = None

    if val_a is not None and val_b is not None:
        diff = val_a - val_b
        metrics_table.append({
            'Metric': label,
            f'Cycle {cycle_a["cycle_id"]}': f"{val_a:{fmt}} {unit}".strip(),
            f'Cycle {cycle_b["cycle_id"]}': f"{val_b:{fmt}} {unit}".strip(),
            'Difference': f"{diff:+{fmt}} {unit}".strip(),
        })


add_metric("Duration (min)", ['time_range', 'duration_minutes'], ['time_range', 'duration_minutes'], ".0f", "min")
add_metric("Peak Pressure", ['discharge_pressure', 'peak'], ['discharge_pressure', 'peak'], ".0f", "bar")
add_metric("Avg Pressure", ['discharge_pressure', 'avg'], ['discharge_pressure', 'avg'], ".1f", "bar")
add_metric("Compression Ratio (avg)", ['compression_ratio', 'avg'], ['compression_ratio', 'avg'], ".1f", "x")
add_metric("Avg Flow", ['flow', 'avg_flow'], ['flow', 'avg_flow'], ".2f", "kg/min")
add_metric("Peak TT110", ['temperature_TT110', 'peak'], ['temperature_TT110', 'peak'], ".1f", "K")
add_metric("Peak TT130", ['temperature_TT130', 'peak'], ['temperature_TT130', 'peak'], ".1f", "K")

if 'motor' in stats_a and 'motor' in stats_b:
    add_metric("Peak Motor Speed", ['motor', 'peak_speed'], ['motor', 'peak_speed'], ".0f", "RPM")

if 'ramp_rate' in stats_a and 'ramp_rate' in stats_b:
    add_metric("Max Ramp Rate", ['ramp_rate', 'max'], ['ramp_rate', 'max'], ".1f", "bar/min")

if metrics_table:
    st.dataframe(pd.DataFrame(metrics_table), use_container_width=True, hide_index=True)

# Cycle metadata from DB (new fields)
st.divider()
st.subheader("Cycle Metadata (from DB)")
meta_cols = st.columns(2)

for meta_col, cycle, label in [(meta_cols[0], cycle_a, "A"), (meta_cols[1], cycle_b, "B")]:
    with meta_col:
        _et = (cycle.get('start_time_et') or cycle['start_time']).strftime('%b %d %H:%M')
        st.markdown(f"**Cycle {cycle['cycle_id']}** ({_et} ET)")
        m_cols = st.columns(2)
        m_cols[0].metric("Dispensed", f"{cycle.get('total_kg_dispensed', 0):.2f} kg")
        m_cols[1].metric("Pump Strokes", f"{cycle.get('total_pump_strokes', 0):.0f}")

st.divider()

# Comparison charts
st.subheader("Overlaid Comparison Charts")

# Find common tags
common_tags = [t for t in df_a.columns if t in df_b.columns and t != 'timestamp']
chart_tags = [t for t in ['PT130', 'TT110', 'TT130', 'FT140', 'M130_Speed', 'MC130_VFD_Speed'] if t in common_tags]

if chart_tags:
    selected_chart_tag = st.selectbox(
        "Select sensor to compare",
        chart_tags,
        key="comparison_tag",
    )

    _et_a = (cycle_a.get('start_time_et') or cycle_a['start_time']).strftime('%b %d %H:%M')
    _et_b = (cycle_b.get('start_time_et') or cycle_b['start_time']).strftime('%b %d %H:%M')
    label_a = f"Cycle {cycle_a['cycle_id']} ({_et_a} ET)"
    label_b = f"Cycle {cycle_b['cycle_id']} ({_et_b} ET)"

    fig = create_comparison_chart(df_a, df_b, selected_chart_tag, label_a, label_b)
    st.plotly_chart(fig, use_container_width=True)

st.divider()

# LLM Comparison
st.subheader("AI Comparison Analysis")

if not llm.api_available:
    st.error("Claude API key not configured. Add ANTHROPIC_API_KEY to your .env file.")
    st.stop()

cache_key = f"comparison_{cycle_a['cycle_id']}_{cycle_b['cycle_id']}"
prompt_key = f"{cache_key}_prompt"
followup_key = f"{cache_key}_followups"
cached_result = st.session_state.get(cache_key)

if cached_result:
    analysis_card("AI Comparison", cached_result)

    # Follow-up section
    original_prompt = st.session_state.get(prompt_key, "")
    follow_up_section(
        session_key=followup_key,
        llm_analyzer=llm,
        original_prompt=original_prompt,
        original_analysis=cached_result,
    )

    if st.button("Re-analyze", key="reanalyze_comparison"):
        del st.session_state[cache_key]
        st.session_state.pop(prompt_key, None)
        st.session_state.pop(followup_key, None)
        st.rerun()
else:
    if st.button("Compare with AI", type="primary", use_container_width=True):
        label_a = f"Cycle {cycle_a['cycle_id']} ({(cycle_a.get('start_time_et') or cycle_a['start_time']).strftime('%b %d')} ET)"
        label_b = f"Cycle {cycle_b['cycle_id']} ({(cycle_b.get('start_time_et') or cycle_b['start_time']).strftime('%b %d')} ET)"

        # Build and store the prompt for follow-up context
        knowledge_text = ""
        if llm.retriever:
            try:
                context = llm.retriever.get_context_for_comparison(stats_a, stats_b)
                knowledge_text = llm.retriever.format_for_prompt(context)
            except Exception:
                pass
        st.session_state[prompt_key] = llm._build_comparison_prompt(
            stats_a, stats_b, label_a, label_b, knowledge_text
        )

        with st.spinner("Claude is comparing the two cycles..."):
            result = llm.compare_cycles(stats_a, stats_b, label_a, label_b)

        st.session_state[cache_key] = result
        st.rerun()