import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

df = pd.read_csv("normalized_train_dataset_with_train_number.csv")


X = df[["departure_station", "arrival_station", "departure_time", "arrival_time", "day_of_week", "train_type", "distance_km"]]
y = df["train_number"]    #Defining features and target

model = LinearRegression()  #Initializing the model
model.fit(X, y)  # Training the model

weights = model.coef_  #Extracting weights and biases
bias = model.intercept_

feature_names = X.columns
print("Learned Weights:")
for feature, weight in zip(feature_names, weights):
    print(f"{feature}: {weight:.4f}")

print(f"\nLearned Bias (Intercept): {bias:.4f}")

y_pred = model.predict(X)   # Evaluate the model
mse = mean_squared_error(y, y_pred)
print(f"\nMean Squared Error on Training Data: {mse:.4f}")