import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error
import gradio as gr
import plotly.express as px
import plotly.graph_objects as go

# Load the dataset
df = pd.read_csv('california_housing_train.csv')

# Select features and target
features = df[['longitude', 'latitude', 'housing_median_age', 'total_rooms',
               'total_bedrooms', 'population', 'households', 'median_income']]
target = df['median_house_value']

# Split the data
X_train, X_test, y_train, y_test = train_test_split(features, target, test_size=0.2, random_state=42)

# Standardize the data
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# Initialize lists to store loss metrics
training_losses = []
validation_losses = []

# Custom MLPRegressor class to capture loss metrics
class CustomMLPRegressor(MLPRegressor):
    def _fit(self, X, y, incremental):
        result = super()._fit(X, y, incremental)
        training_loss = self.loss_
        predictions = self.predict(X_test_scaled)
        validation_loss = mean_squared_error(y_test, predictions)
        training_losses.append(training_loss)
        validation_losses.append(validation_loss)
        return result

# Train the model
model = CustomMLPRegressor(hidden_layer_sizes=(100,), activation='relu', solver='adam', max_iter=1000)
model.fit(X_train_scaled, y_train)

# Create prediction function
def predict_house_price(longitude, latitude, housing_median_age, total_rooms,
                        total_bedrooms, population, households, median_income):
    input_data = scaler.transform([[longitude, latitude, housing_median_age, total_rooms,
                                    total_bedrooms, population, households, median_income]])
    prediction = model.predict(input_data)[0]
    return f"${prediction:,.2f}"

# Create dashboard function
def create_dashboard():
    fig1 = px.scatter(df, x='longitude', y='latitude', color='median_house_value',
                      title="House Prices by Location",
                      labels={'longitude': 'Longitude', 'latitude': 'Latitude', 'median_house_value': 'House Value'})
    
    fig2 = px.histogram(df, x='median_income', nbins=30, title="Distribution of Median Income",
                        labels={'median_income': 'Median Income'})
    
    fig3 = px.histogram(df, x='housing_median_age', nbins=30, title="Distribution of Housing Median Age",
                        labels={'housing_median_age': 'Housing Median Age'})
    
    fig4 = go.Figure()
    fig4.add_trace(go.Scatter(y=training_losses, mode='lines', name='Training Loss'))
    fig4.add_trace(go.Scatter(y=validation_losses, mode='lines', name='Validation Loss'))
    fig4.update_layout(title="Model Loss Over Time", xaxis_title="Epoch", yaxis_title="Loss")
    
    return fig1, fig2, fig3, fig4

# Gradio interface for prediction
iface_predict = gr.Interface(
    fn=predict_house_price,
    inputs=[
        gr.Number(label="Longitude", info="Enter the longitude of the house."),
        gr.Number(label="Latitude", info="Enter the latitude of the house."),
        gr.Number(label="Housing Median Age", info="Enter the median age of the house."),
        gr.Number(label="Total Rooms", info="Enter the total number of rooms."),
        gr.Number(label="Total Bedrooms", info="Enter the total number of bedrooms."),
        gr.Number(label="Population", info="Enter the population in the area."),
        gr.Number(label="Households", info="Enter the number of households in the area."),
        gr.Number(label="Median Income", info="Enter the median income of the households.")
    ],
    outputs="text",
    title="House Price Prediction",
    description="Enter the features to get the predicted house price."
)

# Gradio interface for dashboard
iface_dashboard = gr.Interface(
    fn=create_dashboard,
    inputs=[],
    outputs=[gr.Plot(), gr.Plot(), gr.Plot(), gr.Plot()],
    title="House Price Dashboard",
    description="Visualizations of the housing dataset and model performance."
)

# Launch both interfaces
iface = gr.TabbedInterface([iface_predict, iface_dashboard], ["Prediction", "Dashboard"])

iface.launch()