import requests
import pandas as pd
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import gradio as gr
import datetime
import os

# API Endpoints
WEATHER_API = "https://api.open-meteo.com/v1/forecast"
ELECTRICITY_PRICE_API = "https://www.elprisetjustnu.se/api/v1/prices"
ENERGY_CHARTS_API = "https://energy-charts.info/api/public_power"

# Fetch historical weather data
def fetch_weather_data(start_date="2023-01-01", end_date=None):
    if end_date is None:
        end_date = datetime.datetime.now().strftime('%Y-%m-%d')

    params = {
        "latitude": 59.3293,  # Stockholm latitude
        "longitude": 18.0686,  # Stockholm longitude
        "daily": "temperature_2m_mean,precipitation_sum,wind_speed_10m_max,wind_direction_10m_dominant",
        "start_date": start_date,
        "end_date": end_date,
        "timezone": "Europe/Stockholm"
    }
    response = requests.get(WEATHER_API, params=params)
    response.raise_for_status()
    daily_data = response.json()["daily"]
    return pd.DataFrame(daily_data)

# Fetch historical electricity prices
def fetch_electricity_prices():
    today = datetime.datetime.now().strftime('%Y/%m-%d')
    url = f"{ELECTRICITY_PRICE_API}/{today}_SE3.json"
    response = requests.get(url)
    response.raise_for_status()
    return pd.DataFrame(response.json())

# Fetch energy production data using Energy-Charts API
def fetch_energy_production_data(start_date="2023-01-01", end_date=None):
    if end_date is None:
        end_date = datetime.datetime.now().strftime('%Y-%m-%d')

    params = {
        "country": "se",  # Sweden country code
        "start": start_date,
        "end": end_date
    }
    response = requests.get(ENERGY_CHARTS_API, params=params)
    response.raise_for_status()
    data = response.json()
    production_data = {
        "unix_seconds": data["unix_seconds"],
        **{ptype["name"]: ptype["data"] for ptype in data["production_types"]}
    }
    return pd.DataFrame(production_data)

# Prepare the dataset
def prepare_dataset(weather_data, electricity_data, energy_data):
    dataset = pd.concat([weather_data, electricity_data, energy_data], axis=1)
    dataset = dataset.dropna()
    return dataset

# Train the model
def train_model(dataset):
    X = dataset.drop("price", axis=1)
    y = dataset["price"]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    model = xgb.XGBRegressor(objective="reg:squarederror", n_estimators=100, learning_rate=0.1)
    model.fit(X_train, y_train)
    predictions = model.predict(X_test)

    rmse = mean_squared_error(y_test, predictions, squared=False)
    model.save_model("electricity_price_model.json")
    print(f"Model trained with RMSE: {rmse}")
    return model

# Load the model and make predictions
def predict_price(features):
    model = xgb.XGBRegressor()
    model.load_model("electricity_price_model.json")
    prediction = model.predict(pd.DataFrame([features]))
    return prediction[0]

# Update predictions and save to file
def update_predictions():
    # Fetch the latest data
    weather_data = fetch_weather_data()
    electricity_data = fetch_electricity_prices()
    energy_data = fetch_energy_production_data()

    # Prepare dataset
    dataset = prepare_dataset(weather_data, electricity_data, energy_data)

    # Load the model
    model = xgb.XGBRegressor()
    model.load_model("electricity_price_model.json")

    # Generate predictions
    predictions = model.predict(dataset.drop("price", axis=1))

    # Save predictions to file
    predictions_output = dataset.copy()
    predictions_output["predicted_price"] = predictions
    predictions_output.to_json("predictions.json", orient="records")
    print("Predictions updated and saved.")

# Gradio Interface
def gradio_interface():
    def wrapper(temp, precip, wind_speed, humidity, energy_price, electricity_price):
        features = {
            "temperature_2m": temp,
            "precipitation": precip,
            "wind_speed_10m": wind_speed,
            "humidity": humidity,
            "energy_price": energy_price,
            "electricity_price": electricity_price
        }
        return predict_price(features)

    interface = gr.Interface(
        fn=wrapper,
        inputs=[
            gr.inputs.Number(label="Temperature (°C)"),
            gr.inputs.Number(label="Precipitation (mm)"),
            gr.inputs.Number(label="Wind Speed (m/s)"),
            gr.inputs.Number(label="Humidity (%)"),
            gr.inputs.Number(label="Energy Production Price"),
            gr.inputs.Number(label="Historical Electricity Price")
        ],
        outputs=gr.outputs.Textbox(label="Predicted Electricity Price"),
        title="Electricity Price Prediction",
        description="Predict future electricity prices based on weather and energy data."
    )

    interface.launch()

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "update":
        update_predictions()
    else:
        # Fetch data
        weather_data = fetch_weather_data()
        electricity_data = fetch_electricity_prices()
        energy_data = fetch_energy_production_data()

        # Prepare dataset and train the model
        dataset = prepare_dataset(weather_data, electricity_data, energy_data)
        train_model(dataset)

        # Launch Gradio interface
        gradio_interface()