import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error

df = pd.read_csv("/content/Salary Data.csv")

print("\nDimensions: (Rows, Columns)", df.shape)
print("\nFirst 5 rows: \n", df.head())
print("\nData types and non-null count: \n", df.info())
print("\nMissing values: \n", df.isnull().sum())
print("\nStatistics: \n", df.describe(include="all").T)

print("\n\nColumn Names:", df.columns)

df = df.dropna()
print(df.shape)
plt.scatter(df["YearsExperience"], df['Salary'], s = 5)
plt.xlabel("Years of Experience")
plt.ylabel("Salary")
plt.title("Years of Experience vs Salary")
plt.show()

X = df[["YearsExperience"]]
y = df["Salary"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2,
random_state = 42)

model = LinearRegression()
model.fit(X_train, y_train)

y_pred = model.predict(X_test)
comparison = pd.DataFrame(
    {
        "Actual Salary": y_test.values,
        "Predicted Salary": y_pred
    }
)

print("\nActual vs Predicted Salary: \n", comparison.head())

r2 = r2_score(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))

print("\nMODEL PERFORMANCE")
print("R2 Score: ", r2)
print("MAE: ", mae)
print("RMSE: ", rmse)


# 5. REGRESSION EQUATION

m = model.coef_[0]
c = model.intercept_

print("Salary =", m, "* Experience +", c)

# FITTED REGRESSION LINE

x = np.sort(df["YearsExperience"])

plt.scatter(X["YearsExperience"], y)
plt.plot(x, m*x + c, color="red")
plt.xlabel("Experience")
plt.ylabel("Salary")
plt.show()

# 6. USER INPUT PREDICTION

years = float(input("Enter years of experience: "))
salary = m * years + c

print("Predicted Salary:", salary)
