import gradio as gr
from gradio.components import Markdown, Textbox, Button
import matplotlib.pyplot as plt
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.svm import SVR
from sklearn.pipeline import make_pipeline
from sunpy.net import Fido
from sunpy.net import attrs as a
from sunpy.timeseries import TimeSeries




def process_data():
    # Define the time range for data retrieval
    tstart = "2015-06-21 01:00"
    tend = "2015-06-21 23:00"

    # Query and fetch GOES XRS data
    result_goes15 = Fido.search(a.Time(tstart, tend), a.Instrument("XRS"), a.goes.SatelliteNumber(15), a.Resolution("flx1s"))
    files = Fido.fetch(result_goes15)

    # Load the data into a TimeSeries
    goes_15 = TimeSeries(files, concatenate=True)

    # Extract X-ray flux and time data
    flux_data = goes_15.quantity("xrsb").value
    time_data = goes_15.time.datetime

    # Create a feature matrix with time data (as numerical values)
    X = np.array([(t - time_data[0]).total_seconds() for t in time_data]).reshape(-1, 1)

    # Split the data into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(X, flux_data, test_size=0.2, random_state=42)

    # Train a linear regression model
    linear_model = LinearRegression()
    linear_model.fit(X_train, y_train)

    # Train a quadratic regression model
    quadratic_model = make_pipeline(PolynomialFeatures(degree=2), LinearRegression())
    quadratic_model.fit(X_train, y_train)

    # Train a cubic regression model
    cubic_model = make_pipeline(PolynomialFeatures(degree=3), LinearRegression())
    cubic_model.fit(X_train, y_train)

    # Train a support vector regression (SVR) model
    svr_model = SVR(kernel='linear')
    svr_model.fit(X_train, y_train)

    # Make predictions using all models
    y_pred_linear = linear_model.predict(X_test)
    y_pred_quadratic = quadratic_model.predict(X_test)
    y_pred_cubic = cubic_model.predict(X_test)
    y_pred_svr = svr_model.predict(X_test)

    # Plot the actual and predicted data from all models
    plt.figure(figsize=(12, 6))
    plt.scatter(X_test, y_test, color='blue', label='Actual Data')
    plt.plot(X_test, y_pred_linear, color='red', linewidth=2, label='Linear Prediction')
    plt.plot(X_test, y_pred_quadratic, color='green', linewidth=2, label='Quadratic Prediction')
    plt.plot(X_test, y_pred_cubic, color='orange', linewidth=2, label='Cubic Prediction')
    plt.plot(X_test, y_pred_svr, color='purple', linewidth=2, label='SVR Prediction')

    # Include solar flux data as an additional line in the plot
    plt.plot(X, flux_data, color='cyan', linestyle='dashed', label='Solar Flux')

    plt.title('GOES XRS Space Weather Forecast')
    plt.xlabel('Time (seconds since start)')
    plt.ylabel('X-ray Flux / Solar Flux')
    plt.legend()

    # Save the image
    plt.savefig('space_weather_forecast.png')

    # Display the plot
    #plt.show()
    fig = plt.figure()


process_data()

with gr.Blocks(title="Space Weather Forecast", analytics_enabled=False) as spaceml:
    gr.Markdown("# Space Weather Forecast")
    gr.Markdown("Welcome to the Space Weather Forecast!")
    with gr.Row():
        with gr.Column(scale=1):
            gradio_plot = gr.Image('space_weather_forecast.png')

spaceml.queue().launch(show_api=True, share=True)