"""pwml Demo Web Application.

A Streamlit app that demonstrates the core pwml use cases with strong
visualization components:

- Train a hierarchical classifier on uploaded CSV data
- Run batch predictions and display results with confidence heatmaps
- Evaluate model accuracy per hierarchy level with bar charts
- Monitor input distribution drift with PSI/JS gauges

Requirements (in addition to pwml):
    pip install streamlit matplotlib seaborn

Run:
    streamlit run examples/webapp/app.py
"""

import io
import warnings
import tempfile
import os

import numpy as np
import pandas as pd
import streamlit as st
import matplotlib.pyplot as plt
import seaborn as sns

# ---------------------------------------------------------------------------
# Page config
# ---------------------------------------------------------------------------

st.set_page_config(
    page_title='pwml Demo',
    page_icon='',
    layout='wide',
)

st.title('pwml - Hierarchical Classification Demo')
st.markdown(
    'This app lets you train, evaluate, and inspect a **hierarchical '
    'classifier** built with the pwml library.  Upload your own CSV or use '
    'the built-in synthetic dataset to get started.'
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


@st.cache_data
def load_synthetic():
    """Return a small synthetic two-level classification DataFrame."""
    rng = np.random.default_rng(0)
    n = 120
    animals = ['mammal'] * 40 + ['bird'] * 40 + ['reptile'] * 40
    species_map = {
        'mammal': ['cat', 'dog'],
        'bird': ['eagle', 'parrot'],
        'reptile': ['snake', 'lizard'],
    }
    species = [rng.choice(species_map[a]) for a in animals]
    return pd.DataFrame({
        'weight': rng.uniform(0.5, 80.0, n),
        'length': rng.uniform(0.1, 3.0, n),
        'habitat': rng.choice(['forest', 'desert', 'ocean'], n),
        'animal': animals,
        'species': species,
    })


def build_model(input_features, output_hierarchy, cache_path):
    from pwml.classifiers.features import InputFeature, OutputFeature
    from pwml.classifiers.hierarchical import HierarchicalClassifierModel

    pwml_input = [InputFeature(name, ftype) for name, ftype in input_features]

    # Build nested OutputFeature chain from list of level names
    root = None
    for level_name in reversed(output_hierarchy):
        root = OutputFeature(level_name, child_feature=root) if root else OutputFeature(level_name)
    # reverse was needed to build from deepest up; rebuild correctly
    root = None
    for level_name in output_hierarchy:
        new = OutputFeature(level_name, child_feature=root)
        root = new
    # Actually build from innermost outward
    root = None
    for level_name in reversed(output_hierarchy):
        if root is None:
            root = OutputFeature(level_name)
        else:
            root = OutputFeature(level_name, child_feature=root)

    return HierarchicalClassifierModel(
        model_name='webapp_model',
        experiment_name='webapp',
        input_features=pwml_input,
        output_feature_hierarchy=root,
        cache_path=cache_path,
    )


def plot_confidence_heatmap(predictions_df, predicted_cols):
    """Return a matplotlib figure: confidence heatmap for predicted columns."""
    conf_cols = [c.replace('_predicted', '_confidence')
                 for c in predicted_cols
                 if c.replace('_predicted', '_confidence') in predictions_df.columns]
    if not conf_cols:
        return None
    conf_data = predictions_df[conf_cols].copy()
    conf_data.columns = [c.replace('_confidence', '') for c in conf_cols]

    fig, ax = plt.subplots(figsize=(max(4, len(conf_cols) * 2), 4))
    sns.heatmap(
        conf_data.head(50).T,
        ax=ax,
        vmin=0, vmax=1,
        cmap='YlGn',
        linewidths=0.3,
        cbar_kws={'label': 'Confidence'},
    )
    ax.set_title('Prediction Confidence (first 50 rows)', fontsize=12)
    ax.set_xlabel('Row index')
    ax.set_ylabel('Hierarchy level')
    plt.tight_layout()
    return fig


def plot_accuracy_bar(metrics):
    """Return a matplotlib figure: accuracy per hierarchy level."""
    fig, ax = plt.subplots(figsize=(max(4, len(metrics) * 1.5), 4))
    levels = list(metrics.keys())
    accs = [metrics[l] for l in levels]
    bars = ax.bar(levels, accs, color=sns.color_palette('Blues_d', len(levels)))
    ax.set_ylim(0, 1.05)
    ax.set_ylabel('Accuracy')
    ax.set_title('Accuracy per Hierarchy Level')
    for bar, acc in zip(bars, accs):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.02,
            '{:.1%}'.format(acc),
            ha='center', va='bottom', fontsize=10)
    plt.tight_layout()
    return fig


def plot_drift_gauges(report):
    """Return a matplotlib figure: PSI bar chart per feature."""
    features = list(report.keys())
    psi_values = [report[f]['psi'] for f in features]
    colors = ['#d62728' if report[f]['drifted'] else '#2ca02c' for f in features]

    fig, ax = plt.subplots(figsize=(max(5, len(features) * 1.5), 4))
    bars = ax.bar(features, psi_values, color=colors)
    ax.axhline(0.1, color='orange', linestyle='--', linewidth=1, label='PSI=0.1 (moderate)')
    ax.axhline(0.2, color='red', linestyle='--', linewidth=1, label='PSI=0.2 (significant)')
    ax.set_ylabel('PSI')
    ax.set_title('Population Stability Index per Feature')
    ax.legend(fontsize=8)
    for bar, psi in zip(bars, psi_values):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.005,
            '{:.3f}'.format(psi),
            ha='center', va='bottom', fontsize=9)
    plt.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Session state
# ---------------------------------------------------------------------------

if 'model' not in st.session_state:
    st.session_state['model'] = None
if 'train_df' not in st.session_state:
    st.session_state['train_df'] = None

# ---------------------------------------------------------------------------
# Sidebar - Data source
# ---------------------------------------------------------------------------

st.sidebar.header('1. Data')
data_source = st.sidebar.radio('Data source', ['Synthetic dataset', 'Upload CSV'])

if data_source == 'Upload CSV':
    uploaded = st.sidebar.file_uploader('Upload CSV', type='csv')
    if uploaded:
        df_raw = pd.read_csv(uploaded)
    else:
        df_raw = None
else:
    df_raw = load_synthetic()

if df_raw is not None:
    st.sidebar.success('{:,} rows loaded'.format(len(df_raw)))

# ---------------------------------------------------------------------------
# Sidebar - Feature configuration
# ---------------------------------------------------------------------------

st.sidebar.header('2. Features')

if df_raw is not None:
    all_cols = list(df_raw.columns)

    output_cols = st.sidebar.multiselect(
        'Output (hierarchy levels, top to bottom)',
        options=all_cols,
        default=['animal', 'species'] if 'animal' in all_cols else all_cols[-2:],
    )

    remaining = [c for c in all_cols if c not in output_cols]

    numeric_cols = st.sidebar.multiselect(
        'Numeric input features',
        options=remaining,
        default=[c for c in ['weight', 'length'] if c in remaining],
    )
    category_cols = st.sidebar.multiselect(
        'Category input features',
        options=[c for c in remaining if c not in numeric_cols],
        default=[c for c in ['habitat'] if c in remaining],
    )

    input_features = (
        [(c, 'numeric') for c in numeric_cols]
        + [(c, 'category') for c in category_cols]
    )

# ---------------------------------------------------------------------------
# Sidebar - Train
# ---------------------------------------------------------------------------

st.sidebar.header('3. Train')
train_btn = st.sidebar.button('Train model', disabled=(df_raw is None))

if train_btn and df_raw is not None and output_cols and input_features:
    with st.spinner('Training...'):
        tmpdir = tempfile.mkdtemp(prefix='pwml_webapp_')
        try:
            model = build_model(input_features, output_cols, cache_path=tmpdir)
            model.load_from_dataframe(df_raw)
            st.session_state['model'] = model
            st.session_state['train_df'] = df_raw.copy()
            st.sidebar.success('Model trained!')
        except Exception as exc:
            st.sidebar.error('Training failed: {}'.format(exc))

# ---------------------------------------------------------------------------
# Main area - tabs
# ---------------------------------------------------------------------------

tab_data, tab_predict, tab_evaluate, tab_drift = st.tabs([
    'Data Explorer', 'Predictions', 'Evaluation', 'Drift Monitor'
])

# --- Tab 1: Data Explorer ---
with tab_data:
    if df_raw is not None:
        st.subheader('Dataset preview')
        st.dataframe(df_raw.head(100), use_container_width=True)
        st.subheader('Numeric distributions')
        num_cols_in_df = df_raw.select_dtypes(include='number').columns.tolist()
        if num_cols_in_df:
            fig, axes = plt.subplots(1, len(num_cols_in_df),
                                     figsize=(4 * len(num_cols_in_df), 3))
            if len(num_cols_in_df) == 1:
                axes = [axes]
            for ax, col in zip(axes, num_cols_in_df):
                ax.hist(df_raw[col].dropna(), bins=20, color='steelblue', edgecolor='white')
                ax.set_title(col, fontsize=10)
                ax.set_yticks([])
            plt.tight_layout()
            st.pyplot(fig)
            plt.close(fig)
    else:
        st.info('Load a dataset using the sidebar.')

# --- Tab 2: Predictions ---
with tab_predict:
    model = st.session_state.get('model')
    train_df = st.session_state.get('train_df')

    if model is None:
        st.info('Train a model first (sidebar).')
    else:
        st.subheader('Batch predictions')
        preds_df = model.predict_dataframe(train_df)
        st.dataframe(preds_df.head(100), use_container_width=True)

        predicted_cols = [c for c in preds_df.columns if c.endswith('_predicted')]
        fig_heat = plot_confidence_heatmap(preds_df, predicted_cols)
        if fig_heat:
            st.pyplot(fig_heat)
            plt.close(fig_heat)

        csv_buf = io.StringIO()
        preds_df.to_csv(csv_buf, index=False)
        st.download_button(
            'Download predictions as CSV',
            data=csv_buf.getvalue(),
            file_name='predictions.csv',
            mime='text/csv',
        )

# --- Tab 3: Evaluation ---
with tab_evaluate:
    model = st.session_state.get('model')
    train_df = st.session_state.get('train_df')

    if model is None:
        st.info('Train a model first (sidebar).')
    else:
        st.subheader('Accuracy per hierarchy level (training set)')
        with st.spinner('Evaluating...'):
            metrics, _ = model.evaluate(train_df)
        st.metric('Overall levels evaluated', len(metrics))
        cols_met = st.columns(len(metrics))
        for col_ui, (level, acc) in zip(cols_met, metrics.items()):
            col_ui.metric(level, '{:.1%}'.format(acc))

        fig_acc = plot_accuracy_bar(metrics)
        st.pyplot(fig_acc)
        plt.close(fig_acc)

        st.subheader('Per-node inference latency profile')
        _, latency = model.predict_dataframe(train_df.head(20), profile=True)
        if latency:
            lat_df = pd.DataFrame([
                {'node': k, 'total_s': v, 'avg_ms': v / 20 * 1000}
                for k, v in sorted(latency.items(), key=lambda x: -x[1])
            ])
            st.dataframe(lat_df, use_container_width=True)

            fig_lat, ax_lat = plt.subplots(figsize=(max(5, len(lat_df) * 1.5), 3))
            ax_lat.bar(lat_df['node'], lat_df['avg_ms'], color='mediumpurple')
            ax_lat.set_ylabel('Avg latency (ms / row)')
            ax_lat.set_title('Per-node inference latency (20 rows)')
            plt.xticks(rotation=30, ha='right')
            plt.tight_layout()
            st.pyplot(fig_lat)
            plt.close(fig_lat)

# --- Tab 4: Drift Monitor ---
with tab_drift:
    model = st.session_state.get('model')
    train_df = st.session_state.get('train_df')

    if model is None or train_df is None:
        st.info('Train a model first (sidebar).')
    else:
        from pwml.utilities.driftmonitor import DriftMonitor

        st.subheader('Drift Monitor')
        st.markdown(
            'Compare the feature distributions of an uploaded "live" batch '
            'against the training data.  Upload a CSV that shares the same '
            'input columns as the training set.'
        )

        live_file = st.file_uploader('Upload live batch CSV', type='csv', key='drift_upload')
        psi_thresh = st.slider('PSI drift threshold', 0.05, 0.5, 0.2, 0.05)

        if live_file is not None:
            live_df = pd.read_csv(live_file)
            monitor = DriftMonitor(
                numeric_features=[c for c, t in input_features if t == 'numeric'],
                category_features=[c for c, t in input_features if t == 'category'],
                psi_threshold=psi_thresh,
            )
            monitor.fit(train_df)
            report = monitor.compute(live_df)

            st.subheader('Drift report')
            report_df = pd.DataFrame([
                {'feature': f,
                 'psi': v['psi'],
                 'js_divergence': v['js'],
                 'drifted': v['drifted']}
                for f, v in report.items()
            ])
            st.dataframe(report_df, use_container_width=True)

            if report:
                fig_drift = plot_drift_gauges(report)
                st.pyplot(fig_drift)
                plt.close(fig_drift)

                drifted = [f for f, v in report.items() if v['drifted']]
                if drifted:
                    st.warning(
                        'Drift detected in: **{}**'.format(', '.join(drifted))
                    )
                else:
                    st.success('No significant drift detected (PSI < {:.2f}).'.format(psi_thresh))
        else:
            st.info(
                'Upload a CSV to compare against training data.  '
                'Tip: try modifying some column distributions to see drift detected.'
            )
