import streamlit as st
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score

st.title("Compare Your Algorithm")

default_dataOne = {
    "X": [1, 2, 3, 4, 5],
    "Y": [2.2, 4.4, 6.5, 8.0, 10.1],
    "Select": [True, True, True, True, True] 
}
default_dataTwo = {
    "X": [1, 10, 100, 1000, 10000],
    "Y": [3.3, 6.6, 12.21, 24.84, 48.25],
    "Select": [True, True, True, True, True]
}
dataOne = pd.DataFrame(default_dataOne)
dataTwo = pd.DataFrame(default_dataTwo)

xlabel = st.text_input("X-axis", "X")
ylabel = st.text_input("Y-axis", "Y")

col1, col2 = st.columns(2)

with col1:
    st.subheader("Enter Your Data One")
    cola, colb = st.columns(2)
    with cola:
        user_dataOne = st.data_editor(dataOne, num_rows="dynamic", key="data_editor_one")
    with colb:
        fit_typeOne = st.radio(
            "Choose the Type of Fit",
            options=["Logarithmic", "Linear", "Linearithmic", "Quadratic", "Cubic", "Exponential"],
            index=1,
            key="one"
        )

with col2:
    st.subheader("Enter Your Data Two")
    colc, cold = st.columns(2)
    with colc:
        user_dataTwo = st.data_editor(dataTwo, num_rows="dynamic", key="data_editor_two")
    with cold:
        fit_typeTwo = st.radio(
            "Choose the Type of Fit",
            options=["Logarithmic", "Linear", "Linearithmic", "Quadratic", "Cubic", "Exponential"],
            index=0,
            key="two"
        )

try:
    selected_dataOne = user_dataOne[user_dataOne["Select"]]
    x = np.array(selected_dataOne["X"], dtype=float)
    y = np.array(selected_dataOne["Y"], dtype=float)

    if len(x) < 2 and len(y) < 2:
        st.warning("Please enter at least 2 data points.")
        st.stop()
except ValueError:
    st.error("Invalid data entered. Please ensure all values are numeric.")
    st.stop()

try:
    selected_dataTwo = user_dataTwo[user_dataTwo["Select"]]
    u = np.array(selected_dataTwo["X"], dtype=float)
    v = np.array(selected_dataTwo["Y"], dtype=float)

    if len(u) < 2 and len(v) < 2:
        st.warning("Please enter at least 2 data points.")
        st.stop()
except ValueError:
    st.error("Invalid data entered. Please ensure all values are numeric.")
    st.stop()

if fit_typeOne == "Logarithmic":
    try:
        log_x = np.log(x)
        coefficients = np.polyfit(log_x, y , 1)
        y_fit = coefficients[0] * log_x + coefficients[1]
        r2 = r2_score(y, y_fit)
        equation = f"y = {coefficients[0]:.4f}*log(x) + {coefficients[1]:.4f}"
    except ValueError:
        st.error("Logarithmic fit failed. Ensure all X values are positive.")
        st.stop()

elif fit_typeOne == "Linear":
    degree = 1
    coefficients = np.polyfit(x, y, degree)
    y_fit = np.polyval(coefficients, x)
    r2 = r2_score(y, y_fit)
    equation = f"y = {coefficients[0]:.4f}*x + {coefficients[1]:.4f}"

elif fit_typeOne == "Linearithmic":
    try:
        x_log_x = x * np.log(x)
        A = np.column_stack((x_log_x, x, np.ones_like(x)))
        coefficients, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
        a, b, c = coefficients
        y_fit = a * x_log_x + b * x + c
        r2 = r2_score(y, y_fit)
        equation = f"y = {a:.4f}*x*log(x) + {b:.4f}*x + {c:.4f}"
    except ValueError:
        st.error("Linearithmic fir failed. Ensure all X values are positive.")
        st.stop()

elif fit_typeOne == "Quadratic":
    degree = 2
    coefficients = np.polyfit(x, y, degree)
    y_fit = np.polyval(coefficients, x)
    r2 = r2_score(y, y_fit)
    equation = f"y = {coefficients[0]:.4f}*x² + {coefficients[1]:.4f}*x + {coefficients[2]:.4f}"

elif fit_typeOne == "Cubic":
    degree = 3
    coefficients = np.polyfit(x, y, degree)
    y_fit = np.polyval(coefficients, x)
    r2 = r2_score(y, y_fit)
    equation = f"y = {coefficients[0]:.4f}*x³ + {coefficients[1]:.4f}*x² + {coefficients[2]:.4f}*x + {coefficients[3]:.4f}"

elif fit_typeOne == "Exponential":
    try:
        log_y = np.log(y)
        coefficients = np.polyfit(x, log_y, 1)
        a = np.exp(coefficients[1])
        b = coefficients[0]
        y_fit = a * np.exp(b * x)
        r2 = r2_score(y, y_fit)
        equation = f"y = {a:.4f}*exp({b:.4f}*x)"
    except ValueError:
        st.error("Exponential fit failed. Ensure all Y values are positive.")
        st.stop()

if fit_typeTwo == "Logarithmic":
    try:
        log_u = np.log(u)
        coefficients_Two = np.polyfit(log_u, v , 1)
        v_fit = coefficients_Two[0] * log_u + coefficients_Two[1]
        r2_Two = r2_score(v, v_fit)
        equation_Two = f"y = {coefficients_Two[0]:.4f}*log(x) + {coefficients_Two[1]:.4f}"
    except ValueError:
        st.error("Logarithmic fit failed. Ensure all X values are positive.")
        st.stop()

elif fit_typeTwo == "Linear":
    degree_Two = 1
    coefficients_Two = np.polyfit(u, v, degree_Two)
    v_fit = np.polyval(coefficients_Two, u)
    r2_Two = r2_score(v, v_fit)
    equation_Two = f"y = {coefficients_Two[0]:.4f}*x + {coefficients_Two[1]:.4f}"

elif fit_typeTwo == "Linearithmic":
    try:
        u_log_u = u * np.log(u)
        B = np.column_stack((u_log_u, u, np.ones_like(u)))
        coefficients_Two, _, _, _ = np.linalg.lstsq(B, v, rcond=None)
        d, e, f = coefficients_Two
        v_fit = d * u_log_u + e * u + f
        r2_Two = r2_score(v, v_fit)
        equation_Two = f"y = {d:.4f}*x*log(x) + {e:.4f}*x + {f:.4f}"
    except ValueError:
        st.error("Linearithmic fir failed. Ensure all X values are positive.")
        st.stop()

elif fit_typeTwo == "Quadratic":
    degree_Two = 2
    coefficients_Two = np.polyfit(u, v, degree_Two)
    v_fit = np.polyval(coefficients_Two, u)
    r2_Two = r2_score(v, v_fit)
    equation_Two = f"y = {coefficients_Two[0]:.4f}*x² + {coefficients_Two[1]:.4f}*x + {coefficients_Two[2]:.4f}"

elif fit_typeTwo == "Cubic":
    degree_Two = 3
    coefficients_Two = np.polyfit(u, v, degree_Two)
    v_fit = np.polyval(coefficients_Two, u)
    r2_Two = r2_score(v, v_fit)
    equation_Two = f"y = {coefficients_Two[0]:.4f}*x³ + {coefficients_Two[1]:.4f}*x² + {coefficients_Two[2]:.4f}*x + {coefficients_Two[3]:.4f}"

elif fit_typeTwo == "Exponential":
    try:
        log_v = np.log(v)
        coefficients_Two = np.polyfit(u, log_v, 1)
        d = np.exp(coefficients_Two[1])
        e = coefficients_Two[0]
        v_fit = d * np.exp(e * u)
        r2_Two = r2_score(v, v_fit)
        equation_Two = f"y = {d:.4f}*exp({e:.4f}*x)"
    except ValueError:
        st.error("Exponential fit failed. Ensure all Y values are positive.")
        st.stop()

minimum = min(min(x), min(u))
maximum = max(max(x), max(u))
x_smooth = np.linspace(minimum, maximum, 500)
if fit_typeOne == "Logarithmic":
    y_smooth = coefficients[0] * np.log(x_smooth) + coefficients[1]
elif fit_typeOne == "Linearithmic":
    y_smooth = a * x_smooth * np.log(x_smooth) + b * x_smooth + c
elif fit_typeOne == "Exponential":
    y_smooth = a * np.exp(b * x_smooth)
else:
    y_smooth = np.polyval(coefficients, x_smooth)

u_smooth = np.linspace(minimum, maximum, 500)
if fit_typeTwo == "Logarithmic":
    v_smooth = coefficients_Two[0] * np.log(u_smooth) + coefficients_Two[1]
elif fit_typeTwo == "Linearithmic":
    v_smooth = d * u_smooth * np.log(u_smooth) + e * u_smooth + f
elif fit_typeTwo == "Exponential":
    v_smooth = d * np.exp(e * u_smooth)
else:
    v_smooth = np.polyval(coefficients_Two, u_smooth)

fig, ax = plt.subplots()
ax.scatter(x, y, color="red", label="Original Data One")
ax.scatter(u, v, color="blue", label="Original Data Two")
ax.plot(x_smooth, y_smooth, color="pink", label=f"{fit_typeOne} Fit (R²={r2:.4f})")
ax.plot(u_smooth, v_smooth, color="purple", label=f"{fit_typeTwo} Fit (R²={r2_Two:.4f})")
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.legend()
ax.set_title("fit")

st.pyplot(fig)

st.write(f"**Fitted Equation One**: {equation}")
st.write(f"**R² Value One**: {r2:.6f}")

st.write(f"**Fitted Equation Two**: {equation_Two}")
st.write(f"**R² Value Two**: {r2_Two:.6f}")