app.py

import pandas as pd
import streamlit as st
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
import io
import base64
import os
from datetime import datetime, timedelta
import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg
from pypinyin import lazy_pinyin, Style
from matplotlib.backends.backend_pdf import PdfPages

def get_font(size=14):
    """Loads the specified font, with a fallback."""
    font_path = "simHei.ttc"
    if not os.path.exists(font_path):
        font_path = "SimHei.ttf"  # Fallback font
    if not os.path.exists(font_path):
        st.warning("Font file (simHei.ttc or SimHei.ttf) not found. Display may be incorrect.")
        return font_manager.FontProperties(size=size)
    return font_manager.FontProperties(fname=font_path, size=size)

def get_pinyin_abbr(text):
    """Gets the first letter of the Pinyin for the first two Chinese characters of a text."""
    if not isinstance(text, str):
        return ""
    # Extract the first two Chinese characters
    chars = [c for c in text if '\u4e00' <= c <= '\u9fff'][:2]
    if not chars:
        return ""
    pinyin_list = lazy_pinyin(chars, style=Style.FIRST_LETTER)
    return ''.join(pinyin_list).upper()

def process_schedule(file):
    """
    Processes the uploaded Excel file to extract and clean the movie schedule.
    This version also prepares all data fields needed for the new layout.
    """
    try:
        # Try to read the date from the specified cell
        date_df = pd.read_excel(file, header=None, skiprows=7, nrows=1, usecols=[3])
        date_str = pd.to_datetime(date_df.iloc[0, 0]).strftime('%Y-%m-%d')
        base_date = pd.to_datetime(date_str).date()
    except Exception:
        # Fallback to today's date if reading fails
        date_str = datetime.today().strftime('%Y-%m-%d')
        base_date = datetime.today().date()
    
    try:
        df = pd.read_excel(file, header=9, usecols=[1, 2, 4, 5])
        df.columns = ['Hall', 'StartTime', 'EndTime', 'Movie']
        
        # Data Cleaning
        df['Hall'] = df['Hall'].ffill()
        df.dropna(subset=['StartTime', 'EndTime', 'Movie'], inplace=True)
        df['Hall'] = df['Hall'].astype(str).str.extract(r'(\d+号)')
        df.dropna(subset=['Hall'], inplace=True) # Ensure rows without a hall number are dropped

        # Convert times to datetime objects
        df['StartTime_dt'] = pd.to_datetime(df['StartTime'], format='%H:%M', errors='coerce').apply(
            lambda t: t.replace(year=base_date.year, month=base_date.month, day=base_date.day) if pd.notnull(t) else t
        )
        df['EndTime_dt'] = pd.to_datetime(df['EndTime'], format='%H:%M', errors='coerce').apply(
            lambda t: t.replace(year=base_date.year, month=base_date.month, day=base_date.day) if pd.notnull(t) else t
        )
        
        df.dropna(subset=['StartTime_dt', 'EndTime_dt'], inplace=True)
        
        # Handle screenings that cross midnight
        df.loc[df['EndTime_dt'] < df['StartTime_dt'], 'EndTime_dt'] += timedelta(days=1)
        df = df.sort_values(['Hall', 'StartTime_dt'])
        
        # Merge consecutive screenings of the same movie
        merged_rows = []
        for _, group in df.groupby('Hall'):
            group = group.sort_values('StartTime_dt')
            current = None
            for _, row in group.iterrows():
                if current is None:
                    current = row.copy()
                else:
                    if row['Movie'] == current['Movie']:
                        current['EndTime_dt'] = row['EndTime_dt'] # Extend the end time
                    else:
                        merged_rows.append(current)
                        current = row.copy()
            if current is not None:
                merged_rows.append(current)
        
        if not merged_rows:
             return None, date_str

        merged_df = pd.DataFrame(merged_rows).reset_index(drop=True)
        
        # Adjust times as per original logic
        merged_df['StartTime_dt'] -= timedelta(minutes=10)
        merged_df['EndTime_dt'] -= timedelta(minutes=5)
        
        # --- New Data Preparation for Layout ---
        # 1. Create Index (序号)
        merged_df['Index'] = merged_df.groupby('Hall').cumcount() + 1
        # 2. Create Pinyin Abbreviation (拼音缩写)
        merged_df['Pinyin'] = merged_df['Movie'].apply(get_pinyin_abbr)
        # 3. Create Time String (时间)
        merged_df['TimeStr'] = merged_df['StartTime_dt'].dt.strftime('%H:%M') + ' - ' + merged_df['EndTime_dt'].dt.strftime('%H:%M')
        # 4. Clean Hall Number for display
        merged_df['Hall'] = merged_df['Hall'].str.replace('号', '')

        # Select and reorder columns as per requirement
        final_df = merged_df[['Hall', 'Index', 'Movie', 'Pinyin', 'TimeStr']]
        
        return final_df, date_str

    except Exception as e:
        st.error(f"An error occurred during file processing: {e}")
        return None, date_str


def create_print_layout(data, date_str):
    """
    Creates the print layout on an A4 page based on a dynamic grid system.
    """
    if data is None or data.empty:
        return None

    # --- 1. Layout Constants ---
    A4_WIDTH_IN, A4_HEIGHT_IN = 8.27, 11.69
    MARGIN_IN = 0.4
    USABLE_WIDTH_IN = A4_WIDTH_IN - (2 * MARGIN_IN)
    USABLE_HEIGHT_IN = A4_HEIGHT_IN - (2 * MARGIN_IN)

    # --- 2. Row and Font Calculation ---
    num_content_rows = len(data)
    total_grid_rows = num_content_rows + 2  # Add 2 for top/bottom padding rows
    row_height_in = USABLE_HEIGHT_IN / total_grid_rows
    # Calculate font size in points (1 inch = 72 points) to be 90% of row height
    font_size_pt = (row_height_in * 72) * 0.9
    content_font = get_font(font_size_pt)
    date_font = get_font(12)

    # --- 3. Column Width Calculation ---
    # Create a temporary figure to calculate text widths accurately
    temp_fig = plt.figure(figsize=(A4_WIDTH_IN, A4_HEIGHT_IN))
    canvas = FigureCanvasAgg(temp_fig)
    
    cols_to_measure = ['Hall', 'Index', 'Movie', 'Pinyin', 'TimeStr']
    col_widths_in = []

    for col in cols_to_measure:
        # Find the longest string in the column for measurement
        longest_item = max(data[col].astype(str).tolist(), key=len, default="")
        # Create a temporary text object to measure its width
        t = plt.text(0, 0, longest_item, fontproperties=content_font)
        # Get the bounding box of the text in display units and convert to inches
        bbox = t.get_window_extent(renderer=canvas.get_renderer())
        width_in = bbox.width / temp_fig.dpi
        col_widths_in.append(width_in * 1.1) # Add 10% padding
        t.remove()

    plt.close(temp_fig) # Close the temporary figure

    # Scale column widths to fit the usable page width
    total_calculated_width = sum(col_widths_in)
    scale_factor = USABLE_WIDTH_IN / total_calculated_width if total_calculated_width > 0 else 1
    final_col_widths_in = [w * scale_factor for w in col_widths_in]

    # --- 4. Figure and PDF/PNG Generation ---
    def process_figure(fig, ax):
        # Calculate grid coordinates in Axes units (0 to 1)
        col_widths_ax = [w / USABLE_WIDTH_IN for w in final_col_widths_in]
        row_height_ax = 1.0 / total_grid_rows
        
        x_coords_ax = [0] + np.cumsum(col_widths_ax).tolist()
        y_coords_ax = [1 - i * row_height_ax for i in range(total_grid_rows + 1)]

        # Add date string at the top-left of the usable area
        ax.text(0, 1, date_str, transform=ax.transAxes, fontproperties=date_font,
                ha='left', va='bottom', color='#A9A9A9')

        # --- Draw Grid and Content ---
        for i, row in data.iterrows():
            grid_row_index = i + 1  # Offset by 1 for the top padding row
            y_bottom = y_coords_ax[grid_row_index + 1]
            y_center = y_bottom + row_height_ax / 2

            # Draw bottom dotted line for the current row's cells
            ax.plot([0, 1], [y_bottom, y_bottom], transform=ax.transAxes,
                    linestyle=':', color='gray', linewidth=0.7)

            # Draw content for each cell in the row
            content_list = [row['Hall'], row['Index'], row['Movie'], row['Pinyin'], row['TimeStr']]
            for j, content in enumerate(content_list):
                x_left = x_coords_ax[j]
                x_center = x_left + col_widths_ax[j] / 2
                ax.text(x_center, y_center, content, transform=ax.transAxes,
                        fontproperties=content_font, ha='center', va='center')
        
        # --- Draw Vertical Grid Lines ---
        content_area_top_y = y_coords_ax[1]
        content_area_bottom_y = y_coords_ax[-2]
        for x in x_coords_ax[1:-1]:
            ax.plot([x, x], [content_area_bottom_y, content_area_top_y], transform=ax.transAxes,
                    linestyle=':', color='gray', linewidth=0.7)

        # --- Draw Black Separator Lines Between Halls ---
        hall_change_indices = data.index[data['Hall'] != data['Hall'].shift(-1)]
        for idx in hall_change_indices:
            # The line is at the bottom of the current row
            y_line = y_coords_ax[idx + 2] # +1 for top margin, +1 to get bottom of current row
            ax.plot([0, 1], [y_line, y_line], transform=ax.transAxes,
                    linestyle='-', color='black', linewidth=1.2)


    # Create figures for PNG and PDF
    png_fig = plt.figure(figsize=(A4_WIDTH_IN, A4_HEIGHT_IN), dpi=300)
    pdf_fig = plt.figure(figsize=(A4_WIDTH_IN, A4_HEIGHT_IN), dpi=300)
    
    # Configure axes to fill the usable area defined by margins
    ax_rect = [
        MARGIN_IN / A4_WIDTH_IN, MARGIN_IN / A4_HEIGHT_IN,
        USABLE_WIDTH_IN / A4_WIDTH_IN, USABLE_HEIGHT_IN / A4_HEIGHT_IN
    ]
    png_ax = png_fig.add_axes(ax_rect)
    pdf_ax = pdf_fig.add_axes(ax_rect)
    png_ax.axis('off')
    pdf_ax.axis('off')

    # Process both figures
    process_figure(png_fig, png_ax)
    process_figure(pdf_fig, pdf_ax)

    # Save PNG to buffer
    png_buffer = io.BytesIO()
    png_fig.savefig(png_buffer, format='png', pad_inches=0)
    png_buffer.seek(0)
    image_base64 = base64.b64encode(png_buffer.getvalue()).decode()
    plt.close(png_fig)

    # Save PDF to buffer
    pdf_buffer = io.BytesIO()
    pdf_fig.savefig(pdf_buffer, format='pdf', pad_inches=0)
    pdf_buffer.seek(0)
    pdf_base64 = base64.b64encode(pdf_buffer.getvalue()).decode()
    plt.close(pdf_fig)

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

def display_pdf(base64_pdf):
    """Embeds the PDF in the Streamlit app for display."""
    pdf_display = f'<iframe src="{base64_pdf}" width="100%" height="800" type="application/pdf"></iframe>'
    return pdf_display

# --- Streamlit App Main Body ---
st.set_page_config(page_title="LED 屏幕时间表打印", layout="wide")
st.title("LED 屏幕时间表打印")

uploaded_file = st.file_uploader("选择打开【放映时间核对表.xls】文件", type=["xls"])

if uploaded_file:
    with st.spinner("文件正在处理中，请稍候..."):
        schedule, date_str = process_schedule(uploaded_file)
        if schedule is not None and not schedule.empty:
            output = create_print_layout(schedule, date_str)
            if output:
                # Create tabs to switch between PDF and PNG previews
                tab1, tab2 = st.tabs(["PDF 预览", "PNG 预览"])
                
                with tab1:
                    st.markdown(display_pdf(output['pdf']), unsafe_allow_html=True)
                
                with tab2:
                    st.image(output['png'], use_container_width=True)
            else:
                st.error("生成打印布局失败。")
        elif schedule is None:
             st.error("无法处理文件，请检查文件格式或内容是否正确。")
        else: # schedule is empty
             st.warning("处理完成，但文件中没有找到有效的排片数据。")