import gradio as gr
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
import numpy as np
import statsmodels.api as sm
import os

# Load main metrics
df = pd.read_csv("data/india_growth_metrics.csv")
df.columns = df.columns.str.strip()
cities = df['City'].unique()
metrics = df.columns[1:-1]  # Exclude City and Gini Coefficient

# Cluster setup
scaler = StandardScaler()
cluster_data = scaler.fit_transform(df[metrics])
kmeans = KMeans(n_clusters=4, random_state=0).fit(cluster_data)
df['Cluster'] = kmeans.labels_

def create_bar_chart(selected_cities, selected_metric):
    """Create bar chart for selected cities and metric"""
    if not selected_cities or not selected_metric:
        return None
    
    filtered = df[df['City'].isin(selected_cities)]
    fig = px.bar(filtered, x='City', y=selected_metric, color='City', 
                 title=f"{selected_metric} Comparison")
    return fig

def create_radar_chart(selected_cities):
    """Create radar chart for selected cities"""
    if not selected_cities:
        return None
    
    radar_df = df[df['City'].isin(selected_cities)].set_index('City')
    melted = radar_df[metrics].reset_index().melt(id_vars='City', var_name='Metric', value_name='Value')
    fig = px.line_polar(melted, r='Value', theta='Metric', color='City', 
                        line_close=True, title="Radar View")
    fig.update_traces(fill='toself')
    return fig

def create_correlation_matrix(selected_cities):
    """Create correlation matrix for selected cities"""
    if not selected_cities:
        return None
    
    corr = df[df['City'].isin(selected_cities)][metrics].corr()
    fig = px.imshow(corr, text_auto=True, title="Correlation Matrix")
    return fig

def generate_ai_insights(selected_cities):
    """Generate AI insights for selected cities"""
    if not selected_cities:
        return "Please select cities to generate insights."
    
    insights = []
    for city in selected_cities:
        city_data = df[df['City'] == city]
        if not city_data.empty:
            highest = city_data[metrics].idxmax(axis=1).values[0]
            lowest = city_data[metrics].idxmin(axis=1).values[0]
            insights.append(f"🏙️ **{city}** excels in **{highest}** but needs improvement in **{lowest}**.")
    
    return "\n".join(insights) if insights else "No insights available."

def find_twin_cities(selected_cities):
    """Find twin cities for selected cities"""
    if not selected_cities:
        return "Please select cities to find twin cities."
    
    twin_info = []
    for city in selected_cities:
        city_data = df[df['City'] == city]
        if not city_data.empty:
            city_vec = city_data[metrics].values
            df_temp = df.copy()
            df_temp['distance'] = np.linalg.norm(df_temp[metrics].values - city_vec, axis=1)
            nearest = df_temp[df_temp['City'] != city].sort_values('distance').iloc[0]['City']
            twin_info.append(f"🏙️ **{city}**'s most similar city is **{nearest}**.")
    
    return "\n".join(twin_info) if twin_info else "No twin cities found."

def create_time_series(selected_metric, selected_city):
    """Create time series forecast for selected metric and city"""
    if not selected_metric or not selected_city:
        return None
    
    try:
        time_df = pd.read_csv("data/timeseries.csv", encoding='utf-8-sig')
        time_df.columns = time_df.columns.str.strip()
        time_df['City'] = time_df['City'].str.strip()
        
        fig = go.Figure()
        city_df = time_df[time_df['City'] == selected_city].sort_values('Year')
        city_df['Year'] = pd.to_datetime(city_df['Year'], format='%Y')

        if selected_metric in city_df.columns:
            ts = city_df.set_index('Year')[selected_metric]
            fig.add_trace(go.Scatter(x=ts.index, y=ts.values, mode='lines+markers', 
                                   name=selected_metric))
            try:
                model = sm.tsa.ARIMA(ts, order=(1, 1, 0)).fit()
                forecast = model.get_forecast(steps=3)
                forecast_index = pd.date_range(start=ts.index[-1] + pd.offsets.YearBegin(), 
                                             periods=3, freq='YS')
                forecast_series = pd.Series(forecast.predicted_mean.values, index=forecast_index)
                fig.add_trace(go.Scatter(x=forecast_series.index, y=forecast_series.values,
                                       mode='lines+markers', name=f"{selected_metric} Forecast",
                                       line=dict(dash='dash')))
            except:
                pass

        fig.update_layout(title=f"Time Series Forecast: {selected_metric} in {selected_city}", 
                         xaxis_title='Year')
        return fig
    except Exception as e:
        return None

def create_cluster_view(metric):
    """Create cluster visualization"""
    if not metric:
        return None
    
    fig = px.scatter(df, x=metric, y='Gini Coefficient', color='Cluster',
                     hover_data=['City'], title="City Clustering")
    return fig

# Create Gradio interface
with gr.Blocks(title="India Growth Metrics Dashboard", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🇮🇳 Decode India: Smart AI-Enhanced Dashboard for Growth Metrics")
    gr.Markdown("A powerful interactive dashboard that visualizes and analyzes growth metrics for 35 major Indian cities using AI, time series forecasting, and cluster-based insights.")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Dashboard Controls")
            
            city_selector = gr.Dropdown(
                choices=list(cities),
                value=list(cities[:3]),
                label="Select Cities",
                multiselect=True
            )
            
            metric_selector = gr.Dropdown(
                choices=list(metrics),
                value=metrics[0],
                label="Select Metric"
            )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📈 Bar Chart Comparison")
            bar_chart = gr.Plot(label="Bar Chart")
        
        with gr.Column(scale=1):
            gr.Markdown("### 🌐 Radar Chart")
            radar_chart = gr.Plot(label="Radar Chart")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🧮 Correlation Matrix")
            correlation_matrix = gr.Plot(label="Correlation Matrix")
        
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 City Clustering")
            cluster_view = gr.Plot(label="Cluster View")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📈 Time Series Forecast")
            time_series = gr.Plot(label="Time Series")
        
        with gr.Column(scale=1):
            gr.Markdown("### 🧠 AI Insights")
            ai_insights = gr.Markdown(label="AI Insights")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🌍 Twin City Recommendations")
            twin_cities = gr.Markdown(label="Twin Cities")
    
    # Event handlers
    def update_all_outputs(selected_cities, selected_metric):
        # Time series uses the first selected city
        time_series_city = selected_cities[0] if selected_cities else None

        bar_fig = create_bar_chart(selected_cities, selected_metric)
        radar_fig = create_radar_chart(selected_cities)
        corr_fig = create_correlation_matrix(selected_cities)
        cluster_fig = create_cluster_view(selected_metric)
        ai_md = generate_ai_insights(selected_cities)
        twin_md = find_twin_cities(selected_cities)
        ts_fig = create_time_series(selected_metric, time_series_city)

        return bar_fig, radar_fig, corr_fig, cluster_fig, ts_fig, ai_md, twin_md

    # A list of all outputs for convenience
    all_outputs = [bar_chart, radar_chart, correlation_matrix, cluster_view, time_series, ai_insights, twin_cities]

    # Connect inputs to outputs
    city_selector.change(
        fn=update_all_outputs,
        inputs=[city_selector, metric_selector],
        outputs=all_outputs
    )
    
    metric_selector.change(
        fn=update_all_outputs,
        inputs=[city_selector, metric_selector],
        outputs=all_outputs
    )
    
    demo.load(
        fn=update_all_outputs,
        inputs=[city_selector, metric_selector],
        outputs=all_outputs
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()