import pandas as pd
import numpy as np

np.random.seed(42)

import warnings
from sklearn.exceptions import ConvergenceWarning
warnings.filterwarnings("ignore", category=ConvergenceWarning)

fish_data = pd.read_csv("fish_data.csv", index_col=0)
fish_data = fish_data.drop(columns=["ID"])
fish_data

y = fish_data["Weight"]
X = fish_data.drop(columns=["Weight"])

from sklearn.model_selection import train_test_split 

X_train_raw, X_test_raw, y_train, y_test = train_test_split(X, y)
X_train_raw, X_test_raw, y_train, y_test = train_test_split(X, y, test_size=0.3, stratify=X["Species"])

from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import make_column_transformer

categorical_columns = ["Species"] 

transformer = make_column_transformer(
    (OneHotEncoder(sparse_output=False), ["Species"]),
    remainder="passthrough"
)

X_train_onehot = transformer.fit_transform(X_train_raw)
X_test_onehot = transformer.transform(X_test_raw)

pd.DataFrame(X_train_onehot, columns=transformer.get_feature_names_out())

from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()
scaler.fit(X_train_onehot)

X_train = scaler.transform(X_train_onehot)
X_test = scaler.transform(X_test_onehot)

from sklearn.linear_model import LinearRegression, Lasso
from sklearn.svm import SVR

model_zoo = {
    "linear_regression": LinearRegression(),
    "lasso_var1": Lasso(alpha=1.0),
    "lasso_var2": Lasso(alpha=1e-03),
    "SVR_rbf": SVR(kernel="rbf", C=1e04),
    "SVR_poly": SVR(kernel="poly",C=1e04),
}

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

def fit_and_eval(X_train, y_train, X_test, y_test, model, name):
    """ 1. Natrénuje model na trénovací množině.
        2. Spočte hodnoty metrik na trénovací i testovací množině.
        vrátí slovník ve tvaru {"název metriky": hodnota} 
    """ 
    
    model.fit(X_train, y_train)
    
    y_train_pred = model.predict(X_train)
    y_test_pred = model.predict(X_test)
    
    return {
        "MAE_train": mean_absolute_error(y_train, y_train_pred),
        "MSE_train": mean_squared_error(y_train, y_train_pred),
        "MAE_test": mean_absolute_error(y_test, y_test_pred),
        "MSE_test": mean_squared_error(y_test, y_test_pred)
    }

results = []
for name, model in model_zoo.items(): 
    result = fit_and_eval(X_train, y_train, X_test, y_test, model, name)
    result["model"] = name
    results.append(result)
    
pd.DataFrame(results)

# doplň jméno modelu, který jsi vybrala

best_model = model_zoo["SVR_rbf"]

test_data = pd.read_csv("fish_data_test.csv", index_col=0)

test_data.pop("ID")
y_real_test = test_data.pop("Weight")
X_real_test = test_data 

X_real_test_transformed = transformer.transform(X_real_test)
X_real_test_scaled = scaler.transform(X_real_test_transformed)

y_pred_test = best_model.predict(X_real_test_scaled)

print(f"MAE {mean_absolute_error(y_real_test, y_pred_test):.3f}")
print(f"MSE {mean_squared_error(y_real_test, y_pred_test):.3f}")
print(f"R2  {r2_score(y_real_test, y_pred_test):.3f}")

for weight, predicted_weight in zip(y_real_test, y_pred_test):
    print(f"{weight:>10.1f}     {predicted_weight:>10.1f}")

is_bream = X_real_test["Species"] == "Bream"
predicted_bream_weights = best_model.predict(X_real_test_scaled[is_bream])

is_roach = X_real_test["Species"] == "Roach"
predicted_roach_weights = best_model.predict(X_real_test_scaled[is_roach])

result_bream = pd.DataFrame()
result_bream["length"] = X_real_test[is_bream]["Length3"]
result_bream["weight"] = y_real_test[is_bream]
result_bream["predicted"] = predicted_bream_weights
result_bream = result_bream.sort_values("length")

result_roach = pd.DataFrame()
result_roach["length"] = X_real_test[is_roach]["Length3"]
result_roach["weight"] = y_real_test[is_roach]
result_roach["predicted"] = predicted_roach_weights
result_roach = result_roach.sort_values("length")

import matplotlib.pyplot as plt 

fig, ax = plt.subplots(1, 2)

ax[0].plot(result_bream["length"], result_bream["weight"], label="true weight", marker="o");
ax[0].plot(result_bream["length"], result_bream["predicted"], label="prediction", marker="o");
ax[0].legend()
ax[0].set_title("Bream")

ax[1].plot(result_roach["length"], result_roach["weight"], label="true weight", marker="o");
ax[1].plot(result_roach["length"], result_roach["predicted"], label="prediction", marker="o");
ax[1].legend()
ax[1].set_title("Roach");