import numpy as np
import matplotlib.pyplot as plt
import streamlit as st

# Define the calculation function
def calculate_relay_curve(tr_kva, tr_volt, ct_pr, ct_sec, fac, tms, i_f_fac, curve):
    # Calculate transformer current
    tr_curr = round(tr_kva / (tr_volt * np.sqrt(3)), 3)
    pick = round(fac * tr_curr / ct_pr, 3)  # Pickup value in multiples of In
    i_f = round(i_f_fac * tr_curr * fac, 3)  # Fault current in amps

    # Select curve coefficients based on relay type
    if curve == 'IEC Normal Inverse':
        b_ = 0.14
        a_ = 0.02
    elif curve == 'IEC Very Inverse':
        b_ = 13.5
        a_ = 1
    elif curve == 'IEC Extremely Inverse':
        b_ = 80
        a_ = 2.0
    elif curve == 'IEC Long Time Inverse':
        b_ = 120
        a_ = 1.0
    else:
        st.error("Invalid Curve Selected.")
        return

    # Calculate fault trip time
    try:
        fault_time = b_ * tms / (((i_f * (ct_sec / ct_pr) / (pick * ct_sec)) ** a_) - 1)
    except ZeroDivisionError:
        st.error("Error: Division by zero in fault trip time calculation. Check input values.")
        return

    # Generate data for the relay curve (current vs time)
    x_curr = np.linspace(pick * ct_pr * 1.1, pick * ct_pr * 5, 500)
    y_time = []
    for x in x_curr:
        try:
            time = b_ * tms / (((x * (ct_sec / ct_pr) / (pick * ct_sec)) ** a_) - 1)
            y_time.append(time if time > 0 else np.nan)  # Avoid negative/undefined times
        except ZeroDivisionError:
            y_time.append(np.nan)

    # Display results
    st.write('### Results:')
    st.write(f'**Rated Current:** {tr_curr} A')
    st.write(f'**Pickup:** {pick} x In')
    st.write(f'**Fault Current:** {i_f} A')
    st.write(f'**Trip Time:** {round(fault_time, 3)} sec')

    # Plot the relay curve
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(x_curr, y_time, label="Relay Curve", color='blue')

    # Mark the fault current and trip time on the graph
    ax.axvline(x=i_f, color='red', linestyle='--', label=f"Fault Current: {i_f} A")
    ax.axhline(y=fault_time, color='green', linestyle='--', label=f"Trip Time: {round(fault_time, 3)} sec")

    # Adding labels and title
    ax.set_xlabel('Current (Amp)')
    ax.set_ylabel('Time (sec)')
    ax.set_title(f"Relay Curve: {curve}")

    # Adding a legend
    ax.legend()

    # Show the plot
    st.pyplot(fig)

# Streamlit app layout
st.title("Relay Curve Calculator")
st.sidebar.header("Input Parameters")

# Input fields
tr_kva = st.sidebar.number_input("Transformer Rating (kVA)", min_value=1.0, value=1000.0, step=1.0)
tr_volt = st.sidebar.number_input("Transformer Voltage (kV)", min_value=1.0, value=11000.0, step=1.0)
ct_pr = st.sidebar.number_input("CT Primary (Amp)", min_value=1.0, value=200.0, step=1.0)
ct_sec = st.sidebar.number_input("CT Secondary (Amp)", min_value=1.0, value=5.0, step=1.0)
fac = st.sidebar.number_input("Fault Anticipation Factor", min_value=1.0, value=1.2, step=0.1)
tms = st.sidebar.number_input("Time Multiplier Setting (TMS)", min_value=0.01, value=0.1, step=0.01)
i_f_fac = st.sidebar.number_input("Fault Current Multiplier", min_value=1.0, value=10.0, step=0.1)
curve = st.sidebar.selectbox("Relay Curve Type", ['IEC Normal Inverse', 'IEC Very Inverse', 'IEC Extremely Inverse', 'IEC Long Time Inverse'])

# Run the calculation when the button is clicked
if st.sidebar.button("Calculate"):
    calculate_relay_curve(tr_kva, tr_volt, ct_pr, ct_sec, fac, tms, i_f_fac, curve)