src/visualization.py

"""
Visualization module for RehabWatch.
Creates maps and charts using Folium and Plotly.
"""

import numpy as np
import xarray as xr
import folium
from folium import plugins
from folium.raster_layers import ImageOverlay
import plotly.graph_objects as go
import streamlit as st
from typing import Dict, Any, List, Optional, Tuple
from matplotlib.colors import LinearSegmentedColormap
import base64
from io import BytesIO
from PIL import Image


# NDVI color palette (brown to green)
NDVI_COLORS = ['#8B4513', '#D2B48C', '#FFFF00', '#90EE90', '#228B22', '#006400']

# Change color palette (red-white-green diverging)
CHANGE_COLORS = ['#B71C1C', '#EF9A9A', '#FFFFFF', '#A5D6A7', '#1B5E20']


def array_to_colored_image(
    data: np.ndarray,
    colors: List[str],
    vmin: float,
    vmax: float
) -> np.ndarray:
    """
    Convert a 2D array to a colored RGBA image.

    Args:
        data: 2D numpy array
        colors: List of hex color strings for colormap
        vmin: Minimum value for normalization
        vmax: Maximum value for normalization

    Returns:
        RGBA numpy array (H, W, 4) with values 0-255
    """
    cmap = LinearSegmentedColormap.from_list('custom', colors)

    # Normalize data
    normalized = (data - vmin) / (vmax - vmin)
    normalized = np.clip(normalized, 0, 1)

    # Handle NaN values
    mask = np.isnan(data)

    # Apply colormap
    rgba = cmap(normalized)
    rgba = (rgba * 255).astype(np.uint8)

    # Set NaN pixels to transparent
    rgba[mask, 3] = 0

    return rgba


def create_image_overlay(
    data: xr.DataArray,
    colors: List[str],
    vmin: float,
    vmax: float,
    bounds: List[List[float]]
) -> str:
    """
    Create a base64-encoded PNG image for Folium overlay.

    Args:
        data: xarray DataArray
        colors: Color palette
        vmin: Min value for normalization
        vmax: Max value for normalization
        bounds: [[south, west], [north, east]]

    Returns:
        Base64 encoded PNG string
    """
    # Get the 2D array
    arr = data.values
    if arr.ndim > 2:
        arr = arr.squeeze()

    # Create colored image
    rgba = array_to_colored_image(arr, colors, vmin, vmax)

    # Flip vertically for correct orientation
    rgba = np.flipud(rgba)

    # Convert to PNG
    img = Image.fromarray(rgba, mode='RGBA')
    buffer = BytesIO()
    img.save(buffer, format='PNG')
    buffer.seek(0)

    # Encode to base64
    img_base64 = base64.b64encode(buffer.getvalue()).decode()
    return f"data:image/png;base64,{img_base64}"


def create_comparison_map(
    bbox: Tuple[float, float, float, float],
    ndvi_before: xr.DataArray,
    ndvi_after: xr.DataArray,
    ndvi_change: xr.DataArray,
    center_coords: Tuple[float, float],
    zoom: int = 12
) -> folium.Map:
    """
    Create an interactive comparison map with multiple layers.

    Args:
        bbox: Bounding box (min_lon, min_lat, max_lon, max_lat)
        ndvi_before: NDVI xarray at start date
        ndvi_after: NDVI xarray at end date
        ndvi_change: NDVI change xarray
        center_coords: Map center (lat, lon)
        zoom: Initial zoom level

    Returns:
        Folium Map object with all layers
    """
    # Create base map
    m = folium.Map(
        location=center_coords,
        zoom_start=zoom,
        tiles=None
    )

    # Add satellite basemap
    folium.TileLayer(
        tiles='https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}',
        attr='Esri',
        name='Satellite Imagery',
        overlay=False
    ).add_to(m)

    # Add OpenStreetMap as alternative
    folium.TileLayer(
        tiles='openstreetmap',
        name='OpenStreetMap',
        overlay=False
    ).add_to(m)

    # Calculate bounds for image overlay
    min_lon, min_lat, max_lon, max_lat = bbox
    bounds = [[min_lat, min_lon], [max_lat, max_lon]]

    # Add NDVI Before layer
    try:
        ndvi_before_img = create_image_overlay(
            ndvi_before, NDVI_COLORS, -0.1, 0.8, bounds
        )
        ImageOverlay(
            image=ndvi_before_img,
            bounds=bounds,
            opacity=0.7,
            name='NDVI Before',
            show=False
        ).add_to(m)
    except Exception as e:
        print(f"Error adding NDVI Before layer: {e}")

    # Add NDVI After layer
    try:
        ndvi_after_img = create_image_overlay(
            ndvi_after, NDVI_COLORS, -0.1, 0.8, bounds
        )
        ImageOverlay(
            image=ndvi_after_img,
            bounds=bounds,
            opacity=0.7,
            name='NDVI After',
            show=False
        ).add_to(m)
    except Exception as e:
        print(f"Error adding NDVI After layer: {e}")

    # Add Change Map layer (shown by default)
    try:
        change_img = create_image_overlay(
            ndvi_change, CHANGE_COLORS, -0.3, 0.3, bounds
        )
        ImageOverlay(
            image=change_img,
            bounds=bounds,
            opacity=0.7,
            name='Vegetation Change',
            show=True
        ).add_to(m)
    except Exception as e:
        print(f"Error adding Change layer: {e}")

    # Add tenement boundary
    boundary_coords = [
        [min_lat, min_lon],
        [min_lat, max_lon],
        [max_lat, max_lon],
        [max_lat, min_lon],
        [min_lat, min_lon]
    ]
    folium.PolyLine(
        locations=boundary_coords,
        color='#000000',
        weight=3,
        fill=False,
        popup='Analysis Boundary'
    ).add_to(m)

    # Add layer control
    folium.LayerControl(position='topright').add_to(m)

    # Add legends
    _add_legends(m)

    return m


def _add_legends(m: folium.Map) -> None:
    """Add color legends to the map."""
    legend_html = '''
    <div style="position: fixed; bottom: 50px; left: 50px; z-index: 1000;
                background-color: white; padding: 10px; border-radius: 5px;
                border: 2px solid grey; font-size: 12px; max-width: 150px;">
        <p style="margin: 0 0 5px 0; font-weight: bold;">NDVI Scale</p>
        <div style="background: linear-gradient(to right, #8B4513, #D2B48C, #FFFF00, #90EE90, #228B22, #006400);
                    width: 100%; height: 15px; border-radius: 3px;"></div>
        <div style="display: flex; justify-content: space-between;">
            <span>-0.1</span><span>0.8</span>
        </div>
        <hr style="margin: 8px 0;">
        <p style="margin: 0 0 5px 0; font-weight: bold;">Change</p>
        <div style="background: linear-gradient(to right, #B71C1C, #EF9A9A, #FFFFFF, #A5D6A7, #1B5E20);
                    width: 100%; height: 15px; border-radius: 3px;"></div>
        <div style="display: flex; justify-content: space-between;">
            <span style="color: #B71C1C;">-0.3</span>
            <span style="color: #1B5E20;">+0.3</span>
        </div>
        <p style="margin: 5px 0 0 0; font-size: 10px; text-align: center;">
            Red=Decline | Green=Growth
        </p>
    </div>
    '''
    m.get_root().html.add_child(folium.Element(legend_html))


def create_simple_map(
    center_coords: Tuple[float, float],
    zoom: int = 10,
    bbox: Optional[Tuple[float, float, float, float]] = None
) -> folium.Map:
    """
    Create a simple map for location preview.

    Args:
        center_coords: Map center (lat, lon)
        zoom: Zoom level
        bbox: Optional bounding box to display

    Returns:
        Folium Map object
    """
    m = folium.Map(location=center_coords, zoom_start=zoom)

    # Add satellite imagery
    folium.TileLayer(
        tiles='https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}',
        attr='Esri',
        name='Satellite',
        overlay=False
    ).add_to(m)

    if bbox is not None:
        min_lon, min_lat, max_lon, max_lat = bbox
        boundary_coords = [
            [min_lat, min_lon],
            [min_lat, max_lon],
            [max_lat, max_lon],
            [max_lat, min_lon],
            [min_lat, min_lon]
        ]
        folium.Polygon(
            locations=boundary_coords,
            color='#1B5E20',
            weight=3,
            fill=True,
            fillColor='#2E7D32',
            fillOpacity=0.2,
            popup='Analysis Area'
        ).add_to(m)

    folium.LayerControl().add_to(m)
    return m


def create_time_series_chart(
    timeseries_data: List[Dict[str, Any]],
    title: str = "NDVI Time Series"
) -> go.Figure:
    """
    Create an interactive NDVI time series chart.

    Args:
        timeseries_data: List of dicts with 'date' and 'ndvi' keys
        title: Chart title

    Returns:
        Plotly Figure object
    """
    if not timeseries_data:
        fig = go.Figure()
        fig.add_annotation(
            text="No time series data available",
            xref="paper", yref="paper",
            x=0.5, y=0.5, showarrow=False,
            font=dict(size=16)
        )
        return fig

    dates = [d['date'] for d in timeseries_data]
    ndvi_values = [d['ndvi'] for d in timeseries_data]

    fig = go.Figure()

    # Add NDVI line
    fig.add_trace(go.Scatter(
        x=dates,
        y=ndvi_values,
        mode='lines+markers',
        name='NDVI',
        line=dict(color='#2E7D32', width=2),
        marker=dict(size=6),
        hovertemplate='Date: %{x}<br>NDVI: %{y:.3f}<extra></extra>'
    ))

    # Add reference lines
    fig.add_hline(y=0.6, line_dash="dash", line_color="#4CAF50",
                  annotation_text="Healthy Vegetation", annotation_position="right")
    fig.add_hline(y=0.2, line_dash="dash", line_color="#FF9800",
                  annotation_text="Sparse Vegetation", annotation_position="right")

    fig.update_layout(
        title=dict(text=title, font=dict(size=18)),
        xaxis_title="Date",
        yaxis_title="NDVI",
        yaxis=dict(range=[0, 1]),
        template="plotly_white",
        hovermode="x unified",
        height=400,
        margin=dict(l=60, r=40, t=60, b=60)
    )

    return fig


def create_stats_display(stats: Dict[str, float], rehab_score: int) -> None:
    """
    Display statistics using Streamlit components.

    Args:
        stats: Statistics dictionary
        rehab_score: Rehabilitation score (0-100)
    """
    # Rehabilitation Score with large display
    st.markdown("### Rehabilitation Score")

    score_color = _get_score_color(rehab_score)
    st.markdown(f"""
    <div style="text-align: center; padding: 20px; background-color: {score_color}20;
                border-radius: 10px; margin-bottom: 20px;">
        <span style="font-size: 72px; font-weight: bold; color: {score_color};">
            {rehab_score}
        </span>
        <span style="font-size: 24px; color: {score_color};">/100</span>
    </div>
    """, unsafe_allow_html=True)

    # Progress bar
    st.progress(rehab_score / 100)

    # Key Metrics in columns
    # Logic: arrow direction = numeric delta; color = good/bad for nature
    st.markdown("### Key Metrics")
    col1, col2, col3 = st.columns(3)

    ndvi_change = stats.get('ndvi_change_mean', 0)
    percent_change = stats.get('percent_change', 0)

    with col1:
        st.metric(
            label="NDVI Before",
            value=f"{stats['ndvi_before_mean']:.3f}",
            help="Normalized Difference Vegetation Index: measures vegetation health (-1 to 1)"
        )
        st.metric(
            label="Area Improved",
            value=f"{stats['area_improved_ha']:.1f} ha",
            delta=f"+{stats['percent_improved']:.1f}%",
            delta_color="normal"  # improvement is always good
        )

    with col2:
        # NDVI: increase = good (green), decrease = bad (red)
        st.metric(
            label="NDVI After",
            value=f"{stats['ndvi_after_mean']:.3f}",
            delta=f"{ndvi_change:+.3f}" if ndvi_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Current vegetation index value"
        )
        st.metric(
            label="Area Degraded",
            value=f"{stats['area_degraded_ha']:.1f} ha",
            delta=f"-{stats['percent_degraded']:.1f}%",
            delta_color="inverse"  # degradation showing as negative is correct
        )

    with col3:
        # Vegetation Change: increase = good (green), decrease = bad (red)
        st.metric(
            label="Vegetation Change",
            value=f"{percent_change:+.1f}%",
            delta=f"{percent_change:+.1f}%" if percent_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Percentage change in vegetation cover"
        )
        st.metric(
            label="Total Area",
            value=f"{stats['total_area_ha']:.1f} ha"
        )


def _get_score_color(score: int) -> str:
    """Get color based on rehabilitation score."""
    if score >= 80:
        return "#1B5E20"
    elif score >= 60:
        return "#4CAF50"
    elif score >= 40:
        return "#FF9800"
    elif score >= 20:
        return "#F57C00"
    else:
        return "#B71C1C"


def create_area_breakdown_chart(stats: Dict[str, float]) -> go.Figure:
    """
    Create a pie chart showing area breakdown.

    Args:
        stats: Statistics dictionary with area values

    Returns:
        Plotly Figure object
    """
    labels = ['Improved', 'Stable', 'Degraded']
    values = [
        stats['area_improved_ha'],
        stats['area_stable_ha'],
        stats['area_degraded_ha']
    ]
    colors = ['#4CAF50', '#FFC107', '#F44336']

    fig = go.Figure(data=[go.Pie(
        labels=labels,
        values=values,
        marker_colors=colors,
        hole=0.4,
        textinfo='label+percent',
        hovertemplate='%{label}<br>%{value:.1f} ha<br>%{percent}<extra></extra>'
    )])

    fig.update_layout(
        title="Area Breakdown",
        annotations=[dict(text='Area', x=0.5, y=0.5, font_size=16, showarrow=False)],
        showlegend=True,
        height=350
    )

    return fig


def create_ndvi_comparison_chart(stats: Dict[str, float]) -> go.Figure:
    """
    Create a bar chart comparing before/after NDVI.

    Args:
        stats: Statistics dictionary

    Returns:
        Plotly Figure object
    """
    fig = go.Figure()

    fig.add_trace(go.Bar(
        x=['Before', 'After'],
        y=[stats['ndvi_before_mean'], stats['ndvi_after_mean']],
        marker_color=['#8B4513', '#228B22'],
        text=[f"{stats['ndvi_before_mean']:.3f}", f"{stats['ndvi_after_mean']:.3f}"],
        textposition='outside'
    ))

    fig.update_layout(
        title="NDVI Comparison",
        yaxis_title="NDVI",
        yaxis=dict(range=[0, max(stats['ndvi_after_mean'], stats['ndvi_before_mean']) * 1.3]),
        template="plotly_white",
        height=350
    )

    return fig


def create_statistics_table(stats: Dict[str, float]) -> None:
    """
    Display full statistics as a formatted table.

    Args:
        stats: Statistics dictionary
    """
    import pandas as pd

    data = {
        'Metric': [
            'NDVI Before (mean)',
            'NDVI After (mean)',
            'NDVI Change (mean)',
            'NDVI Change (std dev)',
            'Relative Change',
            'Area Improved',
            'Area Stable',
            'Area Degraded',
            'Total Area',
            '% Improved',
            '% Stable',
            '% Degraded'
        ],
        'Value': [
            f"{stats['ndvi_before_mean']:.4f}",
            f"{stats['ndvi_after_mean']:.4f}",
            f"{stats['ndvi_change_mean']:.4f}",
            f"{stats['ndvi_change_std']:.4f}",
            f"{stats['percent_change']:.2f}%",
            f"{stats['area_improved_ha']:.2f} ha",
            f"{stats['area_stable_ha']:.2f} ha",
            f"{stats['area_degraded_ha']:.2f} ha",
            f"{stats['total_area_ha']:.2f} ha",
            f"{stats['percent_improved']:.2f}%",
            f"{stats['percent_stable']:.2f}%",
            f"{stats['percent_degraded']:.2f}%"
        ],
        'Description': [
            'Mean vegetation index at analysis start',
            'Mean vegetation index at analysis end',
            'Average change in vegetation index',
            'Variation in vegetation change',
            'Percentage change in mean NDVI',
            'Area with NDVI increase > 0.05',
            'Area with NDVI change between -0.05 and 0.05',
            'Area with NDVI decrease > 0.05',
            'Total analyzed area',
            'Percentage of area showing improvement',
            'Percentage of area remaining stable',
            'Percentage of area showing degradation'
        ]
    }

    df = pd.DataFrame(data)
    st.dataframe(df, use_container_width=True, hide_index=True)


# =============================================================================
# NEW EXTENDED VISUALIZATIONS
# =============================================================================

# Color palettes for different indices
BSI_COLORS = ['#228B22', '#90EE90', '#FFFF00', '#D2B48C', '#8B4513']  # Green to brown
WATER_COLORS = ['#8B4513', '#D2B48C', '#87CEEB', '#4169E1', '#000080']  # Brown to blue
MOISTURE_COLORS = ['#B71C1C', '#FF5722', '#FFEB3B', '#8BC34A', '#1B5E20']  # Dry to wet
SLOPE_COLORS = ['#1B5E20', '#4CAF50', '#FFEB3B', '#FF9800', '#B71C1C']  # Flat to steep
EROSION_COLORS = ['#1B5E20', '#4CAF50', '#FFEB3B', '#FF5722', '#B71C1C']  # Low to high risk

# Land cover color mapping
LULC_COLORS = {
    1: '#0000FF',   # Water - Blue
    2: '#228B22',   # Trees - Forest Green
    4: '#006400',   # Flooded Vegetation - Dark Green
    5: '#FFD700',   # Crops - Gold
    7: '#808080',   # Built Area - Gray
    8: '#D2691E',   # Bare Ground - Chocolate
    9: '#FFFFFF',   # Snow/Ice - White
    10: '#C0C0C0',  # Clouds - Silver
    11: '#9ACD32'   # Rangeland - Yellow Green
}


def create_multi_index_map(
    bbox: Tuple[float, float, float, float],
    indices_after: Dict[str, xr.DataArray],
    index_changes: Dict[str, xr.DataArray],
    center_coords: Tuple[float, float],
    zoom: int = 12
) -> folium.Map:
    """
    Create an interactive map with multiple index layers.
    """
    m = folium.Map(location=center_coords, zoom_start=zoom, tiles=None)

    # Add basemaps
    folium.TileLayer(
        tiles='https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}',
        attr='Esri', name='Satellite', overlay=False
    ).add_to(m)

    folium.TileLayer(tiles='openstreetmap', name='OpenStreetMap', overlay=False).add_to(m)

    min_lon, min_lat, max_lon, max_lat = bbox
    bounds = [[min_lat, min_lon], [max_lat, max_lon]]

    # Index configurations: (data, colors, vmin, vmax, name)
    index_configs = [
        ('ndvi', NDVI_COLORS, -0.1, 0.8, 'NDVI'),
        ('savi', NDVI_COLORS, -0.1, 0.8, 'SAVI'),
        ('evi', NDVI_COLORS, -0.1, 0.8, 'EVI'),
        ('bsi', BSI_COLORS, -0.5, 0.5, 'Bare Soil Index'),
        ('ndwi', WATER_COLORS, -0.5, 0.5, 'Water Index (NDWI)'),
        ('ndmi', MOISTURE_COLORS, -0.5, 0.5, 'Moisture Index (NDMI)'),
    ]

    # Add current state layers
    for idx_key, colors, vmin, vmax, name in index_configs:
        if idx_key in indices_after:
            try:
                img = create_image_overlay(indices_after[idx_key], colors, vmin, vmax, bounds)
                ImageOverlay(
                    image=img, bounds=bounds, opacity=0.7,
                    name=f'{name} (Current)', show=(idx_key == 'ndvi')
                ).add_to(m)
            except Exception:
                pass

    # Add change layers
    for idx_key, _, _, _, name in index_configs:
        if idx_key in index_changes:
            try:
                img = create_image_overlay(index_changes[idx_key], CHANGE_COLORS, -0.3, 0.3, bounds)
                ImageOverlay(
                    image=img, bounds=bounds, opacity=0.7,
                    name=f'{name} Change', show=False
                ).add_to(m)
            except Exception:
                pass

    # Add boundary
    boundary_coords = [
        [min_lat, min_lon], [min_lat, max_lon],
        [max_lat, max_lon], [max_lat, min_lon], [min_lat, min_lon]
    ]
    folium.PolyLine(locations=boundary_coords, color='#000000', weight=3).add_to(m)

    folium.LayerControl(position='topright').add_to(m)
    _add_legends(m)

    return m


def create_terrain_map(
    bbox: Tuple[float, float, float, float],
    slope: xr.DataArray,
    aspect: Optional[xr.DataArray],
    erosion_risk: Optional[xr.DataArray],
    center_coords: Tuple[float, float],
    zoom: int = 12
) -> folium.Map:
    """
    Create an interactive terrain analysis map.
    """
    m = folium.Map(location=center_coords, zoom_start=zoom, tiles=None)

    folium.TileLayer(
        tiles='https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}',
        attr='Esri', name='Satellite', overlay=False
    ).add_to(m)

    min_lon, min_lat, max_lon, max_lat = bbox
    bounds = [[min_lat, min_lon], [max_lat, max_lon]]

    # Add slope layer
    try:
        slope_img = create_image_overlay(slope, SLOPE_COLORS, 0, 45, bounds)
        ImageOverlay(
            image=slope_img, bounds=bounds, opacity=0.7,
            name='Slope (degrees)', show=True
        ).add_to(m)
    except Exception:
        pass

    # Add erosion risk layer
    if erosion_risk is not None:
        try:
            erosion_img = create_image_overlay(erosion_risk, EROSION_COLORS, 0, 1, bounds)
            ImageOverlay(
                image=erosion_img, bounds=bounds, opacity=0.7,
                name='Erosion Risk', show=False
            ).add_to(m)
        except Exception:
            pass

    folium.LayerControl(position='topright').add_to(m)

    return m


def create_land_cover_map(
    bbox: Tuple[float, float, float, float],
    lulc: xr.DataArray,
    center_coords: Tuple[float, float],
    zoom: int = 12,
    year: int = 2023
) -> folium.Map:
    """
    Create a land cover classification map.
    """
    from matplotlib.colors import ListedColormap

    m = folium.Map(location=center_coords, zoom_start=zoom, tiles=None)

    folium.TileLayer(
        tiles='https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}',
        attr='Esri', name='Satellite', overlay=False
    ).add_to(m)

    min_lon, min_lat, max_lon, max_lat = bbox
    bounds = [[min_lat, min_lon], [max_lat, max_lon]]

    # Create categorical colormap
    try:
        arr = lulc.values.squeeze()
        rgba = np.zeros((*arr.shape, 4), dtype=np.uint8)

        for class_id, color in LULC_COLORS.items():
            mask = arr == class_id
            r = int(color[1:3], 16)
            g = int(color[3:5], 16)
            b = int(color[5:7], 16)
            rgba[mask] = [r, g, b, 200]

        rgba = np.flipud(rgba)
        img = Image.fromarray(rgba, mode='RGBA')
        buffer = BytesIO()
        img.save(buffer, format='PNG')
        buffer.seek(0)
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        img_url = f"data:image/png;base64,{img_base64}"

        ImageOverlay(
            image=img_url, bounds=bounds, opacity=0.7,
            name=f'Land Cover {year}', show=True
        ).add_to(m)
    except Exception:
        pass

    folium.LayerControl(position='topright').add_to(m)

    return m


def create_multi_index_chart(stats: Dict[str, float]) -> go.Figure:
    """
    Create a grouped bar chart comparing all indices before/after.
    """
    indices = ['NDVI', 'SAVI', 'EVI', 'NDWI', 'NDMI', 'BSI']
    before_values = []
    after_values = []

    for idx in ['ndvi', 'savi', 'evi', 'ndwi', 'ndmi', 'bsi']:
        before_values.append(stats.get(f'{idx}_before_mean', 0))
        after_values.append(stats.get(f'{idx}_after_mean', 0))

    fig = go.Figure()

    fig.add_trace(go.Bar(
        name='Before', x=indices, y=before_values,
        marker_color='#8B4513', text=[f'{v:.3f}' for v in before_values],
        textposition='outside'
    ))

    fig.add_trace(go.Bar(
        name='After', x=indices, y=after_values,
        marker_color='#228B22', text=[f'{v:.3f}' for v in after_values],
        textposition='outside'
    ))

    fig.update_layout(
        title='Multi-Index Comparison',
        barmode='group',
        yaxis_title='Index Value',
        template='plotly_white',
        height=400,
        legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99)
    )

    return fig


def create_terrain_stats_chart(terrain_stats: Dict[str, float]) -> go.Figure:
    """
    Create a chart showing terrain slope distribution.
    """
    labels = ['Flat (<5°)', 'Gentle (5-15°)', 'Moderate (15-30°)', 'Steep (>30°)']
    values = [
        terrain_stats.get('percent_flat', 0),
        terrain_stats.get('percent_gentle', 0),
        terrain_stats.get('percent_moderate', 0),
        terrain_stats.get('percent_steep', 0)
    ]
    colors = ['#1B5E20', '#4CAF50', '#FF9800', '#B71C1C']

    fig = go.Figure(data=[go.Pie(
        labels=labels, values=values, marker_colors=colors,
        hole=0.4, textinfo='label+percent'
    )])

    fig.update_layout(
        title='Slope Distribution',
        height=350
    )

    return fig


def create_land_cover_chart(land_cover_stats: Dict[str, Any]) -> go.Figure:
    """
    Create a grouped bar chart showing land cover change.
    """
    if 'class_changes' not in land_cover_stats:
        return go.Figure()

    changes = land_cover_stats['class_changes']
    classes = list(changes.keys())
    before = [changes[c].get('before', 0) for c in classes]
    after = [changes[c].get('after', 0) for c in classes]

    # Convert to percentages
    total_before = sum(before) or 1
    total_after = sum(after) or 1
    before_pct = [b / total_before * 100 for b in before]
    after_pct = [a / total_after * 100 for a in after]

    fig = go.Figure()

    fig.add_trace(go.Bar(
        name=f"Year {land_cover_stats.get('year_before', 'Before')}",
        x=classes, y=before_pct, marker_color='#8B4513'
    ))

    fig.add_trace(go.Bar(
        name=f"Year {land_cover_stats.get('year_after', 'After')}",
        x=classes, y=after_pct, marker_color='#228B22'
    ))

    fig.update_layout(
        title='Land Cover Change',
        barmode='group',
        yaxis_title='Percentage (%)',
        template='plotly_white',
        height=400
    )

    return fig


def create_vegetation_health_chart(stats: Dict[str, float]) -> go.Figure:
    """
    Create a chart showing vegetation health distribution.
    """
    labels = ['Sparse (0-0.2)', 'Low (0.2-0.4)', 'Moderate (0.4-0.6)', 'Dense (>0.6)']
    values = [
        stats.get('percent_sparse_veg', 0),
        stats.get('percent_low_veg', 0),
        stats.get('percent_moderate_veg', 0),
        stats.get('percent_dense_veg', 0)
    ]
    colors = ['#D2B48C', '#90EE90', '#228B22', '#006400']

    fig = go.Figure(data=[go.Pie(
        labels=labels, values=values, marker_colors=colors,
        hole=0.4, textinfo='label+percent'
    )])

    fig.update_layout(
        title='Vegetation Health Distribution',
        height=350
    )

    return fig


def create_environmental_indicators_chart(stats: Dict[str, float]) -> go.Figure:
    """
    Create a radar chart showing environmental indicators.
    """
    categories = ['Vegetation', 'Moisture', 'Soil Stability', 'Water Presence', 'Dense Veg']

    # Normalize values to 0-100 scale
    values = [
        min(100, stats.get('ndvi_after_mean', 0) * 100 / 0.6),  # NDVI
        max(0, 100 - stats.get('percent_moisture_stressed', 50)),  # Moisture health
        max(0, 100 - stats.get('percent_bare_soil', 50)),  # Soil stability
        min(100, stats.get('percent_water', 0) * 10),  # Water presence
        stats.get('percent_dense_veg', 0)  # Dense vegetation
    ]

    fig = go.Figure()

    fig.add_trace(go.Scatterpolar(
        r=values + [values[0]],  # Close the polygon
        theta=categories + [categories[0]],
        fill='toself',
        fillcolor='rgba(46, 125, 50, 0.3)',
        line=dict(color='#2E7D32', width=2),
        name='Current State'
    ))

    fig.update_layout(
        polar=dict(
            radialaxis=dict(visible=True, range=[0, 100])
        ),
        title='Environmental Health Indicators',
        height=400,
        showlegend=False
    )

    return fig


def create_comprehensive_stats_display(
    stats: Dict[str, float],
    rehab_score: int,
    terrain_stats: Optional[Dict] = None,
    land_cover_stats: Optional[Dict] = None
) -> None:
    """
    Display comprehensive statistics with all new metrics.
    """
    # Rehabilitation Score
    st.markdown("### Rehabilitation Score")
    score_color = _get_score_color(rehab_score)
    st.markdown(f"""
    <div style="text-align: center; padding: 20px; background-color: {score_color}20;
                border-radius: 10px; margin-bottom: 20px;">
        <span style="font-size: 72px; font-weight: bold; color: {score_color};">
            {rehab_score}
        </span>
        <span style="font-size: 24px; color: {score_color};">/100</span>
    </div>
    """, unsafe_allow_html=True)
    st.progress(rehab_score / 100)

    # Primary Metrics with tooltips
    # Logic:
    # - Arrow direction: based on numeric delta (positive=up, negative=down)
    # - Color: "normal" = green for increase (good), red for decrease (bad)
    #          "inverse" = red for increase (bad), green for decrease (good)
    st.markdown("### Key Metrics")
    col1, col2, col3, col4 = st.columns(4)

    # Get change values for proper arrow direction
    ndvi_change = stats.get('ndvi_change_mean', 0)
    percent_change = stats.get('percent_change', 0)

    with col1:
        # NDVI: increase = good (green), decrease = bad (red)
        st.metric(
            "NDVI",
            f"{stats.get('ndvi_after_mean', 0):.3f}",
            delta=f"{ndvi_change:+.3f}" if ndvi_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Normalized Difference Vegetation Index: measures vegetation health. Values range from -1 to 1, with >0.4 indicating healthy vegetation."
        )

    with col2:
        # Vegetation Change: increase = good (green), decrease = bad (red)
        # Use numeric delta for correct arrow direction
        st.metric(
            "Vegetation Change",
            f"{percent_change:+.1f}%",
            delta=f"{percent_change:+.1f}%" if percent_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Percentage change in vegetation cover between analysis dates."
        )

    with col3:
        bsi_change = stats.get('bsi_change', 0)
        # Bare Soil: increase = bad (red), decrease = good (green)
        st.metric(
            "Bare Soil",
            f"{stats.get('percent_bare_soil', 0):.1f}%",
            delta=f"{bsi_change:+.3f}" if bsi_change != 0 else None,
            delta_color="inverse",  # red for +, green for -
            help="Percentage of area with exposed bare soil. Lower values indicate better vegetation cover."
        )

    with col4:
        st.metric(
            "Water Presence",
            f"{stats.get('percent_water', 0):.1f}%",
            help="Percentage of area with water bodies or saturated soil."
        )

    # Secondary Metrics with tooltips
    st.markdown("### Additional Indices")
    col1, col2, col3 = st.columns(3)

    with col1:
        savi_change = stats.get('savi_change', 0)
        # SAVI: increase = good (green), decrease = bad (red)
        st.metric(
            "SAVI",
            f"{stats.get('savi_after_mean', 0):.3f}",
            delta=f"{savi_change:+.3f}" if savi_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Soil Adjusted Vegetation Index: better for sparse vegetation as it accounts for soil brightness."
        )
        evi_change = stats.get('evi_change', 0)
        # EVI: increase = good (green), decrease = bad (red)
        st.metric(
            "EVI",
            f"{stats.get('evi_after_mean', 0):.3f}",
            delta=f"{evi_change:+.3f}" if evi_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Enhanced Vegetation Index: more sensitive in high-biomass areas and corrects for atmospheric effects."
        )

    with col2:
        ndmi_change = stats.get('ndmi_change', 0)
        # NDMI: increase = good (green), decrease = bad (red)
        st.metric(
            "NDMI",
            f"{stats.get('ndmi_after_mean', 0):.3f}",
            delta=f"{ndmi_change:+.3f}" if ndmi_change != 0 else None,
            delta_color="normal",  # green for +, red for -
            help="Normalized Difference Moisture Index: measures vegetation water content. Higher values = more moisture."
        )
        bsi_val_change = stats.get('bsi_change', 0)
        # BSI: increase = bad (red), decrease = good (green)
        st.metric(
            "BSI",
            f"{stats.get('bsi_after_mean', 0):.3f}",
            delta=f"{bsi_val_change:+.3f}" if bsi_val_change != 0 else None,
            delta_color="inverse",  # red for +, green for -
            help="Bare Soil Index: identifies bare soil areas. Higher values indicate more exposed soil (negative for rehab)."
        )

    with col3:
        st.metric(
            "Moisture Stressed",
            f"{stats.get('percent_moisture_stressed', 0):.1f}%",
            help="Percentage of vegetation showing signs of water stress."
        )
        st.metric(
            "Dense Vegetation",
            f"{stats.get('percent_dense_veg', 0):.1f}%",
            help="Percentage of area with dense, healthy vegetation (NDVI > 0.6)."
        )

    # Terrain stats if available
    if terrain_stats and terrain_stats.get('slope_mean'):
        st.markdown("### Terrain Analysis")
        col1, col2, col3 = st.columns(3)

        with col1:
            st.metric("Mean Slope", f"{terrain_stats.get('slope_mean', 0):.1f}°")

        with col2:
            st.metric("Steep Areas", f"{terrain_stats.get('percent_steep', 0):.1f}%")

        with col3:
            if 'percent_high_erosion_risk' in terrain_stats:
                st.metric("High Erosion Risk", f"{terrain_stats.get('percent_high_erosion_risk', 0):.1f}%",
                          delta_color="inverse")

    # Land cover stats if available
    if land_cover_stats and land_cover_stats.get('vegetation_cover_after'):
        st.markdown("### Land Cover")
        col1, col2 = st.columns(2)

        with col1:
            st.metric("Vegetation Cover",
                      f"{land_cover_stats.get('vegetation_cover_after', 0):.1f}%",
                      delta=f"{land_cover_stats.get('vegetation_cover_change', 0):.1f}%")

        with col2:
            st.metric("Bare Ground",
                      f"{land_cover_stats.get('bare_ground_after', 0):.1f}%",
                      delta=f"{land_cover_stats.get('bare_ground_change', 0):.1f}%",
                      delta_color="inverse")