import os
import requests
import pandas as pd
import plotly.express as px
import gradio as gr
from sklearn.preprocessing import StandardScaler

# Fetch the FRED API key from environment variables
API_KEY = os.getenv("FRED_API_KEY")

def fetch_data(series_id, frequency="m", adjustment="sa"):
    """
    Fetch data from FRED API based on the provided series ID, frequency, and adjustment type.
    """
    if not API_KEY:
        raise ValueError("FRED API key not set. Please set the FRED_API_KEY environment variable.")
    
    url = "https://api.stlouisfed.org/fred/series/observations"
    params = {
        'api_key': API_KEY,
        'series_id': series_id,
        'file_type': 'json',
        'frequency': frequency,
        'seasonal_adjustment': adjustment,
    }
    try:
        response = requests.get(url, params=params)
        response.raise_for_status()
        data = response.json()
        return data
    except requests.exceptions.RequestException as e:
        print(f"Error fetching data for {series_id}: {e}")
        return {}

def process_data(data):
    """
    Process the FRED data into a pandas DataFrame.
    """
    if 'observations' not in data:
        print("No observations found in data.")
        return pd.DataFrame()
    
    df = pd.DataFrame(data['observations'])
    df['date'] = pd.to_datetime(df['date'])
    df['value'] = pd.to_numeric(df['value'], errors='coerce')
    return df

def standardize_series(df):
    """
    Standardize the 'value' column in the dataframe using Z-scores.
    """
    scaler = StandardScaler()
    df['standardized_value'] = scaler.fit_transform(df[['value']])
    return df

def create_combined_2d_visualization(dataframes, labels):
    """
    Generate a combined 2D line plot using Plotly.
    Each dataframe will be plotted on the same axes.
    """
    combined_df = pd.concat(dataframes, keys=labels, names=['series', 'index']).reset_index(level='series')
    fig = px.line(
        combined_df,
        x='date',
        y='standardized_value',
        color='series',
        title="Combined Economic Data 2D Visualization",
        labels={"standardized_value": "Standardized Value", "date": "Date", "series": "Data Type"}
    )
    fig.update_layout(
        width=1200,
        height=600,
        xaxis_title="Date",
        yaxis_title="Standardized Value",
        legend_title="Data Type"
    )
    return fig

def visualize_multiple_series(series_names):
    """
    Fetch, standardize, and combine multiple datasets for visualization.
    """
    dataframes = []
    labels = []
    for series_name in series_names:
        series_id = series_options.get(series_name)
        frequency = default_frequencies.get(series_id, "m")
        adjustment = default_adjustments.get(series_id, "sa")
        data = fetch_data(series_id, frequency, adjustment)
        df = process_data(data)
        if not df.empty:
            standardized_df = standardize_series(df)
            dataframes.append(standardized_df)
            labels.append(series_name)
    
    if not dataframes:
        raise ValueError("No valid data to visualize.")
    return create_combined_2d_visualization(dataframes, labels)

# Define default frequencies and adjustments
default_frequencies = {
    "GDP": "q",
    "UNRATE": "m",
    "CPIAUCSL": "m",
    "FEDFUNDS": "m",
    "MORTGAGE30US": "w"
}

default_adjustments = {
    "GDP": "sa",
    "UNRATE": "sa",
    "CPIAUCSL": "sa",
    "FEDFUNDS": "nsa",
    "MORTGAGE30US": "nsa"
}

# Define options for each dropdown
series_options = {
    "Gross Domestic Product (GDP)": "GDP",
    "Unemployment Rate (UNRATE)": "UNRATE",
    "Consumer Price Index (CPI - All Urban Consumers)": "CPIAUCSL",
    "Federal Funds Rate": "FEDFUNDS",
    "30-Year Fixed Mortgage Rate": "MORTGAGE30US"
}

# Gradio Interface using Blocks
with gr.Blocks() as demo:
    gr.Markdown("# FRED Combined Data 2D Visualizer")
    gr.Markdown("Choose multiple economic indicators to visualize them together in a 2D space.")
    
    with gr.Row():
        series_dropdown = gr.CheckboxGroup(
            choices=list(series_options.keys()),
            label="Select Economic Indicators to Compare"
        )
    
    plot_output = gr.Plot()

    # Explanation Section
    with gr.Accordion("Color Coding Explanation", open=True):
        gr.Markdown("""
        - **Gross Domestic Product (GDP)**: Displayed in **blue**.
        - **Unemployment Rate (UNRATE)**: Displayed in **green**.
        - **Consumer Price Index (CPI - All Urban Consumers)**: Displayed in **red**.
        - **Federal Funds Rate**: Displayed in **purple**.
        - **30-Year Fixed Mortgage Rate**: Displayed in **orange**.
        """)
    
    # Define interaction
    series_dropdown.change(
        visualize_multiple_series, 
        inputs=[series_dropdown], 
        outputs=[plot_output]
    )

demo.launch(debug=True)