app_tsptw.py

"""
Streamlit App for TSPTW - app_tsptw.py

"""

import streamlit as st
import torch
import matplotlib.pyplot as plt
import os
import numpy as np

from models.tsptw_gnn_model import TSPTWGraphEncoder
from models.tsptw_trainer import load_tsptw_model
from utils.tsptw_solver import TSPTWDPDPSolver, GreedyTSPTWSolver
from utils.tsptw_visualization import plot_tour, compare_tsptw_solutions
from data.tsptw_generator import TSPTWDataGenerator

st.set_page_config(page_title="TSPTW Solver", page_icon="⏰", layout="wide")
st.title("⏰ TSP with Time Windows Solver")
st.markdown("### Deep Policy Dynamic Programming (DPDP) with Learned Guidance")

st.sidebar.header("⚙️ Configuration")

st.sidebar.subheader("Problem Settings")
n_customers = st.sidebar.slider("Number of Customers", 5, 30, 15)
grid_size = st.sidebar.slider("Grid Size", 50, 200, 100)

st.sidebar.subheader("DPDP Solver Settings")
beam_size = st.sidebar.select_slider(
    "Beam Size (B)",
    options=[10, 50, 100, 500, 1000, 5000],
    value=1000,
    help="Number of states in DP beam search"
)

st.sidebar.subheader("Visualization")
show_heatmap = st.sidebar.checkbox("Show Heatmap Overlay", value=True)

model_path = "models/pretrained/tsptw_best_model.pt"


@st.cache_resource
def load_trained_model(path, device):
    """Load trained model (cached)."""
    try:
        if not os.path.exists(path):
            return None, f"❌ Model not found: {path}\n\nTrain first:\n  python train_tsptw.py"

        model = TSPTWGraphEncoder(
            embedding_dim=128,
            n_layers=5,
            n_heads=4
        )

        checkpoint = torch.load(path, map_location=device)
        state_dict = checkpoint.get("model_state_dict", checkpoint)
        missing, unexpected = model.load_state_dict(state_dict, strict=False)

        if missing:
            print(f"⚠️ Missing keys: {len(missing)} -> {missing[:5]}")
        if unexpected:
            print(f"⚠️ Unexpected keys: {len(unexpected)} -> {unexpected[:5]}")

        model.to(device)
        model.eval()
        return model, None

    except Exception as e:
        return None, str(e)



def generate_tsptw_instance(n_nodes, grid_size, seed=None):
    gen = TSPTWDataGenerator(n_nodes, grid_size, seed=seed)
    coords, time_windows = gen.generate_instance()
    return coords, time_windows


def solve_instance(model, coords, time_windows, method='dpdp', beam_size=1000):
    """Solve TSPTW instance."""
    if method == 'dpdp':
        solver = TSPTWDPDPSolver(model, beam_size=beam_size, device='cpu', use_learned_guidance=True)
        tour, cost, makespan, heatmap = solver.solve(coords, time_windows)
        return tour, cost, makespan, heatmap
    else:
        solver = GreedyTSPTWSolver()
        tour, cost, makespan = solver.solve(coords, time_windows)
        return tour, cost, makespan, None


def display_tour_details(tour, coords, time_windows, title="Tour Timing Details"):
    """Display detailed time schedule table with correct time-window logic."""
    st.markdown(f"#### 🕒 {title}")

    rows = []
    current_time = time_windows[0, 0].item()

    for idx, node in enumerate(tour[:-1]):
        next_node = tour[idx + 1]
        travel_time = torch.norm(coords[node] - coords[next_node]).item()
        arrival = current_time + travel_time

        lower, upper = time_windows[next_node]
        lower, upper = lower.item(), upper.item()

      
        waiting = max(0, lower - arrival)
        actual_time = max(arrival, lower)

        
        if arrival > upper:
            status = "🔴 Late"
        elif arrival < lower:
            status = "⏱ Wait"
        else:
            status = "🟢 On Time"

        rows.append({
            "Step": f"{idx+1}. {node} → {next_node}",
            "Arrive": f"{arrival:.1f}",
            "Window": f"[{lower:.0f}, {upper:.0f}]",
            "Status": status
        })

        current_time = actual_time

    
    st.dataframe(rows, use_container_width=True, hide_index=True)


def main():
    device = 'cpu'

    
    model, error = load_trained_model(model_path, device)
    if error:
        st.error(error)
        st.stop()
    st.sidebar.success("✅ Model loaded successfully!")

  
    with st.expander("ℹ️ About DPDP TSPTW Solver", expanded=False):
        st.markdown("""
        **TSP with Time Windows (TSPTW)**  
        Each customer has a time window [lower, upper]:  
        - Must arrive before upper bound  
        - Can wait if arrive earlier  
        - Objective: Minimize total distance or makespan  

        **Deep Policy Dynamic Programming (DPDP)**  
        - State = (visited, current, time)  
        - Beam search with dominance pruning  
        - Guided by GNN-predicted edge probabilities  
        - Handles time-window feasibility gracefully  
        """)

    
    col1, col2 = st.columns([3, 1])
    with col2:
        if st.button("🔄 Generate New", use_container_width=True):
            seed = np.random.randint(0, 10000)
            coords, time_windows = generate_tsptw_instance(n_customers + 1, grid_size, seed)
            st.session_state.update({
                "coords": coords,
                "time_windows": time_windows,
                "solution": None,
                "cost": None,
                "makespan": None,
                "heatmap": None,
                "greedy_solution": None,
                "greedy_cost": None
            })
            st.rerun()

    
    if "coords" not in st.session_state:
        coords, time_windows = generate_tsptw_instance(n_customers + 1, grid_size)
        st.session_state.coords = coords
        st.session_state.time_windows = time_windows

    coords = st.session_state.coords
    time_windows = st.session_state.time_windows

    
    tw = time_windows
    avg_window = (tw[:, 1] - tw[:, 0]).mean().item()
    horizon = tw[:, 1].max().item()

    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Nodes", f"{n_customers + 1}")
    col2.metric("Avg Window Size", f"{avg_window:.1f}")
    col3.metric("Time Horizon", f"{horizon:.1f}")
    col4.metric("Grid Size", f"{grid_size}")

    st.markdown("---")

    
    col1, col2, col3 = st.columns(3)
    with col1:
        if st.button("🚀 Solve with DPDP", type="primary", use_container_width=True):
            with st.spinner(f"Running DPDP (beam={beam_size})..."):
                try:
                    tour, cost, makespan, heatmap = solve_instance(model, coords, time_windows, 'dpdp', beam_size)
                    st.session_state.update({
                        "solution": tour,
                        "cost": cost,
                        "makespan": makespan,
                        "heatmap": heatmap,
                        "method": "DPDP"
                    })
                except Exception as e:
                    st.error(f"Error: {e}")

    with col2:
        if st.button("⚡ Solve with Greedy", use_container_width=True):
            with st.spinner("Running greedy baseline..."):
                try:
                    tour, cost, makespan, _ = solve_instance(None, coords, time_windows, 'greedy')
                    st.session_state.update({
                        "greedy_solution": tour,
                        "greedy_cost": cost,
                        "greedy_makespan": makespan
                    })
                except Exception as e:
                    st.error(f"Error: {e}")

    with col3:
        if st.button("🔄 Compare Both", use_container_width=True):
            with st.spinner("Solving with both..."):
                try:
                    tour_dp, cost_dp, makespan_dp, heatmap = solve_instance(model, coords, time_windows, 'dpdp', beam_size)
                    tour_gr, cost_gr, makespan_gr, _ = solve_instance(None, coords, time_windows, 'greedy')
                    st.session_state.update({
                        "solution": tour_dp,
                        "cost": cost_dp,
                        "makespan": makespan_dp,
                        "heatmap": heatmap,
                        "greedy_solution": tour_gr,
                        "greedy_cost": cost_gr,
                        "greedy_makespan": makespan_gr
                    })
                except Exception as e:
                    st.error(f"Error: {e}")

    
    if st.session_state.get("solution") and st.session_state.get("greedy_solution"):
        st.markdown("---")
        
        
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("DPDP Cost", f"{st.session_state.cost:.2f}")
        with col2:
            st.metric("Greedy Cost", f"{st.session_state.greedy_cost:.2f}")
        with col3:
            improvement = ((st.session_state.greedy_cost - st.session_state.cost) / st.session_state.greedy_cost) * 100
            st.metric("Improvement", f"{improvement:+.1f}%", delta=f"{improvement:.1f}%")

        
        st.markdown("### 📊 Solution Comparison")
        fig = compare_tsptw_solutions(
            coords,
            [st.session_state.solution, st.session_state.greedy_solution],
            time_windows,
            [f"DPDP (Cost: {st.session_state.cost:.2f})",
             f"Greedy (Cost: {st.session_state.greedy_cost:.2f})"]
        )
        st.pyplot(fig)
        plt.close()

        
        
        

    elif st.session_state.get("solution"):
        st.info("⚡ Run Greedy solver to see comparison")
        
    elif st.session_state.get("greedy_solution"):
        st.info("🚀 Run DPDP solver to see comparison")
        
    else:
        st.info("👆 Click 'Compare Both' to see side-by-side results!")


if __name__ == "__main__":
    main()