### SWEDEN TAC FOR 100 TRAIN-KM AND PR ###
### WORKS PERFECT ###


import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from datetime import datetime, timedelta

# ==========================================
# 1. CLASS DEFINITIONS & INFRASTRUCTURE
# ==========================================

class Train:
    def __init__(self, id, direction, priority, stops, start_time_offset, speed_kmh, dwell_min):
        self.id = str(id)
        self.direction = direction
        self.priority = int(priority)
        self.speed_kmh = float(speed_kmh)
        self.dwell_min = float(dwell_min)
        self.stops = stops
        self.start_time_offset = float(start_time_offset)
        self.schedule = []

    def copy(self):
        new_t = Train(self.id, self.direction, self.priority, self.stops,
                      self.start_time_offset, self.speed_kmh, self.dwell_min)
        new_t.schedule = []
        for s in self.schedule:
            node_copy = s.copy()
            if 'min_arr_constraint' not in node_copy: node_copy['min_arr_constraint'] = None
            new_t.schedule.append(node_copy)
        return new_t

class Infrastructure:
    def __init__(self, station_df, spacing_km=5.0):
        # Access by index to be robust against header names
        self.stations = station_df.iloc[:, 0].astype(str).tolist()
        self.distances = pd.to_numeric(station_df.iloc[:, 1], errors='coerce').fillna(0.0).tolist()
        self.station_map = {name: i for i, name in enumerate(self.stations)}
        self.dist_map = {name: d for name, d in zip(self.stations, self.distances)}
        self.spacing_km = float(spacing_km)

    def get_travel_time(self, s1, s2, speed_kmh):
        d1 = self.dist_map.get(s1, 0.0)
        d2 = self.dist_map.get(s2, 0.0)
        dist_km = abs(d1 - d2)
        if speed_kmh <= 0: return timedelta(minutes=1)
        hours = dist_km / speed_kmh
        return timedelta(hours=hours)

    def get_dynamic_headway(self, speed_kmh):
        if speed_kmh <= 0: return timedelta(minutes=5)
        minutes = (self.spacing_km / speed_kmh) * 60.0
        return timedelta(minutes=max(1.0, minutes))

# ==========================================
# 2. PHYSICS ENGINE
# ==========================================

def update_train_times(train, infra, base_time):
    if not train.schedule:
        current_time = base_time + timedelta(minutes=train.start_time_offset)
        train.schedule = []
        for s in train.stops:
            train.schedule.append({
                'station': s,
                'arr': current_time,
                'dep': current_time,
                'min_arr_constraint': None,
                'manual_delay': timedelta(0)
            })

    for i in range(len(train.schedule)):
        current_node = train.schedule[i]
        manual_delay = current_node.get('manual_delay', timedelta(0))

        # A. ARRIVAL
        if i == 0:
            base_origin_time = base_time + timedelta(minutes=train.start_time_offset)
            current_node['arr'] = base_origin_time + manual_delay
        else:
            prev_node = train.schedule[i-1]
            phys_run_time = infra.get_travel_time(prev_node['station'], current_node['station'], train.speed_kmh)
            min_phys_arr = prev_node['dep'] + phys_run_time
            safety_limit = current_node.get('min_arr_constraint')

            delayed_phys_arr = min_phys_arr + manual_delay

            if safety_limit:
                current_node['arr'] = max(delayed_phys_arr, safety_limit)
            else:
                current_node['arr'] = delayed_phys_arr

        # B. DEPARTURE
        dwell = timedelta(minutes=train.dwell_min) if i < len(train.schedule) - 1 else timedelta(0)
        min_dep = current_node['arr'] + dwell
        current_node['dep'] = max(current_node['dep'], min_dep)

# ==========================================
# 3. CONFLICT DETECTOR
# ==========================================

def detect_conflicts(trains, infra):
    conflicts = []
    for i in range(len(trains)):
        for j in range(i + 1, len(trains)):
            t1 = trains[i]
            t2 = trains[j]
            for k in range(len(t1.schedule) - 1):
                u, v = t1.schedule[k]['station'], t1.schedule[k+1]['station']
                t2_stations = [s['station'] for s in t2.schedule]
                if u in t2_stations and v in t2_stations:
                    idx_u_t2 = t2_stations.index(u)
                    idx_v_t2 = t2_stations.index(v)
                    if abs(idx_u_t2 - idx_v_t2) == 1:
                        t1_ent, t1_ext = t1.schedule[k]['dep'], t1.schedule[k+1]['arr']
                        if idx_u_t2 < idx_v_t2:
                            t2_ent, t2_ext = t2.schedule[idx_u_t2]['dep'], t2.schedule[idx_v_t2]['arr']
                            direction = "Same"
                        else:
                            t2_ent, t2_ext = t2.schedule[idx_v_t2]['dep'], t2.schedule[idx_u_t2]['arr']
                            direction = "Opposite"

                        is_conflict = False
                        conflict_reason = ""

                        if direction == "Opposite":
                            if max(t1_ent, t2_ent) < min(t1_ext, t2_ext):
                                is_conflict = True; conflict_reason = "Head-on"
                        else:
                            if t1_ent <= t2_ent:
                                l_ent, f_ent = t1_ent, t2_ent
                                leader_ext, follower_ext = t1_ext, t2_ext
                                hw = infra.get_dynamic_headway(t2.speed_kmh)
                            elif t1_ent > t2_ent:
                                l_ent, f_ent = t2_ent, t1_ent
                                leader_ext, follower_ext = t2_ext, t1_ext
                                hw = infra.get_dynamic_headway(t1.speed_kmh)

                            tol = timedelta(seconds=1)
                            if f_ent < l_ent + hw - tol:
                                is_conflict = True; conflict_reason = f"Entry HW ({hw.seconds//60}m)"
                            elif follower_ext < leader_ext + hw - tol:
                                is_conflict = True; conflict_reason = f"Exit HW ({hw.seconds//60}m)"

                        if is_conflict:
                            conflicts.append({'time': max(t1_ent, t2_ent), 'trains': (t1, t2), 'segment': (u, v), 'dir': direction, 'desc': f"{conflict_reason} on {u}->{v}"})
    conflicts.sort(key=lambda x: x['time'])
    return conflicts

# ==========================================
# 4. RESOLVER (With Max Buffer Logic)
# ==========================================

def resolve_all_conflicts(trains, infra, base_time, selected_train_id=None, max_buffer_min=30.0):
    iteration = 0
    max_iter = 200
    logs = []
    max_buffer_td = timedelta(minutes=float(max_buffer_min))

    while iteration < max_iter:
        conflicts = detect_conflicts(trains, infra)
        if not conflicts:
            logs.append("No conflicts found.")
            break

        c = conflicts[0]
        t1, t2 = c['trains']
        u, v = c['segment']
        direction = c['dir']

        # If the user selected a train in the View & Optimization panel,
        # treat it as the optimization target for this conflict. This makes the
        # Select Train list-box functional: the selected service is protected
        # before normal priority/tie-break logic is applied.
        selected_id = str(selected_train_id) if selected_train_id not in (None, "", "None") else None
        if selected_id and str(t1.id) == selected_id and str(t2.id) != selected_id:
            winner, loser = t1, t2
            logs.append(f"Iter {iteration}: Optimization target {winner.id} given dispatch priority.")
        elif selected_id and str(t2.id) == selected_id and str(t1.id) != selected_id:
            winner, loser = t2, t1
            logs.append(f"Iter {iteration}: Optimization target {winner.id} given dispatch priority.")
        elif t1.priority > t2.priority:
            winner, loser = t1, t2
        elif t2.priority > t1.priority:
            winner, loser = t2, t1
        else:
            winner, loser = (t1, t2) if t1.id < t2.id else (t2, t1)

        def get_idx(train, station): return [s['station'] for s in train.schedule].index(station)

        if direction == "Same":
            u_w, v_w = get_idx(winner, u), get_idx(winner, v)
            u_l, v_l = get_idx(loser, u), get_idx(loser, v)
            if u_w > v_w: u_w, v_w = v_w, u_w
            if u_l > v_l: u_l, v_l = v_l, u_l

            if winner.schedule[u_w]['dep'] <= loser.schedule[u_l]['dep']:
                leader, follower = winner, loser
                l_u, l_v = u_w, v_w
                f_u, f_v = u_l, v_l
                follower_is_winner = False
            else:
                leader, follower = loser, winner
                l_u, l_v = u_l, v_l
                f_u, f_v = u_w, v_w
                follower_is_winner = True

            hw = infra.get_dynamic_headway(follower.speed_kmh)

            applied_buffer = False
            if follower_is_winner and selected_train_id and str(follower.id) == str(selected_train_id):
                 gap_needed = (follower.schedule[f_u]['dep'] + hw) - leader.schedule[l_u]['dep']
                 current_buffer = leader.schedule[l_u].get('manual_delay', timedelta(0))
                 total_new_buffer = current_buffer + gap_needed

                 if gap_needed > timedelta(0):
                     if total_new_buffer <= max_buffer_td:
                         logs.append(f"Iter {iteration}: Selected Train Optimization -> Applying {gap_needed.total_seconds()//60:.1f}m Buffer to {leader.id} at {u}. Total: {total_new_buffer.total_seconds()//60:.1f}m")
                         leader.schedule[l_u]['manual_delay'] = total_new_buffer
                         update_train_times(leader, infra, base_time)
                         applied_buffer = True
                     else:
                         logs.append(f"Iter {iteration}: Selected Train Optimization Rejected. Total buffer needed ({total_new_buffer.total_seconds()//60:.1f}m) exceeds Limit ({max_buffer_min}m).")

            if not applied_buffer:
                req_dep = leader.schedule[l_u]['dep'] + hw
                if follower.schedule[f_u]['dep'] < req_dep:
                    logs.append(f"Iter {iteration}: Holding {follower.id} at {u}. Req HW: {hw}")
                    follower.schedule[f_u]['dep'] = req_dep

                req_arr = leader.schedule[l_v]['arr'] + hw
                current_const = follower.schedule[f_v].get('min_arr_constraint')
                if current_const is None or req_arr > current_const:
                    logs.append(f"Iter {iteration}: Slowing {follower.id} on {u}->{v}. Req HW: {hw}")
                    follower.schedule[f_v]['min_arr_constraint'] = req_arr

                update_train_times(follower, infra, base_time)

                if follower_is_winner:
                     if l_v < len(leader.schedule):
                         follower_next_dep = follower.schedule[f_v]['dep']
                         req_leader_dep = follower_next_dep + hw
                         if leader.schedule[l_v]['dep'] < req_leader_dep:
                             logs.append(f"Iter {iteration}: Overtake at {v}. Holding {leader.id}.")
                             leader.schedule[l_v]['dep'] = req_leader_dep
                             update_train_times(leader, infra, base_time)

        else:
            u_w, v_w = get_idx(winner, u), get_idx(winner, v)
            t_clear = max(winner.schedule[u_w]['arr'], winner.schedule[v_w]['arr'])
            u_l, v_l = get_idx(loser, u), get_idx(loser, v)
            l_ent = min(u_l, v_l)
            if loser.schedule[l_ent]['dep'] < t_clear:
                logs.append(f"Iter {iteration}: Opp Dir. Holding {loser.id} at {u}.")
                loser.schedule[l_ent]['dep'] = t_clear
                update_train_times(loser, infra, base_time)
        iteration += 1
    return logs

# ==========================================
# 5. DATA EXPORT & VISUALS
# ==========================================

def generate_timetable_matrix(trains, infra):
    data = []
    columns = ["Train ID", "Direction"]
    for s in infra.stations:
        columns.append(f"{s} Arr")
        columns.append(f"{s} Dep")

    for t in trains:
        row = {"Train ID": t.id, "Direction": t.direction}
        for s_data in t.schedule:
            s_name = s_data['station']
            row[f"{s_name} Arr"] = s_data['arr'].strftime("%H:%M:%S")
            row[f"{s_name} Dep"] = s_data['dep'].strftime("%H:%M:%S")

        for s in infra.stations:
            if f"{s} Arr" not in row: row[f"{s} Arr"] = "-"
            if f"{s} Dep" not in row: row[f"{s} Dep"] = "-"

        data.append(row)

    df = pd.DataFrame(data)
    final_cols = ["Train ID", "Direction"]
    for s in infra.stations:
        final_cols.append(f"{s} Arr")
        final_cols.append(f"{s} Dep")

    return df[final_cols]

def calculate_delay_stats(nom_trains, adj_trains, primary_delay_df):
    stats = []
    primary_map = {}
    if primary_delay_df is not None and not primary_delay_df.empty:
        for _, row in primary_delay_df.iterrows():
             if not pd.isna(row.iloc[0]):
                 tid = str(row.iloc[0])
                 try: val = float(row.iloc[2])
                 except: val = 0.0
                 primary_map[tid] = primary_map.get(tid, 0) + val

    for t_nom, t_adj in zip(nom_trains, adj_trains):
        nom_end = t_nom.schedule[-1]['arr']
        adj_end = t_adj.schedule[-1]['arr']
        total_delay = (adj_end - nom_end).total_seconds() / 60.0
        p_delay = primary_map.get(t_nom.id, 0.0)
        s_delay = max(0.0, total_delay - p_delay)

        if p_delay > 0 and s_delay > 0.1: d_type = "Both"
        elif p_delay > 0: d_type = "Primary Only"
        elif s_delay > 0.1: d_type = "Secondary"
        else: d_type = "On Time"

        stats.append({
            "Train ID": t_nom.id,
            "Total Delay": round(total_delay, 1),
            "Primary": round(p_delay, 1),
            "Secondary": round(s_delay, 1),
            "Type": d_type
        })
    return pd.DataFrame(stats)

def generate_buffer_report(trains, primary_delay_df):
    report_data = []
    requested_delays = {}
    if primary_delay_df is not None and not primary_delay_df.empty:
        for _, row in primary_delay_df.iterrows():
            if not pd.isna(row.iloc[0]):
                tid = str(row.iloc[0])
                st = str(row.iloc[1])
                try: val = float(row.iloc[2])
                except: val = 0.0
                requested_delays[(tid, st)] = val

    for t in trains:
        for node in t.schedule:
            st = node['station']
            total_manual = node.get('manual_delay', timedelta(0)).total_seconds() / 60.0
            requested = requested_delays.get((t.id, st), 0.0)
            buffer_val = total_manual - requested

            if buffer_val > 0.1:
                report_data.append({
                    "Train ID": t.id,
                    "Station": st,
                    "Requested Delay (min)": round(requested, 1),
                    "Total Added Delay (min)": round(total_manual, 1),
                    "System Buffer (min)": round(buffer_val, 1)
                })

    if not report_data:
        return pd.DataFrame(columns=["Train ID", "Station", "Requested Delay (min)", "Total Added Delay (min)", "System Buffer (min)"])

    return pd.DataFrame(report_data)

# ==========================================
# 3B. ICONOGRAPHY & LIVE METRIC STRIP
# ==========================================
# Gradio has no npm/CDN icon library, so section glyphs are a small set of
# hand-drawn, single-stroke geometric marks (uniform stroke-width, one accent
# color) instead of emoji - consistent, offline-safe, and not decorative.

_ICON_PATHS = {
    "train": '<rect x="5" y="3" width="14" height="13" rx="3"/><path d="M5 11h14"/>'
             '<circle cx="8.5" cy="19" r="1.3"/><circle cx="15.5" cy="19" r="1.3"/>'
             '<path d="M8 16l-2 3M16 16l2 3"/>',
    "clock": '<circle cx="12" cy="12" r="9"/><path d="M12 7v5l3.5 2"/>',
    "target": '<circle cx="12" cy="12" r="8.5"/><circle cx="12" cy="12" r="4.5"/>'
              '<circle cx="12" cy="12" r=".8" fill="currentColor" stroke="none"/>',
    "chart": '<path d="M4 19V5"/><path d="M4 19h16"/><path d="M7 15l4-5 3 3 5-7"/>',
    "coin": '<circle cx="12" cy="12" r="8.5"/><path d="M14.6 8.3a4.6 4.6 0 1 0 0 7.4"/>'
            '<path d="M8 11h5M8 13.2h4.2"/>',
    "building": '<rect x="4" y="3" width="10" height="18"/><rect x="14" y="9" width="6" height="12"/>'
                '<path d="M7 7h1M10 7h1M7 11h1M10 11h1M7 15h1M10 15h1"/>',
    "dice": '<path d="M3 6h4l7 12h4"/><path d="M14 6h4l3 3"/><path d="M3 17h4l3-4"/>'
            '<path d="M18 3l3 3-3 3"/><path d="M18 14l3 3-3 3"/>',
    "scale": '<path d="M12 4v16M6 8h12"/><path d="M4 8l2.5 5a2.5 2.5 0 0 0 5 0L9 8"/>'
             '<path d="M15 8l2.5 5a2.5 2.5 0 0 0 5 0L20 8"/><path d="M9 20h6"/>',
    "list": '<path d="M8 6h12M8 12h12M8 18h12"/><path d="M4 6h.01M4 12h.01M4 18h.01"/>',
    "bars": '<path d="M4 20V10M10 20V4M16 20v-7M4 20h16"/>',
    "alert": '<path d="M10.3 3.86 1.8 18a1 1 0 0 0 .9 1.5h18.6a1 1 0 0 0 .9-1.5L13.7 3.86a1 1 0 0 0-1.72 0Z"/>'
             '<path d="M12 9v4"/><path d="M12 16.3h.01"/>',
    "check-circle": '<circle cx="12" cy="12" r="9"/><path d="M8.5 12.5l2.5 2.5 5-5"/>',
}

def icon_svg(name, size=18):
    return (f'<svg class="hort-icon" width="{size}" height="{size}" viewBox="0 0 24 24" '
            f'fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" '
            f'stroke-linejoin="round">{_ICON_PATHS[name]}</svg>')

def section_header(icon, text, level=3):
    return f'<h{level} class="hort-heading">{icon_svg(icon)}<span>{text}</span></h{level}>'

def build_metric_html(df_stats, logs_df):
    """Live KPI row: recomputed after every simulation run instead of
    showing fixed placeholder copy."""
    if df_stats is None or df_stats.empty:
        trains_n, conflicts_n, avg_delay, on_time_pct = "-", "-", "-", "-"
        flag_conflicts = False
    else:
        trains_n = str(len(df_stats))
        avg_delay = f"{df_stats['Total Delay'].mean():.1f}"
        on_time = int((df_stats['Type'] == 'On Time').sum())
        on_time_pct = str(round(on_time / len(df_stats) * 100))
        if logs_df is not None and not logs_df.empty:
            conflicts = sum(1 for v in logs_df.iloc[:, 0] if isinstance(v, str) and v.startswith("Iter"))
        else:
            conflicts = 0
        conflicts_n = str(conflicts)
        flag_conflicts = conflicts > 0

    conflict_icon_cls = "kpi-icon is-flagged" if flag_conflicts else "kpi-icon"
    conflict_val_cls = "kpi-value is-flagged" if flag_conflicts else "kpi-value"

    def card(icon, icon_cls, label_text, value_cls, value_html, caption):
        return f'''
        <div class="kpi-card">
            <div class="kpi-head">
                <div class="{icon_cls}">{icon_svg(icon, 15)}</div>
                <span class="kpi-label">{label_text}</span>
            </div>
            <div>
                <div class="{value_cls}">{value_html}</div>
                <div class="kpi-caption">{caption}</div>
            </div>
        </div>'''

    return f'''
    <div class="kpi-grid">
        {card("train", "kpi-icon", "Trains simulated", "kpi-value", trains_n, "from Train Requests")}
        {card("alert", conflict_icon_cls, "Conflicts resolved", conflict_val_cls, conflicts_n, "auto-resolved")}
        {card("clock", "kpi-icon is-neutral", "Avg total delay", "kpi-value", f'{avg_delay}<span class="unit"> min</span>', "across all trains")}
        {card("check-circle", "kpi-icon", "On-time trains", "kpi-value", f'{on_time_pct}<span class="unit">%</span>', "within buffer cap")}
    </div>
    '''

def create_filtered_marey(t_nom, t_del, t_adj, infra, view_mode, selected_train_id):
    fig = go.Figure()
    # Categorical palette tuned bright enough to read on a near-black surface.
    colors = ['#5DB8FF', '#FF8A8A', '#FFD166', '#7EE9A8', '#C9A6FF', '#FFA45C']
    def get_dist(s): return infra.dist_map.get(s, 0)

    def add_lines(trains, color, style, name_suffix, opacity, visible, width=2):
        for i, t in enumerate(trains):
            is_selected = (str(t.id) == str(selected_train_id))
            x, y, txt = [], [], []
            if is_selected and style == 'solid':
                c = '#3DDC84'; w = 5; op = 1.0
            else:
                c = colors[i % len(colors)] if color == 'auto' else color
                w = width; op = opacity

            for s in t.schedule:
                d = get_dist(s['station'])
                x.extend([s['arr'], s['dep']])
                y.extend([d, d])
                txt.extend([f"Arrival at {s['station']}", f"Departure from {s['station']}"])

            fig.add_trace(go.Scatter(
                x=x,
                y=y,
                mode='lines+markers' if style == 'solid' else 'lines',
                name=f"T{t.id} {name_suffix}" + (" · selected" if is_selected else ""),
                line=dict(color=c, width=w, dash='solid' if style == 'solid' else style, shape='spline'),
                marker=dict(size=7, color=c, line=dict(width=1, color='#141815')) if style == 'solid' else None,
                hovertext=txt,
                hovertemplate="%{hovertext}<br>%{x|%H:%M:%S}<br>%{y} km<extra></extra>",
                opacity=op,
                visible=visible
            ))

    if view_mode in ["All Scenarios", "Nominal Only"]: add_lines(t_nom, '#5A6660', 'dash', "(Nominal)", 0.55, True)
    if view_mode in ["All Scenarios", "Conflicted Only"]: add_lines(t_del, '#FF6B6B', 'dot', "(Conflicted)", 0.78, True if view_mode != "All Scenarios" else 'legendonly')
    if view_mode in ["All Scenarios", "Resolved Only"]: add_lines(t_adj, 'auto', 'solid', "(Resolved)", 1.0, True, width=3)

    fig.update_layout(
        title=dict(
            text=f"Marey Diagram · {view_mode}",
            x=0.02,
            y=0.98,
            xanchor='left',
            yanchor='top',
            font=dict(size=14, color='#E7ECE7', family='JetBrains Mono, monospace')
        ),
        xaxis_title="Time",
        yaxis_title="Station / distance",
        yaxis=dict(tickmode='array', tickvals=infra.distances, ticktext=infra.stations, gridcolor='rgba(231,236,231,0.07)', zeroline=False),
        xaxis=dict(tickformat="%H:%M", gridcolor='rgba(231,236,231,0.07)', zeroline=False),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)',
        font=dict(color='#7C8C82', family='JetBrains Mono, monospace'),
        legend=dict(
            orientation='h',
            yanchor='top',
            y=-0.20,
            xanchor='center',
            x=0.5,
            bgcolor='rgba(20,24,21,0.9)',
            bordercolor='rgba(231,236,231,0.12)',
            borderwidth=1
        ),
        hoverlabel=dict(bgcolor='#191F1B', font_color='#E7ECE7', bordercolor='rgba(231,236,231,0.15)'),
        margin=dict(l=58, r=24, t=70, b=120),
        height=650
    )
    return fig

def simulate_capacity_vs_reliability(infra, fleet_profiles, delay_data, max_trains_input, tac_fee, pr_rate, d_mode, r_prob, r_min, r_max, im_resp_share):
    import random
    from datetime import datetime, timedelta
    import pandas as pd
    import plotly.graph_objects as go

    results = []

    try: max_t = int(max_trains_input)
    except: max_t = 20
    if max_t < 2: max_t = 2

    densities = list(range(2, max_t + 1))
    base_time = datetime(2024, 1, 1, 8, 0, 0)

    try: im_share = float(im_resp_share)
    except: im_share = 50.0

    best_profit = -float('inf')
    best_trains = []
    best_delay_df = pd.DataFrame() # <-- Initialize the dataframe to hold the best delays

    for tph in densities:
        scheduled_headway = 60.0 / tph
        trains = []
        current_run_delays = [] # <-- Track delays just for this specific TPH run

        for i in range(tph):
            profile = fleet_profiles[i % len(fleet_profiles)]
            new_train = profile.copy()

            generated_id = str(i + 1)
            new_train.original_id = generated_id
            new_train.id = f"T{generated_id}_{profile.id}"
            new_train.start_time_offset = i * scheduled_headway
            new_train.had_primary_delay = False # Track this cleanly!

            update_train_times(new_train, infra, base_time)

            # --- INJECT DELAYS BASED ON USER MODE ---
            if d_mode == "Predefined (from Delays Tab)" and delay_data is not None and not delay_data.empty:
                for _, d_row in delay_data.iterrows():
                    d_keys = d_row.keys()
                    delay_t_id = ""
                    if "Train ID" in d_keys: delay_t_id = str(d_row["Train ID"]).strip()
                    elif "Trian ID" in d_keys: delay_t_id = str(d_row["Trian ID"]).strip()

                    st_name = str(d_row.get("Station", "")).strip()
                    d_val = 0.0
                    if "Delay" in d_keys: d_val = float(d_row["Delay"])
                    elif "Delay (min)" in d_keys: d_val = float(d_row["Delay (min)"])

                    if delay_t_id == generated_id and d_val > 0:
                        new_train.had_primary_delay = True
                        # Responsibility is decided by the same rule for every
                        # primary delay, regardless of whether the delay is
                        # predefined (from the Delays tab) or randomly generated -
                        # it's set once, before the run, via the "IM Responsibility
                        # Share" slider, so results across capacity levels stay
                        # comparable and consistent between delay modes.
                        party = "IM" if random.uniform(0, 100) <= im_share else "RU"
                        amount = round(d_val * pr_rate, 2)
                        current_run_delays.append({
                            "Train ID": new_train.id, "Station": st_name, "Delay (min)": d_val,
                            "Responsible Party": party, "Amount (€)": amount,
                            "Effect on IM": f"Penalty to IM: -€{amount:,.2f}" if party == "IM" else f"Income for IM: +€{amount:,.2f}"
                        })
                        shift_active = False
                        for stop_event in new_train.schedule:
                            if st_name and st_name.lower() == str(stop_event['station']).lower():
                                shift_active = True
                            if shift_active:
                                stop_event['arr'] += timedelta(minutes=d_val)
                                stop_event['dep'] += timedelta(minutes=d_val)

            elif d_mode == "Randomized":
                # Roll the dice to see if this train gets delayed
                if random.uniform(0, 100) <= float(r_prob):
                    new_train.had_primary_delay = True
                    # Pick a random station (excluding the final destination)
                    rand_idx = random.randint(0, len(new_train.schedule) - 2)
                    rand_station = new_train.schedule[rand_idx]['station']
                    d_val = round(random.uniform(float(r_min), float(r_max)), 1)

                    # Same responsibility rule as the predefined branch above -
                    # the IM Responsibility Share slider, not a fresh 50/50 pick.
                    party = "IM" if random.uniform(0, 100) <= im_share else "RU"
                    amount = round(d_val * pr_rate, 2)

                    # Log the generated delay
                    current_run_delays.append({
                        "Train ID": new_train.id, "Station": rand_station, "Delay (min)": d_val,
                        "Responsible Party": party, "Amount (€)": amount,
                        "Effect on IM": f"Penalty to IM: -€{amount:,.2f}" if party == "IM" else f"Income for IM: +€{amount:,.2f}"
                    })

                    shift_active = False
                    for stop_event in new_train.schedule:
                        if stop_event['station'] == rand_station:
                            shift_active = True
                        if shift_active:
                            stop_event['arr'] += timedelta(minutes=d_val)
                            stop_event['dep'] += timedelta(minutes=d_val)

            trains.append(new_train)

        resolve_all_conflicts(trains, infra, base_time)

        # --- CALCULATE SECONDARY DELAY ---
        total_secondary_delay = 0
        for i, train in enumerate(trains):
            if getattr(train, 'had_primary_delay', False):
                continue

            dummy = fleet_profiles[i % len(fleet_profiles)].copy()
            dummy.start_time_offset = train.start_time_offset
            update_train_times(dummy, infra, base_time)

            last_station_idx = len(dummy.stops) - 1
            expected_arr = dummy.schedule[last_station_idx]['arr']
            actual_arr = train.schedule[last_station_idx]['arr']

            delay_min = (actual_arr - expected_arr).total_seconds() / 60.0
            if delay_min > 0: total_secondary_delay += delay_min

        # --- PRIMARY DELAY RESPONSIBILITY COST (assigned before the run) ---
        # Secondary delays are always paid by IM (PR Penalty above, unchanged).
        # Primary delays are charged to whichever party (IM/RU) was assigned:
        # IM-caused = penalty for IM; RU-caused = income for IM (RU pays IM).
        primary_im_penalty = sum(d["Amount (€)"] for d in current_run_delays if d["Responsible Party"] == "IM")
        primary_ru_income = sum(d["Amount (€)"] for d in current_run_delays if d["Responsible Party"] == "RU")

        revenue = tph * tac_fee
        penalty_cost = total_secondary_delay * pr_rate
        profit = revenue - penalty_cost - primary_im_penalty + primary_ru_income

        # --- CHECK IF THIS IS THE BEST SCENARIO ---
        if profit > best_profit:
            best_profit = profit
            best_trains = [t.copy() for t in trains]
            # Save the delays that were applied to this specific winning scenario!
            if current_run_delays:
                best_delay_df = pd.DataFrame(current_run_delays)
            else:
                best_delay_df = pd.DataFrame(columns=["Train ID", "Station", "Delay (min)", "Responsible Party", "Amount (€)", "Effect on IM"])

        results.append({
            'Capacity (Trains/Hour)': tph,
            'Total Secondary Delay (min)': total_secondary_delay,
            'Revenue (€)': revenue,
            'PR Penalty (€)': penalty_cost,
            'IM Primary Penalty (€)': primary_im_penalty,
            'IM Primary Income (€)': primary_ru_income,
            'Total Profit (€)': profit
        })

    df = pd.DataFrame(results)

    timetable_data = []
    if best_trains:
        for t in best_trains:
            for stop in t.schedule:
                timetable_data.append({
                    "Train ID": t.id,
                    "Station": stop['station'],
                    "Arrival": stop['arr'].strftime("%H:%M:%S"),
                    "Departure": stop['dep'].strftime("%H:%M:%S")
                })
    best_timetable_df = pd.DataFrame(timetable_data)

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df['Capacity (Trains/Hour)'], y=df['Revenue (€)'], mode='lines+markers', name='TAC Revenue', line=dict(color='#5DB8FF', width=3, shape='spline'), marker=dict(size=7)))
    fig.add_trace(go.Scatter(x=df['Capacity (Trains/Hour)'], y=df['PR Penalty (€)'], mode='lines+markers', name='PR Penalty Cost', line=dict(dash='dash', color='#FF6B6B', width=3, shape='spline'), marker=dict(size=7)))
    fig.add_trace(go.Scatter(x=df['Capacity (Trains/Hour)'], y=df['IM Primary Income (€)'] - df['IM Primary Penalty (€)'], mode='lines+markers', name='Primary Delay Net (IM)', line=dict(dash='dot', color='#7EE9A8', width=3, shape='spline'), marker=dict(size=7)))
    fig.add_trace(go.Scatter(x=df['Capacity (Trains/Hour)'], y=df['Total Profit (€)'], mode='lines+markers', name='Total Profit', line=dict(width=4, color='#E7ECE7', shape='spline'), marker=dict(size=8)))

    best_tph = df.loc[df['Total Profit (€)'].idxmax()]['Capacity (Trains/Hour)']
    fig.add_vline(x=best_tph, line_dash="dot", line_color="#3DDC84", annotation_text="Sweet Spot", annotation_position="top left")

    fig.update_layout(
        title=dict(text="Capacity vs. Reliability Sweet Spot", x=0.02, xanchor='left', font=dict(size=18, color='#E7ECE7', family='JetBrains Mono, monospace')),
        xaxis_title="Scheduled trains per hour",
        yaxis_title="Euros (€)",
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)',
        font=dict(color='#7C8C82', family='JetBrains Mono, monospace'),
        xaxis=dict(gridcolor='rgba(231,236,231,0.07)', zeroline=False),
        yaxis=dict(gridcolor='rgba(231,236,231,0.07)', zeroline=False),
        legend=dict(orientation='h', yanchor='bottom', y=1.02, xanchor='right', x=1, bgcolor='rgba(20,24,21,0.9)', bordercolor='rgba(231,236,231,0.12)', borderwidth=1),
        hoverlabel=dict(bgcolor='#191F1B', font_color='#E7ECE7', bordercolor='rgba(231,236,231,0.15)'),
        margin=dict(l=64, r=24, t=88, b=58)
    )

    # --- GENERATE THE TRADITIONAL MAREY DIAGRAM (BOTTOM-LEFT TO TOP-RIGHT) ---
    best_marey_fig = go.Figure()
    dist_map = {st: dist for st, dist in zip(infra.stations, infra.distances)}

    if best_trains:
        for t in best_trains:
            times = []
            locs = []
            for stop in t.schedule:
                times.append(stop['arr'])
                locs.append(dist_map[stop['station']])
                times.append(stop['dep'])
                locs.append(dist_map[stop['station']])

            color = '#FF6B6B' if getattr(t, 'had_primary_delay', False) else '#5DB8FF'

            best_marey_fig.add_trace(go.Scatter(
                x=times,
                y=locs,
                mode='lines+markers',
                name=t.id,
                line=dict(color=color, width=3, shape='spline'),
                marker=dict(size=5, color=color, line=dict(width=1, color='#141815'))
            ))

        station_names = list(dist_map.keys())
        station_distances = list(dist_map.values())

        best_marey_fig.update_layout(
            title=dict(
                text=f"Optimal Space-Time Diagram · {best_tph} Trains/Hour",
                x=0.02,
                y=0.98,
                xanchor='left',
                yanchor='top',
                font=dict(size=14, color='#E7ECE7', family='JetBrains Mono, monospace')
            ),
            xaxis_title="Time",
            yaxis_title="Station / distance",
            xaxis=dict(tickformat="%H:%M", tickangle=-35, gridcolor='rgba(231,236,231,0.07)', zeroline=False),
            yaxis=dict(
                tickmode='array',
                tickvals=station_distances,
                ticktext=[f"{st} ({d}km)" for st, d in zip(station_names, station_distances)],
                gridcolor='rgba(231,236,231,0.07)',
                zeroline=False
            ),
            plot_bgcolor='rgba(0,0,0,0)',
            paper_bgcolor='rgba(0,0,0,0)',
            font=dict(color='#7C8C82', family='JetBrains Mono, monospace'),
            hoverlabel=dict(bgcolor='#191F1B', font_color='#E7ECE7', bordercolor='rgba(231,236,231,0.15)'),
            showlegend=False,
            margin=dict(l=72, r=24, t=70, b=64)
        )

    # Returning 5 elements, matching what run_economic_analysis expects!
    return df, fig, best_timetable_df, best_marey_fig, best_delay_df

# ==========================================
# 6. SIMULATION & GRADIO
# ==========================================

def run_simulation(station_data, train_data, delay_data, spacing_input, plot_view, selected_train_input, max_buffer_input, d_mode, r_prob, r_min, r_max):
    import random
    from datetime import datetime, timedelta
    import pandas as pd # Just in case it's not globally available here

    # ----------------------------------------------------
    # DATA CLEANING & HEADER RECOVERY (Robust to CSV errors)
    # ----------------------------------------------------

    # 1. Clean Train Data
    if train_data is not None:
        train_data = _normalize_uploaded_dataframe(train_data)

    # 2. Clean Delay Data
    if delay_data is not None:
         delay_data = _normalize_uploaded_dataframe(delay_data)

    # 3. Clean Station Data
    if station_data is not None:
         station_data = _normalize_uploaded_dataframe(station_data)

    # ----------------------------------------------------

    try: sp_km = float(spacing_input)
    except: sp_km = 5.0

    try: max_buff = float(max_buffer_input)
    except: max_buff = 15.0

    infra = Infrastructure(station_data, spacing_km=sp_km)
    base_time = datetime.strptime("08:00", "%H:%M")

    # 1. Nominal
    trains_nom = []
    stations_list = infra.stations

    if train_data is None or train_data.empty:
        return pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(["No train data"]), go.Figure(), build_metric_html(None, None)

    for _, row in train_data.iterrows():
        if pd.isna(row.iloc[0]): continue

        try:
            t_id = str(row['Train ID'])
            direction = row['Direction']
            stops = stations_list[:] if str(direction).strip().lower() == "outbound" else stations_list[::-1]

            t = Train(t_id, direction, row['Priority'], stops, row['Start Time Offset (min)'], row['Speed (km/h)'], row['Dwell (min)'])
            update_train_times(t, infra, base_time)
            trains_nom.append(t)
        except KeyError as e:
            return pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame([f"CSV Error: Column missing - {e}"]), go.Figure(), build_metric_html(None, None)

    # 2. Delayed
    trains_del = [t.copy() for t in trains_nom]

    # --- THE BAIT AND SWITCH: GENERATE RANDOM DATAFRAME IF SELECTED ---
    if d_mode == "Randomized (Monte Carlo)":
        random_delays = []
        for t in trains_del:
            if random.uniform(0, 100) <= float(r_prob):
                if len(t.stops) > 1: # Pick a random station (not the final destination)
                    rand_idx = random.randint(0, len(t.stops) - 2)
                    rand_station = t.stops[rand_idx]
                    d_val = round(random.uniform(float(r_min), float(r_max)), 1)
                    random_delays.append({"Train ID": t.id, "Station": rand_station, "Delay": d_val})

        # Override the delay_data variable so your existing code thinks it came from the CSV!
        if random_delays:
            delay_data = pd.DataFrame(random_delays)
        else:
            delay_data = pd.DataFrame(columns=["Train ID", "Station", "Delay"])
    # ---------------------------------------------------------------

    delay_lookup = {}
    if delay_data is not None and not delay_data.empty:
        for _, row in delay_data.iterrows():
            if pd.isna(row.iloc[0]): continue
            tid = str(row.iloc[0])
            st = str(row.iloc[1])
            try: val = float(row.iloc[2])
            except: val = 0.0
            delay_lookup[(tid, st)] = val

    directions = set(t.direction for t in trains_del)
    for d in directions:
        d_trains = [t for t in trains_del if t.direction == d]
        d_trains.sort(key=lambda x: x.start_time_offset)

        last_departure_time = None
        for t in d_trains:
            if not t.schedule: update_train_times(t, infra, base_time)

            origin_s = t.schedule[0]['station']
            prim_delay_min = delay_lookup.get((t.id, origin_s), 0.0)

            nominal_start = base_time + timedelta(minutes=t.start_time_offset)
            earliest_start = nominal_start + timedelta(minutes=prim_delay_min)

            if last_departure_time:
                h_min = infra.get_dynamic_headway(t.speed_kmh)
                min_safe_start = last_departure_time + h_min
                actual_start = max(earliest_start, min_safe_start)
            else:
                actual_start = earliest_start

            t.schedule[0]['manual_delay'] = actual_start - nominal_start

            for node in t.schedule[1:]:
                st_name = node['station']
                d_val = delay_lookup.get((t.id, st_name), 0.0)
                if d_val > 0:
                    node['manual_delay'] = timedelta(minutes=d_val)

            update_train_times(t, infra, base_time)
            last_departure_time = t.schedule[0]['dep']

    # 3. Adjusted (Resolution with Max Buffer)
    trains_adj = [t.copy() for t in trains_del]

    selected_train_id = None
    if selected_train_input and selected_train_input != "None":
        selected_train_id = str(selected_train_input)

    logs = resolve_all_conflicts(trains_adj, infra, base_time, selected_train_id=selected_train_id, max_buffer_min=max_buff)

    # 4. Outputs
    df_nom_matrix = generate_timetable_matrix(trains_nom, infra)
    df_adj_matrix = generate_timetable_matrix(trains_adj, infra)
    df_stats = calculate_delay_stats(trains_nom, trains_adj, delay_data)
    df_buffer = generate_buffer_report(trains_adj, delay_data)
    fig = create_filtered_marey(trains_nom, trains_del, trains_adj, infra, plot_view, selected_train_id)
    logs_df = pd.DataFrame(logs, columns=["Logs"])
    metric_html = build_metric_html(df_stats, logs_df)

    return df_nom_matrix, df_adj_matrix, df_stats, df_buffer, logs_df, fig, metric_html

def get_train_choices(df):
    if df is None or df.empty:
        return gr.update(choices=["None"], value="None", interactive=True)

    choices = ["None"]
    try:
        # Check potential ID columns
        col = None
        if 'Train ID' in df.columns:
            col = 'Train ID'
        elif len(df.columns) > 0:
            col = df.columns[0]  # Fallback to first column

        if col:
            valid_ids = df[col].dropna().astype(str).str.strip().tolist()
            # Keep the list stable and remove duplicates while preserving order.
            seen = set(choices)
            for train_id in valid_ids:
                if train_id and train_id not in seen:
                    choices.append(train_id)
                    seen.add(train_id)
    except Exception:
        pass

    return gr.update(choices=choices, value="None", interactive=True)

# Helper function to load the trains CSV and refresh the selected-train dropdown at the same time.
def load_train_csv(file):
    df = load_csv(file)
    return df, get_train_choices(df)

# --- PRIMARY DELAY RESPONSIBILITY ---
# Responsibility for a primary delay (IM vs RU) is decided by a single shared
# rule - the "IM Responsibility Share" percentage set before running the
# sweet-spot search - applied identically whether the delay is predefined
# (from the Delays tab) or randomly generated during the run. This keeps the
# assumption consistent across delay modes and across every capacity level
# tested (see simulate_capacity_vs_reliability). If IM is responsible it's a
# penalty for IM; if RU is responsible it's income for IM (RU pays IM).
# Secondary delays are untouched - always a cost paid by IM (PR Penalty).

# Add delay_data to the function inputs
def run_economic_analysis(station_data, fleet_data, delay_data, spacing_input, max_trains, tac_val, pr_val, d_mode, r_prob, r_min, r_max, im_resp_share):
    import pandas as pd
    import plotly.graph_objects as go

    try: sp_km = float(spacing_input)
    except: sp_km = 5.0
    if station_data is not None:
        station_data = _normalize_uploaded_dataframe(station_data)
    if delay_data is not None:
        delay_data = _normalize_uploaded_dataframe(delay_data)
    if fleet_data is not None:
        fleet_data = _normalize_uploaded_dataframe(fleet_data)
    infra = Infrastructure(station_data, spacing_km=sp_km)

    if fleet_data is None or fleet_data.empty:
        # We now return 5 empty/error elements
        return pd.DataFrame([{"Error": "No fleet data defined"}]), go.Figure(), pd.DataFrame(), go.Figure(), pd.DataFrame()

    fleet_data = fleet_data.dropna(how='all')
    fleet_data.columns = fleet_data.columns.astype(str).str.strip()

    fleet_profiles = []
    for _, row in fleet_data.iterrows():
        t_id = str(row.get("Type Name", "Standard")).strip()
        speed = float(row.get("Speed (km/h)", 100))
        dwell = float(row.get("Dwell (min)", 2))
        priority = int(row.get("Priority", 1))

        stops = infra.stations[:]
        fleet_profiles.append(Train(t_id, "Outbound", priority, stops, 0, speed, dwell))

    # Now unpack 5 variables from the simulation engine!
    df, fig, best_timetable_df, best_marey_fig, best_delay_df = simulate_capacity_vs_reliability(
        infra, fleet_profiles, delay_data, max_trains, tac_val, pr_val, d_mode, r_prob, r_min, r_max, im_resp_share
    )

    return df, fig, best_timetable_df, best_marey_fig, best_delay_df

# CSV loading and normalization helpers
# ----------------------------------------------------
# Uploaded files often use slightly different header names, for example
# "Trian ID", "Speed", or "Start time". The simulator expects canonical
# internal names, so normalize them immediately when the CSV is loaded.
def _compact_header(name):
    import re
    return re.sub(r"[^a-z0-9]+", "", str(name).strip().lower())

def _normalize_uploaded_dataframe(df):
    if df is None:
        return None

    df = df.copy()
    df = df.dropna(how='all')
    df.columns = df.columns.astype(str).str.strip()

    compact_cols = [_compact_header(c) for c in df.columns]
    compact_set = set(compact_cols)

    train_markers = {
        "trainid", "trianid", "direction", "priority", "speed", "speedkmh",
        "starttime", "startoffset", "starttimeoffset", "starttimeoffsetmin",
        "dwell", "dwellmin", "dwellminutes"
    }
    station_markers = {"distance", "distancekm", "km", "stationdistance"}
    delay_markers = {"delay", "delaymin", "delayminutes", "delaymins"}

    if compact_set & train_markers:
        csv_kind = "trains"
    elif compact_set & station_markers:
        csv_kind = "stations"
    elif compact_set & delay_markers:
        csv_kind = "delays"
    else:
        csv_kind = "unknown"

    base_aliases = {
        "trainid": "Train ID",
        "trianid": "Train ID",      # common typo in uploaded files
        "trainno": "Train ID",
        "trainnumber": "Train ID",
        "id": "Train ID" if csv_kind == "trains" else "ID",
        "direction": "Direction",
        "dir": "Direction",
        "priority": "Priority",
        "speed": "Speed (km/h)",
        "speedkmh": "Speed (km/h)",
        "speedkmph": "Speed (km/h)",
        "start": "Start Time Offset (min)",
        "starttime": "Start Time Offset (min)",
        "startoffset": "Start Time Offset (min)",
        "starttimeoffset": "Start Time Offset (min)",
        "starttimeoffsetmin": "Start Time Offset (min)",
        "starttimemin": "Start Time Offset (min)",
        "startminutes": "Start Time Offset (min)",
        "dwell": "Dwell (min)",
        "dwelltime": "Dwell (min)",
        "dwellmin": "Dwell (min)",
        "dwellminutes": "Dwell (min)",
        "delay": "Delay (min)",
        "delaymin": "Delay (min)",
        "delaymins": "Delay (min)",
        "delayminutes": "Delay (min)",
        "distance": "Distance (km)",
        "distancekm": "Distance (km)",
        "km": "Distance (km)",
        "typename": "Type Name",
        "type": "Type Name",
    }

    station_aliases = {
        "station": "Station Name" if csv_kind == "stations" else "Station",
        "stationname": "Station Name" if csv_kind == "stations" else "Station",
        "stationnames": "Station Name" if csv_kind == "stations" else "Station",
        "stop": "Station",
        "stops": "Station",
    }

    rename_map = {}
    used = set()
    for original in df.columns:
        key = _compact_header(original)
        new_name = station_aliases.get(key, base_aliases.get(key, original.strip()))
        # Avoid duplicate column names after normalization.
        if new_name in used:
            suffix = 2
            candidate = f"{new_name} {suffix}"
            while candidate in used:
                suffix += 1
                candidate = f"{new_name} {suffix}"
            new_name = candidate
        rename_map[original] = new_name
        used.add(new_name)

    df = df.rename(columns=rename_map)

    # Positional fallbacks for common no/weak-header files.
    if csv_kind == "trains" and "Train ID" not in df.columns and len(df.columns) >= 6:
        df.columns = [
            "Train ID", "Direction", "Priority", "Speed (km/h)",
            "Start Time Offset (min)", "Dwell (min)"
        ] + list(df.columns[6:])
    elif csv_kind == "stations" and "Station Name" not in df.columns and len(df.columns) >= 2:
        df.columns = ["Station Name", "Distance (km)"] + list(df.columns[2:])
    elif csv_kind == "delays" and "Train ID" not in df.columns and len(df.columns) >= 3:
        df.columns = ["Train ID", "Station", "Delay (min)"] + list(df.columns[3:])

    # Reorder known tables so the displayed matrix is consistent.
    desired_order = []
    if csv_kind == "trains":
        desired_order = ["Train ID", "Direction", "Priority", "Speed (km/h)", "Start Time Offset (min)", "Dwell (min)"]
    elif csv_kind == "stations":
        desired_order = ["Station Name", "Distance (km)"]
    elif csv_kind == "delays":
        desired_order = ["Train ID", "Station", "Delay (min)"]

    if desired_order:
        ordered = [c for c in desired_order if c in df.columns]
        extras = [c for c in df.columns if c not in ordered]
        df = df[ordered + extras]

    return df.reset_index(drop=True)

# Helper function to load CSV
def load_csv(file):
    if file is None:
        return None

    file_path = getattr(file, "name", file)
    suffix = str(file_path).lower().rsplit(".", 1)[-1] if "." in str(file_path) else ""

    try:
        if suffix in {"xlsx", "xls"}:
            # .xlsx uses openpyxl in normal pandas installations; .xls uses xlrd when available.
            df = pd.read_excel(file_path)
        else:
            # sep=None lets pandas detect comma/semicolon/tab-delimited CSVs.
            df = pd.read_csv(file_path, sep=None, engine="python")
        return _normalize_uploaded_dataframe(df)
    except Exception:
        try:
            if suffix in {"xlsx", "xls"}:
                df = pd.read_excel(file_path)
            else:
                df = pd.read_csv(file_path)
            return _normalize_uploaded_dataframe(df)
        except Exception:
            return None

d_stations = pd.DataFrame({'Station Name': ['A', 'B', 'C', 'D', 'E'], 'Distance (km)': [0, 15, 35, 45, 60]})
d_trains = pd.DataFrame({
    'Train ID': [101, 102],
    'Direction': ['Outbound', 'Outbound'],
    'Priority': [3, 2],
    'Speed (km/h)': [60, 120],
    'Start Time Offset (min)': [0, 5],
    'Dwell (min)': [2, 2]
})
d_delays = pd.DataFrame({'Train ID': [101], 'Station': ['B'], 'Delay (min)': [10]})


CUSTOM_CSS = """
:root {
    /* Dispatcher-console dark theme: near-black canvas, one locked phosphor-
       green accent, everything set in monospace. Amber is reserved strictly
       for the "conflicts > 0" warning state - never decoration. */
    --hort-bg-1: #141815;   /* card / surface, one step above canvas */
    --hort-bg-2: #0A0C0A;   /* page canvas, near-black (never pure #000) */
    --hort-bg-3: #191F1B;   /* table heads, input fields, hover fills */
    --hort-border: rgba(231, 236, 231, 0.12);
    --hort-text: #E7ECE7;      /* phosphor off-white, never pure #fff */
    --hort-muted: #7C8C82;     /* dim sage-gray secondary text */
    --hort-cyan: #3DDC84;      /* locked accent - phosphor green */
    --hort-cyan-bg: rgba(61, 220, 132, 0.10);
    --hort-amber: #FFB100;     /* semantic warning text (conflicts > 0 only) */
    --hort-amber-bg: rgba(255, 177, 0, 0.10);
}

html, body, gradio-app {
    background: var(--hort-bg-2) !important;
    min-height: 100%;
}

.gradio-container {
    max-width: 1440px !important;
    margin: 0 auto !important;
    color: var(--hort-text) !important;
    background: var(--hort-bg-2) !important;
    font-family: 'JetBrains Mono', 'IBM Plex Mono', 'SF Mono', monospace !important;
    overflow: visible !important;
    position: relative;
}
/* Faint ambient glow pooling behind the top bar - depth without the CRT
   scanline texture, which read more "novelty terminal" than commercial
   product chrome. Fixed + pointer-events:none so it never repaints on scroll. */
.gradio-container::before {
    content: "";
    position: fixed;
    inset: 0;
    z-index: 0;
    pointer-events: none;
    background: radial-gradient(ellipse 900px 420px at 50% -8%, rgba(61,220,132,0.04), transparent 70%);
}

.gradio-container h1, .gradio-container h2, .gradio-container h3, .gradio-container h4,
.gradio-container p, .gradio-container label, .gradio-container .label-wrap,
.gradio-container .prose, .gradio-container .prose *, .gradio-container .markdown,
.gradio-container .markdown *, .gradio-container .gr-markdown, .gradio-container .gr-markdown * {
    color: var(--hort-text) !important;
    font-family: 'JetBrains Mono', 'IBM Plex Mono', 'SF Mono', monospace !important;
}
/* The Soft theme renders every field label as a solid primary-hue chip
   (a Gradio default, not something we set) - strip it back to plain
   small muted text so labels read as UI copy, not colored badges. */
.gradio-container [data-testid="block-info"] {
    background: transparent !important;
    color: var(--hort-muted) !important;
    padding: 0 !important;
    font-size: 12.5px !important;
    font-weight: 500 !important;
    box-shadow: none !important;
}
/* Prose (paragraph copy) reads better in a humanist sans than all-mono -
   headings, labels, data and buttons stay mono for the technical voice. */
.gradio-container .prose p, .gradio-container .markdown p, .gradio-container .gr-markdown p {
    font-family: 'Space Grotesk', 'Helvetica Neue', Arial, sans-serif !important;
    color: var(--hort-muted) !important;
    line-height: 1.7;
    font-size: 14px;
}

/* ---------------------------------------------------------------------
   App shell: a slim brand strip + the tab nav directly beneath it, both
   on the same surface, standing in for a real product top bar. Soft
   14px cards below it instead of last build's sharp terminal edges -
   the corner-radius scale for THIS build is: bar/sidebar flush (0),
   inputs/buttons 8px, cards 14px. Held everywhere, no mixing.
--------------------------------------------------------------------- */
.brand-strip {
    display: flex;
    align-items: center;
    gap: 14px;
    padding: 13px 6px;
    background: var(--hort-bg-1);
    border-bottom: 1px solid var(--hort-border);
}
.brand-mark {
    width: 24px;
    height: 24px;
    border-radius: 7px;
    background: var(--hort-cyan);
    display: flex;
    align-items: center;
    justify-content: center;
    flex: 0 0 auto;
    box-shadow: 0 0 14px rgba(61,220,132,0.35);
}
.brand-mark svg { color: #0A0C0A; }
.brand-name {
    font-family: 'JetBrains Mono', monospace;
    font-weight: 600;
    font-size: 14px;
    letter-spacing: -0.01em;
    color: var(--hort-text) !important;
}
.brand-spacer { flex: 1 1 auto; }
.status-pill {
    display: inline-flex;
    align-items: center;
    gap: 6px;
    background: var(--hort-cyan-bg);
    border-radius: 999px;
    padding: 5px 10px 5px 8px;
}
.status-dot {
    width: 6px; height: 6px;
    border-radius: 50%;
    background: var(--hort-cyan);
    box-shadow: 0 0 5px rgba(61,220,132,0.8);
}
.status-pill span.label { font-size: 11px; font-weight: 600; color: var(--hort-cyan); }
.avatar-chip {
    width: 28px; height: 28px;
    border-radius: 50%;
    background: var(--hort-bg-3);
    border: 1px solid var(--hort-border);
    display: flex; align-items: center; justify-content: center;
    font-size: 10.5px;
    color: var(--hort-muted);
}

/* KPI row: real stat cards (icon + label, big number, context line)
   instead of the previous bracket-hairline readout. */
.kpi-grid {
    display: grid;
    grid-template-columns: repeat(4, minmax(0, 1fr));
    gap: 16px;
    margin: 0 0 22px 0;
}
.kpi-card {
    background: var(--hort-bg-1);
    border: 1px solid var(--hort-border);
    border-radius: 14px;
    padding: 18px 20px;
    box-shadow: 0 20px 40px -28px rgba(0,0,0,0.6), inset 0 1px 0 rgba(231,236,231,0.03);
    display: flex;
    flex-direction: column;
    gap: 12px;
}
.kpi-head { display: flex; align-items: center; gap: 10px; }
.kpi-icon {
    width: 30px; height: 30px;
    border-radius: 8px;
    display: flex; align-items: center; justify-content: center;
    flex: 0 0 auto;
    background: var(--hort-cyan-bg);
}
.kpi-icon.is-flagged { background: var(--hort-amber-bg); }
.kpi-icon.is-neutral { background: rgba(231,236,231,0.06); }
.kpi-icon svg { color: var(--hort-cyan); }
.kpi-icon.is-flagged svg { color: var(--hort-amber); }
.kpi-icon.is-neutral svg { color: var(--hort-muted); }
.kpi-label { font-size: 12.5px; color: var(--hort-muted); }
.kpi-value {
    font-family: 'JetBrains Mono', monospace;
    color: var(--hort-text) !important;
    font-size: 26px;
    font-weight: 600;
}
.kpi-value.is-flagged { color: var(--hort-amber) !important; }
.kpi-value .unit { font-size: 14px; font-weight: 500; color: var(--hort-muted); }
.kpi-caption { font-size: 11.5px; color: var(--hort-muted); margin-top: 2px; }
@media (max-width: 900px) {
    .kpi-grid { grid-template-columns: repeat(2, minmax(0, 1fr)); }
}
@media (max-width: 480px) {
    .kpi-grid { grid-template-columns: 1fr; }
}

/* Sidebar (run configuration) + main content columns. */
.app-sidebar {
    background: var(--hort-bg-1);
    border-right: 1px solid var(--hort-border);
    padding: 22px 20px !important;
    display: flex !important;
    flex-direction: column !important;
}
.app-sidebar .side-label {
    font-size: 11px;
    font-weight: 600;
    color: var(--hort-muted) !important;
    text-transform: uppercase;
    letter-spacing: .04em;
    margin: 4px 0 2px;
}
.sidebar-spacer { flex: 1 1 auto; min-height: 12px; }
.app-main { padding: 24px 28px !important; }
.data-row-card { min-height: 0; }

/* Raw editable tables tuck behind a collapsed accordion by default, so the
   main pane reads as header + KPI row + one results panel, not a stack of
   six cards. */
.data-accordion {
    background: transparent !important;
    border: 1px dashed var(--hort-border) !important;
    border-radius: 10px !important;
    margin-bottom: 20px !important;
}
.data-accordion > .label-wrap,
.data-accordion .label-wrap button {
    font-size: 13px !important;
    color: var(--hort-muted) !important;
    font-family: 'Space Grotesk', 'Helvetica Neue', Arial, sans-serif !important;
}

/* Section-header icon + label, replacing emoji glyphs. Bolder stroke reads
   closer to a Phosphor "Bold" weight than a generic thin-line icon set. */
.hort-heading {
    display: flex !important;
    align-items: center;
    gap: 10px;
    font-weight: 600 !important;
    text-transform: uppercase;
    font-size: 14px !important;
    letter-spacing: .03em;
}
.hort-heading .hort-icon {
    flex: 0 0 auto;
    color: var(--hort-cyan);
    opacity: .9;
}

/* ---------------------------------------------------------------------
   Cards: dark surface, hairline border, sharp 4px radius, top bezel
   highlight for depth - a green-tinted glow appears on hover only.
--------------------------------------------------------------------- */
.panel-card, .chart-card, .results-card {
    position: relative !important;
    overflow: visible !important;
    background: var(--hort-bg-1) !important;
    border: 1px solid var(--hort-border) !important;
    border-radius: 14px !important;
    padding: 24px !important;
    box-shadow: inset 0 1px 0 rgba(231,236,231,0.04) !important;
    backdrop-filter: none !important;
    -webkit-backdrop-filter: none !important;
}
.chart-card { min-height: 720px; }
.panel-card, .results-card {
    transition: border-color .2s ease, box-shadow .2s ease;
}
.panel-card:hover, .results-card:hover {
    border-color: rgba(61,220,132,0.30) !important;
    box-shadow: inset 0 1px 0 rgba(231,236,231,0.04), 0 0 24px rgba(61,220,132,0.06) !important;
}
.panel-card h3, .chart-card h3, .results-card h3 {
    color: var(--hort-text) !important;
    font-weight: 600 !important;
    margin-top: 0 !important;
}

/* Keep the visual layer smooth while allowing dropdown menus to open. */
.gradio-container .block,
.gradio-container .form,
.gradio-container .panel,
.gradio-container .gr-box,
.gradio-container .contain,
.gradio-container .wrap {
    border-color: var(--hort-border) !important;
    border-radius: 8px !important;
    box-shadow: none !important;
    overflow: visible !important;
}
.panel-card .block, .panel-card .form, .panel-card .gr-box, .panel-card .contain, .panel-card .wrap,
.chart-card .block, .chart-card .form, .chart-card .gr-box, .chart-card .contain, .chart-card .wrap,
.results-card .block, .results-card .form, .results-card .gr-box, .results-card .contain, .results-card .wrap {
    background: transparent !important;
    border-color: var(--hort-border) !important;
}

.gradio-container input,
.gradio-container textarea,
.gradio-container select,
.gradio-container button {
    border-radius: 8px !important;
    font-family: 'JetBrains Mono', 'IBM Plex Mono', 'SF Mono', monospace !important;
}
.gradio-container input,
.gradio-container textarea,
.gradio-container select {
    background: var(--hort-bg-3) !important;
    color: var(--hort-text) !important;
    border: 1px solid var(--hort-border) !important;
    box-shadow: none !important;
}
.gradio-container input:focus,
.gradio-container textarea:focus,
.gradio-container select:focus {
    border-color: var(--hort-cyan) !important;
    box-shadow: 0 0 0 2px rgba(61, 220, 132, 0.14) !important;
}

/* View & Optimization selectors
   Keep dropdowns native. Do not style Gradio's internal combobox/listbox/menu/button
   elements; doing so can make the menu jump, flicker, or open in the wrong place. */
.hort-listbox {
    border-radius: 8px !important;
}
.hort-listbox label,
.hort-listbox .label-wrap {
    color: var(--hort-text) !important;
    font-weight: 600 !important;
}
.selector-strip {
    padding: 4px 0 2px 0 !important;
    overflow: visible !important;
    position: relative !important;
    z-index: 20 !important;
}
#view_listbox, #train_listbox {
    position: relative !important;
    z-index: 30 !important;
}
/* Important: do not apply blur/filter/transform around Gradio dropdowns.
   These can offset the popover menu and make it appear in the wrong section. */
.hort-view-select, .hort-train-select {
    backdrop-filter: none !important;
    -webkit-backdrop-filter: none !important;
    filter: none !important;
    transform: none !important;
}

.gradio-container input[type="radio"],
.gradio-container input[type="checkbox"] {
    accent-color: var(--hort-cyan) !important;
}
.gradio-container fieldset {
    background: transparent !important;
    border: 1px solid var(--hort-border) !important;
    border-radius: 8px !important;
}
.hort-radio fieldset {
    padding: 6px !important;
    display: flex !important;
    gap: 6px !important;
}
.hort-radio label {
    display: flex !important;
    align-items: center !important;
    gap: 8px;
    margin: 0 !important;
    padding: 10px 14px !important;
    border-radius: 6px !important;
    background: transparent !important;
    border: 1px solid transparent !important;
    color: var(--hort-muted) !important;
    font-weight: 500 !important;
    transition: all .15s ease !important;
}
.hort-radio label:hover {
    background: var(--hort-bg-3) !important;
}
.hort-radio label:has(input[type="radio"]:checked),
.hort-radio label[aria-checked="true"],
.hort-radio .selected {
    background: var(--hort-cyan-bg) !important;
    border: 1px solid rgba(61,220,132,0.35) !important;
    border-radius: 6px !important;
    box-shadow: 0 0 12px rgba(61,220,132,0.08) !important;
    color: var(--hort-cyan) !important;
    font-weight: 700 !important;
}

/* File input laid out as one horizontal row:
   label | compact file picker | helper text. */
.input-file-example {
    width: 100% !important;
    max-width: none !important;
    margin: 0 !important;
    padding: 0 !important;
    display: flex !important;
    flex-direction: row !important;
    align-items: center !important;
    justify-content: flex-start !important;
    flex-wrap: nowrap !important;
    gap: 0.75rem !important;
    background: transparent !important;
    border: 0 !important;
    box-shadow: none !important;
}

/* Keep Gradio's child wrappers from stretching each item onto its own line. */
.input-file-example > * {
    margin: 0 !important;
}
.input-file-example .input-file-label-wrap {
    flex: 0 0 auto !important;
    width: auto !important;
    min-width: max-content !important;
}
.input-file-example .input-file-control {
    flex: 0 0 20rem !important;
    width: 20rem !important;
    max-width: 20rem !important;
    min-width: 16rem !important;
}
.input-file-example .input-file-help-wrap {
    flex: 1 1 auto !important;
    width: auto !important;
    min-width: 0 !important;
}

.input-file-example .input-file-label {
    display: block !important;
    margin: 0 !important;
    color: var(--hort-text) !important;
    font-size: 0.8rem !important;
    line-height: 1.25rem !important;
    font-weight: 500 !important;
    white-space: nowrap !important;
}

/* Gradio component shell: compact native-style file field. */
.input-file-control {
    margin: 0 !important;
    padding: 0 !important;
    background: transparent !important;
    border: 0 !important;
    box-shadow: none !important;
}
.input-file-control > div,
.input-file-control .wrap,
.input-file-control .block,
.input-file-control .container {
    width: 100% !important;
    min-width: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
    background: transparent !important;
    border: 0 !important;
    box-shadow: none !important;
}

/* Empty-state uploader. Gradio still handles click/drop/upload; only its appearance changes. */
.input-file-control button[aria-label*="upload" i],
.input-file-control button[aria-label*="drop" i] {
    width: 100% !important;
    height: 2.25rem !important;      /* compact Input height: 36px */
    min-height: 2.25rem !important;
    max-height: 2.25rem !important;
    padding: 0 0.75rem !important;
    display: flex !important;
    align-items: center !important;
    justify-content: flex-start !important;
    gap: 0 !important;
    overflow: hidden !important;
    border: 1px solid var(--hort-border) !important;
    border-style: solid !important;
    border-radius: 6px !important;
    background: var(--hort-bg-3) !important;
    color: var(--hort-text) !important;
    box-shadow: none !important;
    outline: none !important;
    cursor: pointer !important;
    transition: border-color 120ms ease, box-shadow 120ms ease !important;
}

/* Hide Gradio's upload icon / drag-drop copy and render the native-file-input wording. */
.input-file-control button[aria-label*="upload" i] > *,
.input-file-control button[aria-label*="drop" i] > * {
    display: none !important;
}
.input-file-control button[aria-label*="upload" i]::before,
.input-file-control button[aria-label*="drop" i]::before {
    content: "Choose File";
    display: inline-block;
    flex: 0 0 auto;
    margin-right: 0.75rem;
    color: var(--hort-text);
    font-size: 0.8rem;
    line-height: 2rem;
    font-weight: 500;
    white-space: nowrap;
}
.input-file-control button[aria-label*="upload" i]:focus-visible,
.input-file-control button[aria-label*="drop" i]:focus-visible {
    border-color: var(--hort-cyan) !important;
    box-shadow: 0 0 0 2px rgba(61, 220, 132, 0.14) !important;
}

/* Selected-file state: keep the control compact instead of expanding into a preview card. */
.input-file-control .file,
.input-file-control .file-preview,
.input-file-control [class*="file-preview"] {
    min-height: 2.25rem !important;
    height: 2.25rem !important;
    max-height: 2.25rem !important;
    width: 100% !important;
    padding: 0 0.75rem !important;
    display: flex !important;
    align-items: center !important;
    overflow: hidden !important;
    border: 1px solid var(--hort-border) !important;
    border-radius: 6px !important;
    background: var(--hort-bg-3) !important;
    color: var(--hort-text) !important;
    box-shadow: none !important;
}
.input-file-control .file-name,
.input-file-control [class*="file-name"] {
    min-width: 0 !important;
    overflow: hidden !important;
    text-overflow: ellipsis !important;
    white-space: nowrap !important;
    color: var(--hort-text) !important;
    font-size: 0.8rem !important;
    font-weight: 400 !important;
}
.input-file-control .file svg,
.input-file-control .file-preview svg {
    display: none !important;
}

.input-file-example .input-file-help {
    margin: 0 !important;
    color: var(--hort-muted) !important;
    font-size: 0.75rem !important;
    line-height: 1rem !important;
    font-weight: 400 !important;
    white-space: nowrap !important;
}

.gradio-container table,
.gradio-container thead,
.gradio-container tbody,
.gradio-container td,
.gradio-container th {
    color: var(--hort-text) !important;
    border-color: var(--hort-border) !important;
    font-family: 'JetBrains Mono', 'IBM Plex Mono', 'SF Mono', monospace !important;
}
.gradio-container th {
    background: var(--hort-bg-3) !important;
    color: var(--hort-cyan) !important;
    font-weight: 600 !important;
    font-size: 12px !important;
}
.gradio-container td {
    background: var(--hort-bg-1) !important;
    font-size: 13px !important;
}
.gradio-container td,
.gradio-container th {
    padding-top: 10px !important;
    padding-bottom: 10px !important;
}
.gradio-container .gradio-dataframe,
.gradio-container [data-testid="dataframe"],
.gradio-container .table-wrap {
    border: 1px solid var(--hort-border) !important;
    box-shadow: none !important;
    background: var(--hort-bg-1) !important;
    border-radius: 10px !important;
    overflow: auto !important;
}
/* Do not clip uploaded CSV matrices. Gradio virtualizes dataframe rows inside
   scrollable wrappers; forcing overflow:hidden can hide rows after upload. */
.gradio-container [data-testid="dataframe"] > div,
.gradio-container .gradio-dataframe > div,
.gradio-container .table-wrap > div {
    overflow: auto !important;
}
.gradio-container [data-testid="dataframe"] table,
.gradio-container .gradio-dataframe table,
.gradio-container .table-wrap table {
    overflow: visible !important;
}
.gradio-container .gradio-dataframe button[title*="fullscreen" i],
.gradio-container .gradio-dataframe button[aria-label*="fullscreen" i],
.gradio-container [data-testid="dataframe"] button[title*="fullscreen" i],
.gradio-container [data-testid="dataframe"] button[aria-label*="fullscreen" i] {
    display: none !important;
}

.gradio-container button {
    font-weight: 600 !important;
    font-family: 'Space Grotesk', 'Helvetica Neue', Arial, sans-serif !important;
    font-size: 14px !important;
}
.gradio-container button.primary {
    color: #0A0C0A !important;
    background: var(--hort-cyan) !important;
    border: 1px solid var(--hort-cyan) !important;
    border-radius: 8px !important;
    box-shadow: 0 0 0 rgba(61,220,132,0) !important;
    transition: background-color .15s ease, transform .1s ease, box-shadow .2s ease !important;
}
.gradio-container button.primary:hover {
    background: #57E89C !important;
    border-color: #57E89C !important;
    box-shadow: 0 0 20px rgba(61,220,132,0.35) !important;
}
.gradio-container button.primary:active {
    transform: scale(0.98);
}

/* Underline tabs - the second tier of the top bar, directly under the
   brand strip (same background, hairline bottom border). */
.gradio-container .tab-nav {
    background: var(--hort-bg-1) !important;
    border: none !important;
    border-bottom: 1px solid var(--hort-border) !important;
    border-radius: 0 !important;
    padding: 0 !important;
    gap: 22px;
}
.gradio-container .tab-nav button {
    border-radius: 0 !important;
    background: transparent !important;
    color: var(--hort-muted) !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    margin-bottom: -1px;
    font-family: 'Space Grotesk', 'Helvetica Neue', Arial, sans-serif !important;
    transition: color .15s ease;
}
.gradio-container .tab-nav button:hover {
    color: var(--hort-text) !important;
}
.gradio-container .tab-nav button.selected {
    color: var(--hort-text) !important;
    background: transparent !important;
    border-bottom: 2px solid var(--hort-cyan) !important;
    box-shadow: none !important;
    font-weight: 600 !important;
}
.hort-main-tabs > .tab-nav,
.gradio-container .hort-main-tabs .tab-nav {
    margin-bottom: 22px !important;
    overflow: visible !important;
    padding: 0 6px !important;
}
.hort-main-tabs > .tab-nav button,
.gradio-container .hort-main-tabs .tab-nav button {
    min-height: 44px !important;
    padding: 10px 2px !important;
    font-size: 14px !important;
    font-weight: 500 !important;
}

.plot-container, .js-plotly-plot, .plotly, .gr-plot {
    border-radius: 10px !important;
    overflow: hidden !important;
}

/* Keep the space-time diagram on the same dark surface as its card, while
   allowing Plotly's built-in download/export to use a transparent figure. */
.chart-card .gr-plot,
.chart-card .plot-container,
.chart-card .js-plotly-plot,
.chart-card .plotly {
    background: var(--hort-bg-1) !important;
}
.footer-note {
    text-align: center;
    color: var(--hort-muted) !important;
    font-size: 12px;
    margin: 28px 0 8px 0;
}
footer { display: none !important; }

/* Dropdown overlay stability fix for Hugging Face Spaces / Gradio:
   remove visual effects that can create a wrong containing block for the popup menu. */
.panel-card *, .results-card *, .chart-card * {
    scroll-margin-top: 0;
}

@media (max-width: 900px) {
    .app-sidebar { border-right: none; border-bottom: 1px solid var(--hort-border); }
}
@media (max-width: 560px) {
    .panel-card, .chart-card, .results-card { padding: 16px !important; }
    .kpi-value { font-size: 20px; }
}
"""

# Dispatcher-console type system: JetBrains Mono for system chrome (headings,
# labels, data, buttons, tabs) and Space Grotesk for human-readable prose
# (hero subtitle, section descriptions) - mono everywhere reads efficient but
# hurts paragraph legibility, so prose gets a companion sans built for the
# same technical mood. Loaded via <head> instead of a CSS @import so it
# resolves before first paint on Hugging Face Spaces.
CUSTOM_HEAD = """
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=JetBrains+Mono:ital,wght@0,400;0,500;0,600;0,700;1,400&family=IBM+Plex+Mono:wght@400;500;600&family=Space+Grotesk:wght@400;500;600&display=swap" rel="stylesheet">
"""

try:
    HORT_THEME = gr.themes.Soft(
        primary_hue="green",
        secondary_hue="stone",
        neutral_hue="stone",
        radius_size="sm",
        spacing_size="md",
    )
except Exception:
    HORT_THEME = None

def make_exact_file_input():
    """Render label, file picker, and helper text on one horizontal line."""
    with gr.Row(elem_classes=["input-file-example"]):
        gr.HTML(
            '<label class="input-file-label">Upload a file</label>',
            elem_classes=["input-file-label-wrap"]
        )
        file_component = gr.File(
            label="Upload a file",
            show_label=False,
            file_types=[".csv", ".xlsx", ".xls"],
            height=36,
            elem_classes=["input-file-control"]
        )
        gr.HTML(
            '<p class="input-file-help">You are only allowed to upload .CSV, .XLSX or .XLS files.</p>',
            elem_classes=["input-file-help-wrap"]
        )
    return file_component


with gr.Blocks(title="Rail Delay Lab", theme=HORT_THEME, css=CUSTOM_CSS) as demo:
    gr.HTML(f"""
    <div class="brand-strip">
        <div class="brand-mark">{icon_svg('train', 13)}</div>
        <span class="brand-name">Rail Delay Lab</span>
        <div class="brand-spacer"></div>
        <div class="status-pill">
            <span class="status-dot"></span>
            <span class="label">Live</span>
        </div>
        <div class="avatar-chip">OP</div>
    </div>
    """)

    with gr.Tabs(elem_classes=["hort-main-tabs"]):
        with gr.TabItem("Delay Propagation and Conflict Resolution"):
            gr.Markdown("## Delay Propagation and Conflict Resolution")
            gr.Markdown("Upload CSV files to populate the tables below, or edit them manually, then run the resolver.")

            with gr.Row(equal_height=True):
                with gr.Column(scale=0, min_width=300, elem_classes=["app-sidebar"]):
                    gr.HTML('<p class="side-label">Infrastructure</p>')
                    txt_spacing = gr.Number(value=5.0, label="Min Safety Spacing (km)", step=0.1)

                    gr.HTML('<p class="side-label">Delay Injection</p>')
                    main_delay_mode = gr.Radio(
                        choices=["Use Uploaded/Inserted Delay", "Randomized (Monte Carlo)"],
                        value="Use Uploaded/Inserted Delay",
                        label="Simulation Delay Mode",
                        elem_classes=["hort-radio"]
                    )
                    main_rand_prob = gr.Slider(0, 100, value=25, step=5, label="Delay Probability (%)")
                    main_min_delay = gr.Number(value=5, label="Min Delay (min)")
                    main_max_delay = gr.Number(value=15, label="Max Delay (min)")

                    gr.HTML('<p class="side-label">View &amp; Optimization</p>')
                    with gr.Column(elem_classes=["selector-strip"]):
                        dd_view = gr.Dropdown(
                            choices=["All Scenarios", "Nominal Only", "Conflicted Only", "Resolved Only"],
                            value="All Scenarios",
                            label="Select View",
                            interactive=True,
                            filterable=False,
                            allow_custom_value=False,
                            elem_id="view_listbox",
                            elem_classes=["hort-listbox", "hort-view-select"]
                        )
                        dd_selected_train = gr.Dropdown(
                            choices=["None", "101", "102"],
                            value="None",
                            label="Select Train",
                            interactive=True,
                            filterable=False,
                            allow_custom_value=False,
                            elem_id="train_listbox",
                            elem_classes=["hort-listbox", "hort-train-select"]
                        )
                    txt_max_buffer = gr.Number(value=15.0, label="Max Buffer Cap (min)", step=1.0)

                    gr.HTML('<div class="sidebar-spacer"></div>')
                    btn = gr.Button("Simulate", variant="primary", size="lg")

                with gr.Column(scale=1, elem_classes=["app-main"]):
                    metric_strip = gr.HTML(build_metric_html(None, None))

                    with gr.Accordion("Edit input data", open=False, elem_classes=["data-accordion"]):
                        with gr.Row():
                            with gr.Column():
                                with gr.Group(elem_classes=["panel-card", "data-row-card"]):
                                    gr.HTML(section_header("building", "Infrastructure Data"))
                                    file_infra = make_exact_file_input()
                                    t_stations = gr.Dataframe(value=d_stations, label="Stations Data", interactive=True, max_height=300, wrap=True, show_row_numbers=True, buttons=[])
                                    file_infra.upload(fn=load_csv, inputs=file_infra, outputs=t_stations)
                            with gr.Column():
                                with gr.Group(elem_classes=["panel-card", "data-row-card"]):
                                    gr.HTML(section_header("clock", "Delay Data"))
                                    file_delays = make_exact_file_input()
                                    t_delays = gr.Dataframe(value=d_delays, label="Delay Data", interactive=True, max_height=300, wrap=True, show_row_numbers=True, buttons=[])
                                    file_delays.upload(fn=load_csv, inputs=file_delays, outputs=t_delays)

                        with gr.Group(elem_classes=["panel-card"]):
                            gr.HTML(section_header("train", "Train Requests"))
                            file_trains = make_exact_file_input()
                            t_trains = gr.Dataframe(value=d_trains, label="Train Requests Data", interactive=True, max_height=320, wrap=True, show_row_numbers=True, buttons=[])
                            file_trains.upload(fn=load_train_csv, inputs=file_trains, outputs=[t_trains, dd_selected_train])

                    with gr.Group(elem_classes=["chart-card"]):
                        gr.HTML(section_header("chart", "Operational Space-Time Diagram"))
                        plot = gr.Plot(label="Marey Diagram")

                        gr.HTML(section_header("list", "Delay & Conflict Outputs"))
                        with gr.Tabs():
                            with gr.TabItem("Resolved Timetable"):
                                out_adj = gr.Dataframe(max_height=460, wrap=True, show_row_numbers=True, buttons=[])
                            with gr.TabItem("Nominal Timetable"):
                                out_nom = gr.Dataframe(max_height=460, wrap=True, show_row_numbers=True, buttons=[])
                            with gr.TabItem("Delay Breakdown"):
                                out_stats = gr.Dataframe(max_height=460, wrap=True, show_row_numbers=True, buttons=[])
                            with gr.TabItem("Buffer Report"):
                                out_buffer = gr.Dataframe(max_height=460, wrap=True, show_row_numbers=True, buttons=[])
                            with gr.TabItem("Conflict Logs"):
                                out_logs = gr.Dataframe(max_height=460, wrap=True, show_row_numbers=True, buttons=[])

        with gr.TabItem("Economic Capacity Analysis"):
            gr.Markdown("## Economic Capacity Analysis")
            gr.Markdown("Evaluate the capacity/reliability trade-off, TAC revenue, performance-regime penalties, and the best operating sweet spot.")

            with gr.Row(equal_height=True):
                with gr.Column(scale=0, min_width=300, elem_classes=["app-sidebar"]):
                    gr.HTML('<p class="side-label">Infrastructure</p>')
                    econ_txt_spacing = gr.Number(value=5.0, label="Min Safety Spacing (km)", step=0.1)

                    gr.HTML('<p class="side-label">Delay Injection Strategy</p>')
                    delay_mode = gr.Radio(
                        choices=["Predefined (from Delays Tab)", "Randomized"],
                        value="Predefined (from Delays Tab)",
                        label="Delay Mode",
                        elem_classes=["hort-radio"]
                    )
                    rand_prob = gr.Slider(0, 100, value=25, step=5, label="Probability of Delay per Train (%)")
                    min_delay = gr.Number(value=5, label="Min Random Delay (min)")
                    max_delay = gr.Number(value=15, label="Max Random Delay (min)")

                    gr.HTML('<p class="side-label">Primary Delay Responsibility</p>')
                    gr.Markdown("_Share of every primary delay attributed to the Infrastructure Manager (IM) vs the Railway Undertaking (RU), applied before the search runs._")
                    im_resp_share = gr.Slider(0, 100, value=50, step=5, label="IM Responsibility Share (%)")

                    gr.HTML('<p class="side-label">Economic &amp; Search Parameters</p>')
                    max_trains_in = gr.Number(value=25, label="Max Trains to Evaluate")
                    num_tac = gr.Number(value=115.84, label="TAC Fee (€ per train)")
                    num_pr = gr.Number(value=7.09, label="PR Penalty (€ per min)")

                    gr.HTML('<div class="sidebar-spacer"></div>')
                    btn_econ = gr.Button("Calculate Sweet Spot", variant="primary", size="lg")

                with gr.Column(scale=1, elem_classes=["app-main"]):
                    with gr.Accordion("Edit input data", open=False, elem_classes=["data-accordion"]):
                        with gr.Row():
                            with gr.Column():
                                with gr.Group(elem_classes=["panel-card", "data-row-card"]):
                                    gr.HTML(section_header("building", "Infrastructure Data"))
                                    econ_file_infra = make_exact_file_input()
                                    econ_t_stations = gr.Dataframe(value=d_stations, label="Stations Data", interactive=True, max_height=300, wrap=True, show_row_numbers=True, buttons=[])
                                    econ_file_infra.upload(fn=load_csv, inputs=econ_file_infra, outputs=econ_t_stations)
                            with gr.Column():
                                with gr.Group(elem_classes=["panel-card", "data-row-card"]):
                                    gr.HTML(section_header("clock", "Delay Data"))
                                    econ_file_delays = make_exact_file_input()
                                    econ_t_delays = gr.Dataframe(value=d_delays, label="Delay Data", interactive=True, max_height=300, wrap=True, show_row_numbers=True, buttons=[])
                                    econ_file_delays.upload(fn=load_csv, inputs=econ_file_delays, outputs=econ_t_delays)

                    with gr.Group(elem_classes=["panel-card"]):
                        gr.HTML(section_header("train", "Fleet Mix"))
                        t_fleet = gr.Dataframe(
                            headers=["Type Name", "Speed (km/h)", "Dwell (min)", "Priority"],
                            value=[["High-Speed", 130, 2, 1], ["Regional", 100, 4, 2], ["Freight", 80, 5, 3]],
                            interactive=True,
                            label="Fleet Composition",
                            max_height=320,
                            wrap=True,
                            show_row_numbers=True,
                            buttons=[]
                        )

                    with gr.Group(elem_classes=["chart-card"]):
                        gr.HTML(section_header("bars", "Economic Capacity Outputs"))
                        with gr.Row():
                            out_econ_plot = gr.Plot(label="Profit Curve")
                            out_best_marey = gr.Plot(label="Optimal Timetable (Space-Time Diagram)")
                        with gr.Row():
                            out_econ_df = gr.Dataframe(label="Raw Economic Data", max_height=460, wrap=True, show_row_numbers=True, buttons=[])
                            out_best_timetable = gr.Dataframe(label="Optimal Timetable (Data)", max_height=460, wrap=True, show_row_numbers=True, buttons=[])
                        gr.Markdown("### Best Scenario Primary Delays Applied (with Responsibility)")
                        out_econ_delays = gr.Dataframe(label="Primary Delays & Responsibility (Best Scenario)", max_height=460, wrap=True, show_row_numbers=True, buttons=[])

    gr.HTML("<div class='footer-note'>HoRT RailCap Lab · Educational railway capacity and reliability simulator</div>")

    t_trains.change(fn=get_train_choices, inputs=[t_trains], outputs=[dd_selected_train])

    btn.click(
        fn=run_simulation,
        inputs=[
            t_stations, t_trains, t_delays, txt_spacing, dd_view, dd_selected_train, txt_max_buffer,
            main_delay_mode, main_rand_prob, main_min_delay, main_max_delay
        ],
        outputs=[out_nom, out_adj, out_stats, out_buffer, out_logs, plot, metric_strip]
    )

    dd_view.change(
        fn=run_simulation,
        inputs=[
            t_stations, t_trains, t_delays, txt_spacing, dd_view, dd_selected_train, txt_max_buffer,
            main_delay_mode, main_rand_prob, main_min_delay, main_max_delay
        ],
        outputs=[out_nom, out_adj, out_stats, out_buffer, out_logs, plot, metric_strip]
    )

    dd_selected_train.change(
        fn=run_simulation,
        inputs=[
            t_stations, t_trains, t_delays, txt_spacing, dd_view, dd_selected_train, txt_max_buffer,
            main_delay_mode, main_rand_prob, main_min_delay, main_max_delay
        ],
        outputs=[out_nom, out_adj, out_stats, out_buffer, out_logs, plot, metric_strip]
    )

    txt_max_buffer.change(
        fn=run_simulation,
        inputs=[
            t_stations, t_trains, t_delays, txt_spacing, dd_view, dd_selected_train, txt_max_buffer,
            main_delay_mode, main_rand_prob, main_min_delay, main_max_delay
        ],
        outputs=[out_nom, out_adj, out_stats, out_buffer, out_logs, plot, metric_strip]
    )

    main_delay_mode.change(
        fn=run_simulation,
        inputs=[
            t_stations, t_trains, t_delays, txt_spacing, dd_view, dd_selected_train, txt_max_buffer,
            main_delay_mode, main_rand_prob, main_min_delay, main_max_delay
        ],
        outputs=[out_nom, out_adj, out_stats, out_buffer, out_logs, plot, metric_strip]
    )

    btn_econ.click(
        fn=run_economic_analysis,
        inputs=[
            econ_t_stations, t_fleet, econ_t_delays, econ_txt_spacing, max_trains_in,
            num_tac, num_pr, delay_mode, rand_prob, min_delay, max_delay, im_resp_share
        ],
        outputs=[out_econ_df, out_econ_plot, out_best_timetable, out_best_marey, out_econ_delays]
    )

# Launch the app. Gradio 6 moved css/theme from Blocks() to launch(), while
# Gradio 4/5 accepts them in Blocks(). This keeps the visual theme working on
# both current and older Hugging Face Spaces runtimes.
if __name__ == "__main__":
    import inspect
    launch_kwargs = {"debug": False}
    if "css" in inspect.signature(demo.launch).parameters:
        launch_kwargs.update({"css": CUSTOM_CSS, "theme": HORT_THEME})
    if "head" in inspect.signature(demo.launch).parameters:
        launch_kwargs.update({"head": CUSTOM_HEAD})
    demo.launch(**launch_kwargs)