app.py

import pandas as pd
import streamlit as st
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import io
import base64
import math
from matplotlib.backends.backend_pdf import PdfPages

# --- Constants ---
SPLIT_TIME = "17:30"
BUSINESS_START = "09:30"
BUSINESS_END = "01:30"
BORDER_COLOR = 'grey' # Changed to grey for the new design
DATE_COLOR = '#A9A9A9'

def process_schedule(file):
    """
    Processes the uploaded Excel file to extract, clean, and sort screening times.
    """
    try:
        # Read Excel, skipping header rows
        df = pd.read_excel(file, skiprows=8)

        # Extract required columns (G, H, J)
        df = df.iloc[:, [6, 7, 9]]
        df.columns = ['Hall', 'StartTime', 'EndTime']

        # Clean data: drop rows with missing values
        df = df.dropna(subset=['Hall', 'StartTime', 'EndTime'])

        # Format Hall to "# " format
        df['Hall'] = df['Hall'].str.extract(r'(\d+)').astype(str) + ' '

        # Convert time columns to datetime objects
        base_date = datetime.today().date()
        df['StartTime'] = pd.to_datetime(df['StartTime'])
        df['EndTime'] = pd.to_datetime(df['EndTime'])

        # --- Handle overnight screenings ---
        # If a show ends after midnight (e.g., 1:30 AM), it belongs to the previous day's schedule.
        # We handle this by adding a day to its datetime object.
        for idx, row in df.iterrows():
            if row['EndTime'].hour < 9: # Assuming any end time before 9 AM is part of the previous night
                df.at[idx, 'EndTime'] = row['EndTime'] + timedelta(days=1)

        # Create a comparable time column that correctly handles the business day logic
        df['time_for_comparison'] = df['EndTime']

        # Sort screenings by their end time
        df = df.sort_values('EndTime')

        # Split data into day and night shifts
        split_datetime = datetime.combine(base_date, datetime.strptime(SPLIT_TIME, "%H:%M").time())
        part1 = df[df['time_for_comparison'] <= split_datetime].copy()
        part2 = df[df['time_for_comparison'] > split_datetime].copy()

        # Format the time display string (e.g., "5:30")
        for part in [part1, part2]:
            part['EndTime'] = part['EndTime'].dt.strftime('%-I:%M')

        # Precisely read the date from cell C6
        date_df = pd.read_excel(file, skiprows=5, nrows=1, usecols=[2], header=None)
        date_cell = date_df.iloc[0, 0]

        try:
            if isinstance(date_cell, str):
                date_str = datetime.strptime(date_cell, '%Y-%m-%d').strftime('%Y-%m-%d')
            else:
                date_str = pd.to_datetime(date_cell).strftime('%Y-%m-%d')
        except:
            date_str = datetime.today().strftime('%Y-%m-%d')

        return part1[['Hall', 'EndTime']], part2[['Hall', 'EndTime']], date_str

    except Exception as e:
        st.error(f"Error processing file: {str(e)}")
        return None, None, None

def _draw_grid_on_figure(fig, data, title, date_str):
    """
    Internal helper function to draw the new grid layout onto a Matplotlib figure.
    """
    plt.rcParams['font.family'] = 'sans-serif'
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS'] # Font that supports Chinese characters

    total_items = len(data)
    if total_items == 0:
        return

    num_cols = 3
    num_rows = math.ceil(total_items / num_cols)

    A5_WIDTH_IN = 5.83
    A5_HEIGHT_IN = 8.27

    # 1. Redesign layout based on precise grid calculations
    margin_y = (A5_HEIGHT_IN / num_rows) / 4
    margin_x = margin_y # Use symmetric margins for a cleaner look

    # Prevent margins from becoming too large on pages with few items
    if A5_WIDTH_IN < 2 * margin_x or A5_HEIGHT_IN < 2 * margin_y:
        margin_x = A5_WIDTH_IN / 10
        margin_y = A5_HEIGHT_IN / 10

    printable_width = A5_WIDTH_IN - 2 * margin_x
    printable_height = A5_HEIGHT_IN - 2 * margin_y
    cell_width = printable_width / num_cols
    cell_height = printable_height / num_rows

    # Prepare data: Sort into column-first (Z-style) order for layout
    data_values = data.values.tolist()
    while len(data_values) % num_cols != 0:
        data_values.append(['', '']) # Pad data for a full grid
    rows_per_col_layout = math.ceil(len(data_values) / num_cols)
    sorted_data = [['', '']] * len(data_values)
    for i, item in enumerate(data_values):
        if item[0] and item[1]:
            row_in_col = i % rows_per_col_layout
            col_idx = i // rows_per_col_layout
            new_index = row_in_col * num_cols + col_idx
            if new_index < len(sorted_data):
                sorted_data[new_index] = item

    # --- Draw each cell onto the figure ---
    item_counter = 0
    for idx, (hall, end_time) in enumerate(sorted_data):
        if not (hall and end_time):
            continue

        item_counter += 1
        row_grid = idx // num_cols
        col_grid = idx % num_cols

        # Calculate position for each cell's axes in Figure coordinates [left, bottom, width, height]
        ax_left_in = margin_x + col_grid * cell_width
        ax_bottom_in = margin_y + (num_rows - 1 - row_grid) * cell_height # Y-axis from bottom
        ax_pos = [
            ax_left_in / A5_WIDTH_IN,
            ax_bottom_in / A5_HEIGHT_IN,
            cell_width / A5_WIDTH_IN,
            cell_height / A5_HEIGHT_IN,
        ]
        ax = fig.add_axes(ax_pos)

        # 2. Change Cell Border to a gray, dotted line
        for spine in ax.spines.values():
            spine.set_visible(True)
            spine.set_linestyle(':') # Dotted line style
            spine.set_edgecolor(BORDER_COLOR)
            spine.set_linewidth(1)

        # 4. Add Cell Index Number
        ax.text(0.07, 0.93, str(item_counter),
                transform=ax.transAxes,
                fontsize=9,
                color='grey',
                ha='left',
                va='top')

        # 3. Adjust Cell Content
        display_text = f"{hall}{end_time}"

        # Dynamically estimate font size to fill 90% of the cell width
        # This is a heuristic that provides a good balance of size and spacing.
        font_scale_factor = 1.7
        estimated_fontsize = (cell_width * 72 / len(display_text)) * font_scale_factor * 0.9
        max_fontsize = cell_height * 72 * 0.6 # Cap font size to 60% of cell height
        final_fontsize = min(estimated_fontsize, max_fontsize)

        ax.text(0.5, 0.5, display_text,
                fontsize=final_fontsize,
                fontweight='bold',
                ha='center',
                va='center',
                transform=ax.transAxes)

        ax.set_xticks([])
        ax.set_yticks([])

    # Add date and title to the top margin of the figure
    title_fontsize = margin_y * 72 * 0.3 # Scale title font size with margin
    fig.text(margin_x / A5_WIDTH_IN, 1 - (margin_y * 0.5 / A5_HEIGHT_IN),
             f"{date_str} {title}",
             fontsize=title_fontsize,
             color=DATE_COLOR,
             fontweight='bold',
             ha='left',
             va='center')


def create_print_layout(data, title, date_str):
    """
    Creates the final print-ready output in both PNG and PDF formats using the new grid layout.
    """
    if data.empty:
        return None

    # --- Create separate figures for PNG and PDF to ensure no cross-contamination ---
    png_fig = plt.figure(figsize=(5.83, 8.27), dpi=300)
    pdf_fig = plt.figure(figsize=(5.83, 8.27), dpi=300)

    # --- Draw the layout on both figures ---
    _draw_grid_on_figure(png_fig, data, title, date_str)
    _draw_grid_on_figure(pdf_fig, data, title, date_str)

    # --- Save PNG to a memory buffer ---
    png_buffer = io.BytesIO()
    png_fig.savefig(png_buffer, format='png', bbox_inches='tight', pad_inches=0.02)
    png_buffer.seek(0)
    png_base64 = base64.b64encode(png_buffer.getvalue()).decode()
    plt.close(png_fig)

    # --- Save PDF to a memory buffer ---
    pdf_buffer = io.BytesIO()
    with PdfPages(pdf_buffer) as pdf:
        pdf.savefig(pdf_fig, bbox_inches='tight', pad_inches=0.02)
    pdf_buffer.seek(0)
    pdf_base64 = base64.b64encode(pdf_buffer.getvalue()).decode()
    plt.close(pdf_fig)

    return {
        'png': f'data:image/png;base64,{png_base64}',
        'pdf': f'data:application/pdf;base64,{pdf_base64}'
    }

def display_pdf(base64_pdf):
    """
    Generates the HTML to embed and display a PDF in Streamlit.
    """
    pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="800" type="application/pdf"></iframe>'
    return pdf_display

# --- Streamlit User Interface ---
st.set_page_config(page_title="散厅时间快捷打印", layout="wide")
st.title("散厅时间快捷打印")

uploaded_file = st.file_uploader("上传【放映场次核对表.xls】文件", type=["xls", "xlsx"])

if uploaded_file:
    part1_data, part2_data, date_string = process_schedule(uploaded_file)

    if part1_data is not None and part2_data is not None:
        # Generate layouts for both day and night shifts
        part1_output = create_print_layout(part1_data, "A", date_string)
        part2_output = create_print_layout(part2_data, "C", date_string)

        col1, col2 = st.columns(2)

        with col1:
            st.subheader("白班散场预览 (时间 ≤ 17:30)")
            if part1_output:
                # Use tabs for PDF and PNG previews
                tab1_1, tab1_2 = st.tabs(["PDF 预览", "PNG 预览"])
                with tab1_1:
                    st.markdown(display_pdf(part1_output['pdf']), unsafe_allow_html=True)
                with tab1_2:
                    st.image(part1_output['png'])
            else:
                st.info("白班部分没有数据")

        with col2:
            st.subheader("夜班散场预览 (时间 > 17:30)")
            if part2_output:
                # Use tabs for PDF and PNG previews
                tab2_1, tab2_2 = st.tabs(["PDF 预览", "PNG 预览"])
                with tab2_1:
                    st.markdown(display_pdf(part2_output['pdf']), unsafe_allow_html=True)
                with tab2_2:
                    st.image(part2_output['png'])
            else:
                st.info("夜班部分没有数据")