import streamlit as st
import numpy as np
from eda_ts.constants import DF_KEY
import plotly.graph_objects as go
import pandas as pd
st.set_page_config(
    page_title="Frequencies",
    page_icon="🔄",
    layout="wide"
)
st.title("🔄 Frequencies")


df = st.session_state[DF_KEY]

def format_period_info(hz):
    minute = 60
    hour = 3600
    day = 86400
    return f"""
    - {hz:.5f} cycles/second = {1/hz:.5f} seconds/cycle
    - {hz * minute:.5f} cycles/minute = {1/hz/minute:.5f} minutes/cycle
    - {hz * hour:.5f} cycles/hour = {1/hz/hour:.5f} hours/cycle
    - {hz * day:.5f} cycles/day = {1/hz/day:.5f} days/cycle
    """

inferred_time_interval = pd.Timedelta(df.index.freq).total_seconds() 
if inferred_time_interval is not None:
    st.info(f"🇫 Detected sampling interval: **{inferred_time_interval:.5f} seconds** (_{1/inferred_time_interval:.2f} Hz_)")
else:
    st.warning("⚠️ Could not automatically detect sampling interval")


col1, col2 = st.columns(2)
with col1:
    sampling_rate = st.number_input(
        "Enter Sampling Rate:", 
        min_value=1e-5,
        value=inferred_time_interval if inferred_time_interval is not None else 1.0,
    )
with col2:
    sampling_unit = st.selectbox(
        label="Select the sampling unit:", 
        options=["Hz", "Seconds"],
        index=1
        
    )
if sampling_unit == "Seconds":
    sampling_rate = 1 / sampling_rate
    sampling_unit = "Hz"

st.markdown("---")
col_check_box, col_filters, col_transforms = st.columns(3)

with col_check_box:
    show_dc = st.checkbox(
        "Show DC Component", 
        value=False, 
        help="Show the DC component($f=0$) in the FFT plot"
    )
    real_signal = st.checkbox(
        "Real-valued Signal", 
        value=True, 
        help="Will use `rfft` instead of `fft`"
    )
    colx,coly = st.columns(2)
    with colx:
        log_x = st.checkbox(
            "Log X", 
            value=True, 
            help="Logarithmic scale for the x-axis"
        )
    with coly:
        log_y = st.checkbox(
            "Log Y", 
            value=True, 
            help="Logarithmic scale for the y-axis"
            )
        
with col_filters:
    low_pass = st.number_input(
        "LPF (Hz)", 
        min_value=0.0,
        value=0.01,
        help="Low Pass Filter - cuts off frequencies above the specified value. Set to 0 to show all frequencies"
    )
    high_pass = st.number_input(
        "HPF (Hz)", 
        min_value=0.0,
        value=0.0,
        help="High Pass Filter - cuts off frequencies below the specified value. Set to 0 to show all frequencies"
    )    

with col_transforms:
    win_size = st.number_input(
        "Resample Window Size",
        min_value=0,
        value=0,
        help="Choose window size for rolling mean filter before applying FFT",
        format="%d"
    )
    
    if win_size > 0:
        df = df.rolling(window=int(win_size)).mean()


if real_signal:
    fft_data = np.fft.rfft(df.values, axis=0, norm='ortho')
    fftfreq_data = np.fft.rfftfreq(df.shape[0], d=sampling_rate)
else:
    fft_data = np.fft.fft(df.values, axis=0, norm='ortho')
    fftfreq_data = np.fft.fftfreq(df.shape[0], d=sampling_rate)

fft_angle = np.angle(fft_data)
fft_magnitude = np.abs(fft_data)


if not show_dc:
    fft_magnitude = fft_magnitude[1:]
    fftfreq_data = fftfreq_data[1:]


if low_pass > 0:
    freq_mask = fftfreq_data <= low_pass
    fft_magnitude = fft_magnitude[freq_mask]
    fftfreq_data = fftfreq_data[freq_mask]


if high_pass > 0:
    freq_mask = fftfreq_data >= high_pass
    fft_magnitude = fft_magnitude[freq_mask]
    fftfreq_data = fftfreq_data[freq_mask]


fig = go.Figure()
for i, column in enumerate(df.columns):
    fig.add_trace(go.Scatter(
        x=fftfreq_data,
        y=fft_magnitude[:, i],
        mode='lines',
        name=column
    ))

fig.update_layout(
    title="FFT Magnitude",
    xaxis_title="Frequency (Hz)",
    yaxis_title="Magnitude",
    xaxis_type="log" if log_x else "linear",
    yaxis_type="log" if log_y else "linear",
)

st.plotly_chart(fig, use_container_width=True)


with st.expander("⏱️ Frequency/Period Converter", expanded=True):
    if 'freq_value' not in st.session_state:
        st.session_state.freq_value = 1e-5

    def update_freq_from_input():
        st.session_state.freq_value = st.session_state.freq_input

    def update_freq_from_slider():
        st.session_state.freq_value = st.session_state.freq_slider

    freq_cols = st.columns([0.8, 0.2])
    with freq_cols[0]:
        freq_to_convert = st.slider(
            "Select frequency (Hz)", 
            min_value=0.0,
            max_value=0.1, 
            value=st.session_state.freq_value,
            step=1e-5,
            format="%.5f",
            key="freq_slider",
            on_change=update_freq_from_slider
        )
    with freq_cols[1]:
        freq_to_convert = st.number_input(
            "Or type frequency",
            min_value=1e-5,
            max_value=1.0,
            value=st.session_state.freq_value,
            format="%.5f",
            key="freq_input",
            on_change=update_freq_from_input
        )
    st.markdown(format_period_info(st.session_state.freq_value))
