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("Fit Your Data")

default_data = {
    "X": [1, 2, 3, 4, 5],
    "Y": [2.2, 4.4, 6.5, 8.0, 10.1],
    "Select": [True, True, True, True, True]
}
data = pd.DataFrame(default_data)

with st.sidebar:
    st.subheader("Enter Your Data")
    user_data = st.data_editor(data, num_rows="dynamic", key="data_editor")
    xlabel = st.text_input("X label", "X-axis")
    ylabel = st.text_input("Y label", "Y-axis")
    fit_type = st.radio(
        "Choose the Type of Fit",
        options=["Logarithmic", "Linear", "Linearithmic", "Quadratic", "Cubic", "Exponential"],
        index=0
    )

try:
    selected_data = user_data[user_data["Select"]]
    x = np.array(selected_data["X"], dtype=float)
    y = np.array(selected_data["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()


if fit_type == "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_type == "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_type == "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_type == "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_type == "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_type == "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()

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


fig, ax = plt.subplots()
ax.scatter(x, y, color="red", label="Original Data")
ax.plot(x_smooth, y_smooth, color="blue", label=f"{fit_type} Fit (R²={r2:.4f})")
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.legend()
ax.set_title(equation)

st.pyplot(fig)

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