import dash
from dash import dcc, html, Input, Output, State, ctx
import pandas as pd
import plotly.express as px
from dash.dependencies import ALL

# Load dataset
df = pd.read_csv("data/moths_combined.csv")
# df['country'] = "UK"

# Fields to evaluate completeness
description_fields = [
    'main colors', 'pattern description', 'details colors',
    'antennae description', 'antennae colors',
    'head color', 'abdomen color',
    'forewings description', 'forewing colors',
    'hindwing description', 'hindwing colors'
]

# Compute completeness
def compute_completeness(row):
    filled = sum([pd.notna(row[col]) and str(row[col]).strip() != "" for col in description_fields])
    return filled / len(description_fields)

df["completeness_score"] = df.apply(compute_completeness, axis=1)

def classify(score):
    if score >= 0.8:
        return "High"
    elif score >= 0.6:
        return "Medium"
    else:
        return "Low"

df["completeness_label"] = df["completeness_score"].apply(classify)

# Country-level completeness
def completeness_color(score):
    if score >= 0.8:
        return 'green'
    
    elif score >= 0.6:
        return 'orange'
    else:
        return 'red'

country_completeness = df.groupby('country')['completeness_score'].mean().reset_index()
country_completeness['color'] = country_completeness['completeness_score'].apply(completeness_color)

# Dash app
app = dash.Dash(__name__, suppress_callback_exceptions=True)
server = app.server 
app.title = "Moth Explorer"

app.layout = html.Div([
    html.Div([
        html.H1("🦋 Moth Description Explorer 🦋", style={
            'textAlign': 'center',
            'padding': '20px 0',
            'color': '#1c3b6f'
        }),

        html.P(
            "Explore the morphological completeness of moth species, from order to species level. "
            "Click on the sunburst chart to view detailed color and wing pattern information.",
            style={
                'textAlign': 'center',
                'fontSize': '16px',
                'maxWidth': '900px',
                'margin': '0 auto',
                'paddingBottom': '20px',
                'color': '#333'
            }
        ),

        dcc.Store(id='selected-country', data=None),

        html.Div(id='country-boxes', style={
            'display': 'flex',
            'flexWrap': 'wrap',
            'gap': '10px',
            'justifyContent': 'center',
            'marginBottom': '20px'
        })
    ], style={'fontFamily': '"Segoe UI", Roboto, Inter, sans-serif'}),

    html.Div([
        html.Div([
            dcc.Graph(id='sunburst-plot', config={'scrollZoom': True})
        ], style={'flex': '1', 'padding': '10px'}),

        html.Div(id='info-box', style={
            'flex': '1',
            'padding': '20px',
            'fontSize': '16px',
            'lineHeight': '1.5',
            'overflowY': 'auto'
        })
    ], style={
        'display': 'flex',
        'flexDirection': 'row',
        'flexWrap': 'nowrap',
        'height': '80vh',
        'gap': '20px',
        'fontFamily': '"Segoe UI", Roboto, Inter, sans-serif'
    }),

    html.Div([
        html.Img(src='/assets/images/logo1.png', style={'height': '60px', 'margin': '10px'}),
        html.Img(src='/assets/images/logo2.png', style={'height': '60px', 'margin': '10px'})
    ], style={
        'textAlign': 'center',
        'padding': '20px 0'
    })
])

# Country selection boxes
@app.callback(
    Output('country-boxes', 'children'),
    Input('selected-country', 'data')
)
def render_country_boxes(selected_country):
    boxes = []
    for _, row in country_completeness.iterrows():
        border = '3px solid black' if row['country'] == selected_country else '1px solid #ccc'
        boxes.append(html.Div(
            row['country'],
            id={'type': 'country-box', 'index': row['country']},
            style={
                'backgroundColor': row['color'],
                'padding': '10px 15px',
                'borderRadius': '8px',
                'border': border,
                'cursor': 'pointer',
                'color': 'white',
                'fontWeight': 'bold',
                'textAlign': 'center'
            }
        ))
    boxes.append(html.Div(
        "All Countries",
        id={'type': 'country-box', 'index': 'all'},
        style={
            'backgroundColor': '#999',
            'padding': '10px 15px',
            'borderRadius': '8px',
            'border': '1px solid #ccc',
            'cursor': 'pointer',
            'color': 'white',
            'fontWeight': 'bold',
            'textAlign': 'center'
        }
    ))
    return boxes

# Handle box click
@app.callback(
    Output('selected-country', 'data'),
    Input({'type': 'country-box', 'index': ALL}, 'n_clicks'),
    State({'type': 'country-box', 'index': ALL}, 'id'),
    prevent_initial_call=True
)
def select_country(n_clicks, ids):
    if ctx.triggered_id:
        return ctx.triggered_id['index']
    return dash.no_update

# Sunburst plot
@app.callback(
    Output('sunburst-plot', 'figure'),
    # Input('sunburst-plot', 'clickData'),
    Input('selected-country', 'data')
)
def update_sunburst(selected_country):
    print(f"Selected country: {selected_country}")
    filtered_df = df[df['country'] == selected_country] if selected_country else df
    if selected_country != 'UK' and selected_country != 'Costa Rica':
        print('no UK or Costa Rica selected, showing all data')
        filtered_df = df  # Filter out low completeness
    print(f"Filtered data size: {filtered_df.shape}")
    fig = px.sunburst(
        filtered_df,
        path=['order_name', 'family_name', 'genus_name', 'species_name'],
        values=None,  # Use counts automatically
        color='completeness_score',
        color_continuous_scale=['red', 'orange', 'green'],
        hover_data={'completeness_score': ':.2f'}
    )
    fig.update_traces(insidetextorientation='radial')
    fig.update_layout(margin=dict(t=0, l=0, r=0, b=0))
    return fig

# Info panel with table
@app.callback(
    Output('info-box', 'children'),
    Input('sunburst-plot', 'clickData')
)
def display_info(clickData):
    if clickData:
        label = clickData['points'][0]['label']
        matched = df[df['species_name'] == label]
        if not matched.empty:
            row = matched.iloc[0]
            fields = {
                "Main colors": row['main colors'],
                "Pattern": row['pattern description'],
                "Details colors": row['details colors'],
                "Antennae": f"{row['antennae description']} ({row['antennae colors']})",
                "Head color": row['head color'],
                "Abdomen color": row['abdomen color'],
                "Forewings": f"{row['forewings description']} ({row['forewing colors']})",
                "Hindwings": f"{row['hindwing description']} ({row['hindwing colors']})"
            }

            table_rows = [
                html.Tr([
                    html.Th("Field", style={'textAlign': 'left', 'border': '1px solid #ccc', 'padding': '6px', 'backgroundColor': '#f0f0f0'}),
                    html.Th("Value", style={'textAlign': 'left', 'border': '1px solid #ccc', 'padding': '6px', 'backgroundColor': '#f0f0f0'})
                ])
            ]
            for key, value in fields.items():
                table_rows.append(
                    html.Tr([
                        html.Td(key, style={'border': '1px solid #ccc', 'padding': '6px'}),
                        html.Td(value if value and str(value).strip() else "N/A", style={'border': '1px solid #ccc', 'padding': '6px'})
                    ])
                )

            return html.Div([
                html.H2(f"{row['species_name']}", style={'color': '#1c3b6f'}),
                html.P(f"Order: {row['order_name']}"),
                html.P(f"Family: {row['family_name']}"),
                html.P(f"Genus: {row['genus_name']}"),
                html.P(f"Completeness level: {row['completeness_label']}"),
                html.Table(table_rows, style={
                    'borderCollapse': 'collapse',
                    'width': '100%',
                    'marginTop': '20px'
                })
            ])
    return html.Div("Click on a species to view its descriptions.")

if __name__ == '__main__':
    app.run_server(debug=True)