import gradio as gr
import pandas as pd
import plotly.graph_objects as go
import numpy as np

def plot_real_estate_correlation(state):
    # Read the CSV file
    df = pd.read_csv('https://files.zillowstatic.com/research/public_csvs/zhvi/Zip_zhvi_uc_sfrcondo_tier_0.33_0.67_sm_sa_month.csv')

    # Filter for the given state
    df = df[df['State'] == state.upper()]

    # Extract the list of ZIP codes and filter only columns that are date strings
    zip_codes = df['RegionName'].unique()
    
    # Extract columns that are valid date strings only
    date_columns = []
    for col in df.columns[7:]:
        try:
            # Try to parse column names as dates
            pd.to_datetime(col)
            date_columns.append(col)
        except:
            continue
    
    # Initialize a DataFrame to hold price data for correlation calculation
    price_matrix = []

    # Loop through each ZIP code in the state
    for zip_code in zip_codes:
        df_zip = df[df['RegionName'] == zip_code]

        # Extract only the columns with valid date data (price values)
        prices = df_zip.loc[:, date_columns].values.flatten()

        # Append prices to the matrix if there are no missing values
        if not np.isnan(prices).all():
            price_matrix.append(prices)

    # Convert to DataFrame for easier manipulation
    price_matrix_df = pd.DataFrame(price_matrix, index=zip_codes, columns=date_columns)

    # Transpose to align for correlation calculation (each column = ZIP code)
    price_matrix_df = price_matrix_df.T.dropna()

    # Calculate the correlation matrix for ZIP codes
    corr_matrix = price_matrix_df.corr()

    # Prepare the grid data for 3D plot
    z_data = corr_matrix.values
    x_data, y_data = np.meshgrid(zip_codes, zip_codes)

    # Create the 3D surface plot
    fig = go.Figure(data=[go.Surface(z=z_data, x=x_data, y=y_data)])

    # Update plot layout
    fig.update_layout(
        title=f'3D Correlation Matrix of Housing Prices in {state}',
        scene=dict(
            xaxis_title='ZIP Code',
            yaxis_title='ZIP Code',
            zaxis_title='Correlation',
        ),
        autosize=True
    )

    return fig

iface = gr.Interface(fn=plot_real_estate_correlation,
                     inputs=[gr.components.Textbox(label="State (e.g., 'NJ' for New Jersey)")],
                     outputs=gr.Plot())

iface.launch(share=False, debug=True)