# импортируем необходимые библиотеки
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from datetime import datetime, timedelta
import json
import requests

# загружаем данные
url = "https://services.swpc.noaa.gov/json/ovation_aurora_latest.json"
response = requests.get(url)
data = json.loads(response.content)

# преобразуем данные в формат DataFrame
df = pd.DataFrame(data["coordinates"], columns=["longitude", "latitude", "aurora"])
df["datetime"] = pd.to_datetime(data["Forecast Time"])
df.set_index("datetime", inplace=True)
df = df.loc[:, ["latitude", "aurora"]]
df.columns = ["latitude", "kp"]
df.index.name = "datetime"
df = df.resample("H").mean()
df.fillna(method="ffill", inplace=True)

# выбираем данные для Санкт-Петербурга
data_spb = df.loc[(df["latitude"] >= 59.7) & (df["latitude"] <= 60.2)]

# создаем признаки для прогнозирования
data_spb["hour"] = data_spb.index.hour
data_spb["day_of_week"] = data_spb.index.dayofweek

# обучаем модель
X = data_spb[["hour", "day_of_week"]]
y = data_spb["kp"]
model = LinearRegression()

if len(X) > 0:
    model.fit(X, y)
else:
    print("No data available for training the model")

# делаем прогноз на следующие 24 часа
last_datetime = data_spb.index[-1]
next_datetime = last_datetime + timedelta(hours=1)
forecast = pd.DataFrame(index=pd.date_range(next_datetime, periods=24, freq="H"))
forecast["hour"] = forecast.index.hour
forecast["day_of_week"] = forecast.index.dayofweek
forecast["kp"] = model.predict(forecast[["hour", "day_of_week"]])

# выводим график прогноза
if len(data_spb) > 0:
    plt.plot(data_spb.index, data_spb["kp"], label="actual")
    plt.plot(forecast.index, forecast["kp"], label="forecast")
    plt.legend()
    plt.title("KP Index Forecast for St. Petersburg")
    plt.xlabel("Date")
    plt.ylabel("KP Index")
    plt.show()
else:
    print("No data available for the forecast")