import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
import datetime as dt

# Importing the dataset
data = pd.read_csv("https://raw.githubusercontent.com/iulianastroia/csv_data/master/final_dataframe.csv")

data['day'] = pd.to_datetime(data['day'], dayfirst=True)
# sort dates by day
data = data.sort_values(by=['day'])
print("sorted days", data.day)

group_by_df = pd.DataFrame(
    [name, group.mean().pm25] for name, group in data.groupby('day')
)
group_by_df.columns = ['day', 'pm25']

group_by_df['day'] = pd.to_datetime(group_by_df['day'])
group_by_df['day'] = group_by_df['day'].map(dt.datetime.toordinal)
print(group_by_df)

X = group_by_df['day'].values.reshape(-1, 1)
y = group_by_df['pm25'].values.reshape(-1, 1)
# # Splitting the dataset into training(70%) and test(30%)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3,
                                                    random_state=0
                                                    )

lin_reg = LinearRegression()
lin_reg.fit(X_train, y_train)  # training the algorithm

# Fitting Polynomial Regression to the dataset
poly_reg = PolynomialFeatures(degree=10)
X_poly = poly_reg.fit_transform(X_train)
print("x poly", X_poly)
pol_reg = LinearRegression()
pol_reg.fit(X_poly, y_train)

# print(group_by_df)
group_by_df['day'] = group_by_df['day'].map(dt.datetime.fromordinal)
print(group_by_df)


print("REAL DATA:", group_by_df['pm25'])
print("PREDICTED DATA:", pol_reg.predict(poly_reg.fit_transform(X)))

# Fitting Polynomial Regression to the dataset
from sklearn.preprocessing import PolynomialFeatures

poly_reg = PolynomialFeatures(degree=4)
X_poly = poly_reg.fit_transform(X)
pol_reg = LinearRegression()
pol_reg.fit(X_poly, y)

# Visualizing the Polymonial Regression results
import matplotlib.pyplot as plt


def viz_polymonial():
    plt.plot(X, y, color='red',marker='o')
    plt.plot(X, pol_reg.predict(poly_reg.fit_transform(X)), color='blue',marker='o')
    plt.title('Polynomial Regression')
    plt.xlabel('Day')
    plt.ylabel('Pm2.5')
    plt.show()
    return

viz_polymonial()