import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import pandas as pd

# Set page config
st.set_page_config(
    page_title="Signal Generator Calculator",
    page_icon="📊",
    layout="wide"
)

# Custom CSS for styling
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 30px;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #ff7f0e;
        margin-top: 20px;
        margin-bottom: 10px;
    }
    .info-box {
        background-color: #f0f2f6;
        padding: 20px;
        border-radius: 10px;
        margin: 10px 0;
    }
    .warning-box {
        background-color: #ffe6e6;
        padding: 15px;
        border-radius: 5px;
        margin: 10px 0;
        border-left: 5px solid #ff4444;
    }
</style>
""", unsafe_allow_html=True)

st.markdown('<h1 class="main-header">🔧 AD9833 Signal Generator Design Calculator</h1>', unsafe_allow_html=True)
st.markdown("*For Reddit user Schniedelholz*")

# Sidebar for inputs
with st.sidebar:
    st.header("⚙️ Input Parameters")

    clock_freq = st.number_input(
        "Clock Frequency (MHz)",
        min_value=1.0,
        max_value=50.0,
        value=24.0,
        step=0.1
    )

    max_output_freq = st.number_input(
        "Maximum Output Frequency (MHz)",
        min_value=0.01,
        max_value=clock_freq/2,
        value=1.0,
        step=0.1
    )

    st.markdown("### Filter Design")
    filter_type = st.selectbox(
        "Filter Type",
        ["Butterworth", "Chebyshev", "Elliptic", "Bessel"]
    )

    filter_order = st.slider("Filter Order", 3, 9, 5, 2)

    st.markdown("### Output Stage")
    supply_voltage = st.number_input("Op-Amp Supply (±V)", 5.0, 20.0, 15.0, 0.5)
    output_impedance = st.selectbox("Output Impedance", ["50Ω", "75Ω", "High-Z"])

# Main content area
col1, col2 = st.columns(2)

with col1:
    st.markdown('<h2 class="sub-header">📐 Calculated Parameters</h2>', unsafe_allow_html=True)

    # Calculate key parameters
    nyquist_freq = clock_freq / 2
    recommended_cutoff = max_output_freq * 2.5

    if recommended_cutoff > nyquist_freq * 0.8:
        recommended_cutoff = nyquist_freq * 0.8

    # AD9833 specific calculations
    freq_resolution = clock_freq * 1e6 / (2**28)
    max_theoretical_freq = clock_freq / 2

    # Display results in styled boxes
    st.markdown('<div class="info-box">', unsafe_allow_html=True)
    st.write(f"**Nyquist Frequency:** {nyquist_freq:.2f} MHz")
    st.write(f"**Recommended Filter Cutoff:** {recommended_cutoff:.2f} MHz")
    st.write(f"**Frequency Resolution:** {freq_resolution:.3f} Hz")
    st.write(f"**Phase Resolution:** {360/4096:.3f}°")
    st.markdown('</div>', unsafe_allow_html=True)

    # Filter component values (example for Butterworth)
    if filter_type == "Butterworth":
        st.markdown('<div class="info-box">', unsafe_allow_html=True)
        st.write("**Example Component Values (1kΩ impedance):**")
        # Simplified calculation for demonstration
        C = 1 / (2 * np.pi * recommended_cutoff * 1e6 * 1000)
        st.write(f"Capacitor: {C*1e12:.0f} pF")
        st.write(f"Resistor: 1 kΩ")
        st.markdown('</div>', unsafe_allow_html=True)

    # Warnings
    if max_output_freq > nyquist_freq * 0.4:
        st.markdown('<div class="warning-box">', unsafe_allow_html=True)
        st.warning("⚠️ Output frequency is close to Nyquist limit. Consider lower frequencies for better signal quality.")
        st.markdown('</div>', unsafe_allow_html=True)

with col2:
    st.markdown('<h2 class="sub-header">📊 Filter Response</h2>', unsafe_allow_html=True)

    # Create filter
    if filter_type == "Butterworth":
        b, a = signal.butter(filter_order, 2*recommended_cutoff/clock_freq)
    elif filter_type == "Chebyshev":
        b, a = signal.cheby1(filter_order, 0.5, 2*recommended_cutoff/clock_freq)
    elif filter_type == "Elliptic":
        b, a = signal.ellip(filter_order, 0.5, 40, 2*recommended_cutoff/clock_freq)
    else:  # Bessel
        b, a = signal.bessel(filter_order, 2*recommended_cutoff/clock_freq)

    # Frequency response
    w, h = signal.freqz(b, a, worN=2000)
    freq = w * clock_freq / (2 * np.pi)

    # Plot
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6))

    # Magnitude
    ax1.plot(freq, 20 * np.log10(abs(h)), 'b', linewidth=2)
    ax1.axvline(max_output_freq, color='g', linestyle='--', label='Max Output')
    ax1.axvline(recommended_cutoff, color='r', linestyle='--', label='Filter Cutoff')
    ax1.set_ylabel('Magnitude (dB)')
    ax1.set_xlim(0, nyquist_freq)
    ax1.grid(True, alpha=0.3)
    ax1.legend()
    ax1.set_title(f'{filter_type} Filter Response')

    # Phase
    ax2.plot(freq, np.angle(h) * 180/np.pi, 'r', linewidth=2)
    ax2.set_xlabel('Frequency (MHz)')
    ax2.set_ylabel('Phase (degrees)')
    ax2.set_xlim(0, nyquist_freq)
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    st.pyplot(fig)

# Additional calculations section
st.markdown('<h2 class="sub-header">🔢 Additional Calculations</h2>', unsafe_allow_html=True)

col3, col4 = st.columns(2)

with col3:
    st.markdown("### Digital Potentiometer Settings")
    desired_gain = st.slider("Desired Voltage Gain", 0.1, 10.0, 1.0, 0.1)

    # Calculate pot settings (assuming 10kΩ digital pot)
    pot_steps = 256  # 8-bit pot
    pot_max_r = 10000  # 10kΩ

    # For non-inverting amp: Gain = 1 + (Rf/Rin)
    if desired_gain >= 1:
        rf_ratio = desired_gain - 1
        pot_setting = int(rf_ratio * pot_steps / 10)  # Assuming max gain of 10
        pot_resistance = pot_setting * pot_max_r / pot_steps

        st.markdown('<div class="info-box">', unsafe_allow_html=True)
        st.write(f"**Digital Pot Setting:** {pot_setting}/256")
        st.write(f"**Resistance:** {pot_resistance:.0f} Ω")
        st.write(f"**Actual Gain:** {1 + pot_resistance/1000:.2f}")
        st.markdown('</div>', unsafe_allow_html=True)

with col4:
    st.markdown("### Output Protection")
    max_current = st.slider("Maximum Output Current (mA)", 10, 100, 25, 5)

    # Protection resistor calculation
    short_circuit_r = (supply_voltage / (max_current / 1000))
    power_rating = (max_current / 1000)**2 * short_circuit_r

    st.markdown('<div class="info-box">', unsafe_allow_html=True)
    st.write(f"**Series Protection Resistor:** {short_circuit_r:.0f} Ω")
    st.write(f"**Minimum Power Rating:** {power_rating:.2f} W")
    st.write(f"**Suggested Resistor:** {short_circuit_r*1.2:.0f} Ω, {power_rating*2:.1f} W")
    st.markdown('</div>', unsafe_allow_html=True)

# Summary table
st.markdown('<h2 class="sub-header">📋 Design Summary</h2>', unsafe_allow_html=True)

summary_data = {
    "Parameter": [
        "Clock Frequency",
        "Max Output Frequency",
        "Filter Type",
        "Filter Order",
        "Filter Cutoff",
        "Frequency Resolution",
        "Output Impedance",
        "Supply Voltage"
    ],
    "Value": [
        f"{clock_freq} MHz",
        f"{max_output_freq} MHz",
        filter_type,
        str(filter_order),  # Convert to string
        f"{recommended_cutoff:.2f} MHz",
        f"{freq_resolution:.3f} Hz",
        output_impedance,
        f"±{supply_voltage} V"
    ]
}

df = pd.DataFrame(summary_data)
st.table(df)

# Footer
st.markdown("---")
st.markdown("*Happy building! Remember to use proper decoupling capacitors and a good PCB layout.*")