import dash
from dash import html, dcc, Input, Output, callback
import dash_bootstrap_components as dbc
import pandas as pd
import plotly.express as px
import os

dash.register_page(__name__, path='/analysis')

# --- DATA LOADING LOGIC ---
# Using absolute paths to ensure the app finds the CSVs inside the /pages folder
current_dir = os.path.dirname(__file__)
features_path = os.path.join(current_dir, 'sel_features.csv')
target_path = os.path.join(current_dir, 'sel_target.csv')

try:
    # sel_features is tab-separated based on your Jupyter logic
    df_features = pd.read_csv(features_path, sep='\t')
    # sel_target is comma-separated
    df_target = pd.read_csv(target_path)
    
    # Merge for multivariate analysis
    df = pd.concat([df_features, df_target], axis=1)
    
    # List of bands for dropdowns and melting
    band_columns = [col for col in df_features.columns]
    data_loaded = True
except Exception as e:
    print(f"Error loading data: {e}")
    data_loaded = False

# --- LAYOUT ---
layout = html.Div([
    html.Div([
        html.H1("Data Analysis: Exploring the Galaxy Sample", className="text-white fw-bold mb-2"),
        html.P("Multivariate Analysis: Distributions, Correlations, and Ranges", className="lead text-info"),
    ], className="mb-5"),

    # SECTION 1: VIOLIN PLOT (Multicolored)
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H4("Photometric Band Distributions", className="text-info mb-3"),
                    html.P("Comparative density and magnitude ranges for all filter bands (u, g, r, i, z, y).", 
                           className="text-muted small"),
                    dcc.Graph(id='violin-plot'),
                ])
            ], className="modern-card mb-4"),
        ], width=12)
    ]),

    # SECTION 2: TARGET DISTRIBUTION (Viridis Purple)
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H4("Redshift Distribution (zhelio)", className="text-info mb-3"),
                    html.P("Ground-truth redshift distribution from DEEP2/3 and 3D-HST surveys.", 
                           className="text-muted small"),
                    dcc.Graph(
                        id='zhelio-dist',
                        figure=px.histogram(
                            df, x="zhelio", nbins=50, 
                            template="plotly_dark",
                            labels={'zhelio': 'True Redshift (z)'},
                            color_discrete_sequence=['#440154'] # Deep Purple from Viridis
                        ).update_layout(
                            paper_bgcolor='rgba(0,0,0,0)', 
                            plot_bgcolor='rgba(0,0,0,0)',
                        ) if data_loaded else {}
                    )
                ])
            ], className="modern-card mb-4"),
        ], width=12),
    ]),

    # SECTION 3: CORRELATION HEATMAP (Viridis Palette)
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H4("Feature Correlation Matrix", className="text-info mb-3"),
                    html.P("Mathematical relationship between photometric features and the target redshift.", 
                           className="text-muted small"),
                    dcc.Graph(id='correlation-heatmap'),
                ])
            ], className="modern-card mb-4"),
        ], width=12),
    ]),

    # SECTION 4: FEATURE ANALYSIS (Interactive Scatter)
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardBody([
                    html.H5("Feature Analysis", className="text-info mb-3"),
                    html.Label("Select Photometric Band to Analyze:", className="text-light"),
                    dcc.Dropdown(
                        id='band-selector',
                        options=[{'label': b, 'value': b} for b in band_columns] if data_loaded else [],
                        value=band_columns[0] if data_loaded else None,
                        className="mb-4", 
                        style={'color': '#000'}
                    ),
                    dcc.Graph(id='redshift-scatter-plot'),
                ])
            ], className="modern-card mb-4"),
        ], width=12),
    ]),
])

# --- CALLBACKS ---

# 1. Callback for Multicolored Violin Plot
@callback(
    Output('violin-plot', 'figure'),
    Input('band-selector', 'value')
)
def update_violin(_):
    if not data_loaded: return {}
    df_long = pd.melt(df, value_vars=band_columns, var_name='Band', value_name='Magnitude')
    fig = px.violin(
        df_long, x='Band', y='Magnitude', color='Band',
        box=True, points="all", template="plotly_dark",
        color_discrete_sequence=px.colors.qualitative.Vivid # Distinct colors per band
    )
    fig.update_layout(
        paper_bgcolor='rgba(0,0,0,0)', 
        plot_bgcolor='rgba(0,0,0,0)', 
        showlegend=False,
        yaxis_title="Magnitude (Brightness)"
    )
    return fig

# 2. Callback for Viridis Heatmap
@callback(
    Output('correlation-heatmap', 'figure'),
    Input('band-selector', 'value')
)
def update_heatmap(_):
    if not data_loaded: return {}
    corr = df.corr()
    fig = px.imshow(
        corr, text_auto=".2f",
        color_continuous_scale='Viridis',
        template="plotly_dark"
    )
    fig.update_layout(
        paper_bgcolor='rgba(0,0,0,0)', 
        plot_bgcolor='rgba(0,0,0,0)',
        coloraxis_colorbar=dict(title="Correlation")
    )
    return fig

# 3. Callback for Viridis Scatter Plot
@callback(
    Output('redshift-scatter-plot', 'figure'),
    Input('band-selector', 'value')
)
def update_scatter(selected_band):
    if not data_loaded or not selected_band: return {}
    fig = px.scatter(
        df, x=selected_band, y="zhelio", 
        color="zhelio",
        color_continuous_scale='Viridis', 
        template="plotly_dark",
        labels={selected_band: f"Magnitude ({selected_band})", "zhelio": "True Redshift (z)"}
    )
    fig.update_layout(
        paper_bgcolor='rgba(0,0,0,0)', 
        plot_bgcolor='rgba(0,0,0,0)',
        font=dict(color="white")
    )
    return fig