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
from pypinyin import lazy_pinyin, Style
from matplotlib.backends.backend_pdf import PdfPages

def get_font(size=14):
    """
    Retrieves the specified font properties. Looks for 'simHei.ttc' first,
    then falls back to 'SimHei.ttf'.
    """
    font_path = "simHei.ttc"
    if not os.path.exists(font_path):
        font_path = "SimHei.ttf"
    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 text:
        return ""
    # Extract the first two Chinese characters
    chars = [c for c in text if '\u4e00' <= c <= '\u9fff']
    chars = chars[:2]
    # Get the first letter of the Pinyin for each character
    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 data.
    """
    try:
        # Try to read the date from the Excel file
        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()
        file.seek(0) # Reset file pointer after failed read attempt

    try:
        df = pd.read_excel(file, header=9, usecols=[1, 2, 4, 5])
        df.columns = ['Hall', 'StartTime', 'EndTime', 'Movie']
        
        df['Hall'] = df['Hall'].ffill()
        df.dropna(subset=['StartTime', 'EndTime', 'Movie'], inplace=True)
        
        df['Hall_Num'] = df['Hall'].astype(str).str.extract(r'(\d+)').astype(int)
        df['Hall'] = df['Hall_Num'].astype(str) + '号'

        # 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.loc[df['EndTime_dt'] < df['StartTime_dt'], 'EndTime_dt'] += timedelta(days=1)
        df = df.sort_values(['Hall_Num', 'StartTime_dt'])
        
        # Merge consecutive screenings of the same movie
        merged_rows = []
        for _, group in df.groupby('Hall_Num'):
            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']
                    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)
        
        # Adjust start and end times
        merged_df['StartTime_dt'] = merged_df['StartTime_dt'] - timedelta(minutes=10)
        merged_df['EndTime_dt'] = merged_df['EndTime_dt'] - timedelta(minutes=5)
        
        # Create final data columns for the table
        merged_df['Sequence'] = merged_df.groupby('Hall_Num').cumcount() + 1
        merged_df['Pinyin'] = merged_df['Movie'].apply(get_pinyin_abbr)
        merged_df['Time'] = merged_df['StartTime_dt'].dt.strftime('%H:%M') + ' - ' + merged_df['EndTime_dt'].dt.strftime('%H:%M')
        
        final_df = merged_df[['Hall', 'Hall_Num', 'Sequence', 'Movie', 'Pinyin', 'Time']].copy()
        final_df = final_df.sort_values(['Hall_Num', 'Sequence']).reset_index(drop=True)

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

def create_print_layout(data, date_str):
    """
    Generates the print layout as PNG and PDF based on the processed data.
    """
    if data is None or data.empty:
        return None

    def draw_table_on_ax(fig, ax):
        """Helper function to draw the table on a given Matplotlib axis."""
        ax.set_axis_off()
        
        num_rows = len(data)
        if num_rows == 0:
            return

        # --- Layout & Sizing Calculations ---
        # Define printable area using relative coordinates (fractions of the page)
        TOP_MARGIN, BOTTOM_MARGIN = 0.96, 0.04
        LEFT_MARGIN, RIGHT_MARGIN = 0.04, 0.96
        TABLE_HEIGHT = TOP_MARGIN - BOTTOM_MARGIN
        TABLE_WIDTH = RIGHT_MARGIN - LEFT_MARGIN

        # Total rows for calculation: content rows + 2 for top/bottom buffer
        total_layout_rows = num_rows + 2
        row_height = TABLE_HEIGHT / total_layout_rows
        
        # Font size is 90% of the calculated row height
        # (row_height is a fraction of figure height, convert to points)
        font_size_pt = (row_height * fig.get_figheight() * 0.9) * 72
        font = get_font(size=font_size_pt)
        
        # Relative column widths
        col_relative_widths = {'hall': 1.2, 'seq': 0.8, 'movie': 5.0, 'pinyin': 1.5, 'time': 2.5}
        total_rel_width = sum(col_relative_widths.values())
        
        col_widths = {k: (v / total_rel_width) * TABLE_WIDTH for k, v in col_relative_widths.items()}
        
        # Calculate the absolute x-position for the left edge of each column
        x_pos = {}
        current_x = LEFT_MARGIN
        col_order = ['hall', 'seq', 'movie', 'pinyin', 'time']
        for col_name in col_order:
            x_pos[col_name] = current_x
            current_x += col_widths[col_name]
        x_pos['end'] = current_x  # The right edge of the last column

        # Add date header
        ax.text(LEFT_MARGIN, TOP_MARGIN, date_str, fontsize=12, color='#A9A9A9',
                ha='left', va='top', fontproperties=get_font(12), transform=ax.transAxes)

        # --- Drawing Loop ---
        for i, row in data.iterrows():
            # Calculate y-position for the center of the current row
            y_center = TOP_MARGIN - (i + 1.5) * row_height # +1.5 to account for header space
            y_top = y_center + row_height / 2
            y_bottom = y_center - row_height / 2

            # --- 1. Draw Cell Content ---
            # Column content is ordered as per requirement
            ax.text(x_pos['hall'] + col_widths['hall']/2, y_center, f"{row['Hall']}", va='center', ha='center', fontproperties=font, transform=ax.transAxes)
            ax.text(x_pos['seq'] + col_widths['seq']/2, y_center, f"{row['Sequence']}", va='center', ha='center', fontproperties=font, transform=ax.transAxes)
            ax.text(x_pos['movie'] + 0.01, y_center, f"{row['Movie']}", va='center', ha='left', fontproperties=font, transform=ax.transAxes, clip_on=True)
            ax.text(x_pos['pinyin'] + col_widths['pinyin']/2, y_center, f"{row['Pinyin']}", va='center', ha='center', fontproperties=font, transform=ax.transAxes)
            ax.text(x_pos['time'] + col_widths['time']/2, y_center, f"{row['Time']}", va='center', ha='center', fontproperties=font, transform=ax.transAxes)
            
            # --- 2. Draw Cell Borders ---
            # Draw all vertical cell lines with gray dots
            for col_name in col_order:
                ax.plot([x_pos[col_name], x_pos[col_name]], [y_bottom, y_top], color='gray', linestyle=':', linewidth=0.7)
            ax.plot([x_pos['end'], x_pos['end']], [y_bottom, y_top], color='gray', linestyle=':', linewidth=0.7)

            # Draw top border of the cell
            ax.plot([LEFT_MARGIN, x_pos['end']], [y_top, y_top], color='gray', linestyle=':', linewidth=0.7)
            
            # --- 3. Draw Hall Separator or Bottom Border ---
            # Check if this is the last movie for the current hall
            is_last_in_hall = (i == num_rows - 1) or (row['Hall_Num'] != data.iloc[i + 1]['Hall_Num'])
            
            if is_last_in_hall:
                # If it's the last entry for a hall, draw a solid black line below it
                ax.plot([LEFT_MARGIN, x_pos['end']], [y_bottom, y_bottom], color='black', linestyle='-', linewidth=1.2, zorder=3)
            else:
                # Otherwise, draw a standard gray dotted line
                ax.plot([LEFT_MARGIN, x_pos['end']], [y_bottom, y_bottom], color='gray', linestyle=':', linewidth=0.7)

    # --- Generate PNG and PDF Outputs ---
    # Create a figure for PNG output
    png_fig = plt.figure(figsize=(8.27, 11.69), dpi=300)
    png_ax = png_fig.add_subplot(111)
    draw_table_on_ax(png_fig, png_ax)
    
    # Create a separate figure for PDF output
    pdf_fig = plt.figure(figsize=(8.27, 11.69), dpi=300)
    pdf_ax = pdf_fig.add_subplot(111)
    draw_table_on_ax(pdf_fig, pdf_ax)
    
    # Save PNG to a buffer
    png_buffer = io.BytesIO()
    png_fig.savefig(png_buffer, format='png', bbox_inches='tight', pad_inches=0.05)
    png_buffer.seek(0)
    image_base64 = base64.b64encode(png_buffer.getvalue()).decode()
    plt.close(png_fig)
    
    # Save PDF to a buffer
    pdf_buffer = io.BytesIO()
    pdf_fig.savefig(pdf_buffer, format='pdf', bbox_inches='tight', pad_inches=0.05)
    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 UI ---
st.set_page_config(page_title="LED 屏幕时间表打印", layout="wide")
st.title("LED 屏幕时间表打印")

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

if uploaded_file:
    with st.spinner("文件正在处理中，请稍候..."):
        schedule, date_str = process_schedule(io.BytesIO(uploaded_file.getvalue()))
        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("生成打印布局时出错。")
        else:
            st.error("无法处理文件，或文件中没有找到有效排片数据。请检查文件格式或内容是否正确。")