from collections import defaultdict
from dataclasses import dataclass
from typing import List, Dict, Tuple
import pandas as pd
import gradio as gr
import zipfile
import os
import shutil
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment

# --- Global Data Storage ---
all_semesters = defaultdict(lambda: defaultdict(list))
all_halls = []
current_mode = "major"  # default mode

# --- Subject Mapping ---
subject_map = {
    "BN": "Bengali", "CH": "Chemistry", "CS": "Computer Science", "EC": "Economics",
    "HI": "History", "MT": "Mathematics", "PI": "Philosophy", "PL": "Political Science",
    "PH": "Physics", "SN": "Sanskrit", "ST": "Statistics", "EN": "English"
}

# --- Dataclass for Halls ---
@dataclass
class ExamHall:
    id: int
    name: str
    capacity: int
    rows: int
    seats_per_row: int
    seating_arrangement: List[List[str]]

# --- Helper Functions ---
def clear_all():
    global all_semesters, all_halls
    all_semesters = defaultdict(lambda: defaultdict(list))
    all_halls = []
    return "All data cleared!", "", "", "", ""

def set_mode(mode_choice):
    global current_mode
    current_mode = mode_choice
    return f"Mode set to: {mode_choice.title()}"

def add_hall(hall_name, capacity, rows, seats_per_row):
    hall_id = len(all_halls) + 1
    hall = ExamHall(
        id=hall_id,
        name=hall_name,
        capacity=int(capacity),
        rows=int(rows),
        seats_per_row=int(seats_per_row),
        seating_arrangement=[[f"{r+1}{chr(65+c)}" for c in range(int(seats_per_row))] for r in range(int(rows))]
    )
    all_halls.append(hall)
    hall_data = [(h.name, h.capacity, h.rows, h.seats_per_row) for h in all_halls]
    hall_df = pd.DataFrame(hall_data, columns=["Hall Name", "Capacity", "Rows", "Seats per Row"])
    return f"{len(all_halls)} hall(s) added.", hall_df.to_html(index=False, escape=False)

def upload_semester(day, slot, sem_name, excel_file, num_subjects, subject_input):
    if not day or not slot:
        return "Day and Slot are required.", ""
    df = pd.read_excel(excel_file.name)
    subjects = []
    try:
        for line in subject_input.strip().split("\n"):
            subject, paper, duration = [x.strip() for x in line.split(",")]
            subjects.append((subject, paper, int(duration)))
    except Exception as e:
        return f"Error parsing subjects: {e}", ""
    all_semesters[day][slot].append((sem_name, subjects, df))
    preview_data = []
    for d, slots in all_semesters.items():
        for s, sem_data in slots.items():
            for sem_name_, subjects_, _ in sem_data:
                preview_data.append([d, s, sem_name_, ", ".join([f"{sub} - {pap}" for sub, pap, _ in subjects_])])
    preview_df = pd.DataFrame(preview_data, columns=["Day", "Slot", "Semester", "Subjects"])
    return f"Uploaded semester {sem_name} for {day}, {slot}", preview_df.to_html(index=False, escape=False)

def assign_papers_with_duration(df: pd.DataFrame, subjects_info: Dict[str, List[Dict]]) -> pd.DataFrame:
    rows = []
    for subject_code, paper_info_list in subjects_info.items():
        subject_students = df[df['Subject Code'] == subject_code]
        for paper_info in paper_info_list:
            temp_df = subject_students.copy()
            temp_df['Paper Name'] = paper_info['paper']
            temp_df['Duration'] = paper_info['duration']
            rows.append(temp_df)
    return pd.concat(rows, ignore_index=True)

def allocate_students(students: pd.DataFrame, halls: List[ExamHall], used_hall_ids: set, duration: int, semester_row_parity: Dict[str, int], mode: str) -> Tuple[Dict[int, List[tuple]], set, List[dict]]:
    seating_map = {}
    used_seats = {}
    unallocated_students = []
    halls_sorted = [h for h in sorted(halls, key=lambda h: h.capacity, reverse=True) if h.id not in used_hall_ids]
    updated_used_halls = set()

    for hall in halls_sorted:
        seating_map[hall.id] = []
        used_seats[hall.id] = set()

    subject_paper_groups = students.groupby(['Subject Code', 'Paper Name', 'Semester'])
    subject_blocks = list(subject_paper_groups)

    for i, ((sub_code, paper_name, semester), group) in enumerate(subject_blocks):
        group = group.sample(frac=1).reset_index(drop=True)
        idx = 0
        hall_idx = 0

        while idx < len(group) and hall_idx < len(halls_sorted):
            hall = halls_sorted[hall_idx]
            used = used_seats[hall.id]
            assignments = seating_map[hall.id]
            cols = hall.seats_per_row
            rows = hall.rows
            row_parity = semester_row_parity.get(semester, 0)

            col_step = 2  # Default for minor
            row_step = 2
            col_start = i % 2

            if mode == "major":
                if cols == 6:
                    col_step = 3
                    col_start = i % 3
                elif cols == 4:
                    col_step = 2
                    row_step = 2
                    col_start = i % 2

            for col in range(col_start, cols, col_step):
                for row in range(row_parity, rows, row_step):
                    if idx >= len(group):
                        break
                    seat = hall.seating_arrangement[row][col]
                    if seat in used:
                        continue
                    student = group.iloc[idx]
                    assignments.append((seat, student['Name'], student['Roll Number'], student['Subject'], paper_name, student['Semester']))
                    used.add(seat)
                    idx += 1
                if idx >= len(group):
                    break

            seating_map[hall.id] = assignments
            if assignments:
                updated_used_halls.add(hall.id)
            if idx < len(group):
                hall_idx += 1

        for j in range(idx, len(group)):
            student = group.iloc[j]
            unallocated_students.append({
                "Name": student['Name'],
                "Roll Number": student['Roll Number'],
                "Semester": student['Semester'],
                "Subject": student['Subject'],
                "Paper": paper_name
            })

    return seating_map, updated_used_halls, unallocated_students

def save_seating(seating_map, path, day_name, slot, duration, summary):
    os.makedirs(path, exist_ok=True)
    wb = Workbook()
    wb.remove(wb.active)
    has_data = False

    for hall_id, seats in seating_map.items():
        if not seats:
            continue
        has_data = True
        ws = wb.create_sheet(title=f"Hall {hall_id}")
        ws.append(["Seat", "Name", "Roll Number", "Subject", "Paper Name"])
        for col in range(1, 6):
            ws.cell(row=1, column=col).font = Font(bold=True)
            ws.cell(row=1, column=col).alignment = Alignment(horizontal="center")
        for seat in seats:
            ws.append(seat[:5])

        for sem in set(s[-1] for s in seats):
            subject_info = {}
            count = 0
            for s in seats:
                if s[-1] != sem:
                    continue
                count += 1
                subject_info.setdefault(s[3], set()).add(s[4])
            summary.append({
                "Day": day_name,
                "Slot": slot,
                "Duration": duration,
                "Room Number": hall_id,
                "Student Number": count,
                "Semester": sem,
                "Subjects and Papers": "; ".join(f"{k}: {', '.join(sorted(v))}" for k, v in subject_info.items())
            })

    if has_data:
        wb.save(os.path.join(path, f"{day_name}_{slot}_D{duration}.xlsx"))

def generate_summary(summary, path):
    df = pd.DataFrame(summary)
    df.sort_values(by=["Day", "Slot", "Room Number"], inplace=True)
    wb = Workbook()
    ws = wb.active
    ws.title = "Exam Schedule"
    headers = ["Day", "Slot", "Duration", "Room Number", "Student Number", "Semester", "Subjects and Papers"]
    ws.append(headers)
    for i, col in enumerate(headers, 1):
        ws.cell(row=1, column=i).font = Font(bold=True)
    for row in df.itertuples(index=False):
        ws.append(list(row))
    wb.save(path)

def save_unallocated_students(unallocated, path):
    if not unallocated:
        return
    df = pd.DataFrame(unallocated)
    df = df[["Name", "Roll Number", "Semester", "Subject", "Paper"]]
    df.to_excel(path, index=False)

def create_zip(root_folder, zip_name):
    with zipfile.ZipFile(zip_name, 'w') as zipf:
        for root, _, files in os.walk(root_folder):
            for file in files:
                filepath = os.path.join(root, file)
                arcname = os.path.relpath(filepath, root_folder)
                zipf.write(filepath, arcname)

def allocate_and_generate():
    if not all_semesters or not all_halls:
        return "Please upload semester data and halls first!", None, None, None
    try:
        if os.path.exists("seating_plans"):
            shutil.rmtree("seating_plans")
        halls = all_halls
        exam_summary = []
        all_unallocated = []
        for day, slots in all_semesters.items():
            for slot, semesters in slots.items():
                semester_data = []
                for sem_name, subjects, df in semesters:
                    subjects_info = defaultdict(list)
                    for subject, paper, duration in subjects:
                        subjects_info[subject].append({'paper': paper, 'duration': duration})
                    df['Semester'] = sem_name
                    df['Subject Code'] = df['Roll Number'].str[:2]
                    df['Subject'] = df['Subject Code'].map(subject_map)
                    full_data = assign_papers_with_duration(df, subjects_info)
                    semester_data.append((sem_name, full_data))
                semester_row_parity = {semester: idx % 2 for idx, (semester, _) in enumerate(semester_data)}
                used_halls = set()
                for duration in [2, 3]:
                    all_duration_data = [(semester, full_data[full_data['Duration'] == duration]) for semester, full_data in semester_data if not full_data[full_data['Duration'] == duration].empty]
                    if all_duration_data:
                        combined_df = pd.concat([d[1] for d in all_duration_data], ignore_index=True)
                        seating, used, unallocated = allocate_students(
                            combined_df, halls, used_halls, duration, semester_row_parity, mode=current_mode
                        )
                        used_halls.update(used)
                        all_unallocated.extend(unallocated)
                        save_seating(seating, f"seating_plans/{day}_{slot}_D{duration}", day, slot, duration, exam_summary)
        summary_path = "exam_schedule_summary.xlsx"
        unallocated_path = "unallocated_students.xlsx"
        generate_summary(exam_summary, summary_path)
        save_unallocated_students(all_unallocated, unallocated_path)
        create_zip("seating_plans", "exam_seating_plans.zip")
        return "Seating plan generated successfully!", "exam_seating_plans.zip", summary_path, unallocated_path
    except Exception as e:
        return f"Error: {str(e)}", None, None, None

# --- Gradio Interface ---
with gr.Blocks(title="Exam Seating Scheduler") as demo:
    gr.Image(
            value="https://drive.google.com/uc?id=1nmhNC3JjyU9YsVh-bzgzAmQZgloqNtEQ",
            height=158,
            width=158,
            show_label=False,
            container=False
     
        )
    gr.Markdown("<h1 style='text-align: center;'>📘 RKMRC Exam Seating Scheduler</h1>")
    #gr.Markdown("# 📘 Exam Seating Scheduler")

    with gr.Tab("🏫 Halls Setup"):
        gr.Markdown("### Add Exam Halls")
        with gr.Row():
            hall_name = gr.Textbox(label="Hall Name")
            capacity = gr.Number(label="Total Capacity", precision=0)
            rows = gr.Number(label="Number of Rows", precision=0)
            seats_per_row = gr.Number(label="Seats per Row", precision=0)
            add_hall_btn = gr.Button("Add Hall")
        hall_status = gr.Textbox(label="Hall Status", interactive=False)
        hall_preview = gr.HTML(label="Hall Preview")
        add_hall_btn.click(add_hall, inputs=[hall_name, capacity, rows, seats_per_row], outputs=[hall_status, hall_preview])

    with gr.Tab("📅 Upload Semester Data"):
        gr.Markdown("### Upload Semester Info")
        with gr.Row():
            day_input = gr.Textbox(label="Day Name (e.g. Day 1)")
            slot_dropdown = gr.Dropdown(choices=["10AM", "2PM"], label="Slot")
        sem_name = gr.Textbox(label="Semester Name")
        excel_input = gr.File(label="Upload Excel", file_types=[".xlsx"])
        num_subjects = gr.Number(label="Number of Subjects", precision=0)
        subject_input = gr.Textbox(label="Subjects (comma-separated per line)", lines=5)
        upload_btn = gr.Button("Upload Semester Data")
        upload_status = gr.Textbox(label="Upload Status", interactive=False)
        parsed_subjects_preview = gr.HTML(label="Subjects Preview")
        upload_btn.click(upload_semester, inputs=[day_input, slot_dropdown, sem_name, excel_input, num_subjects, subject_input], outputs=[upload_status, parsed_subjects_preview])

    with gr.Tab("📦 Generate Output"):
        gr.Markdown("### Select Seating Mode")
        mode_dropdown = gr.Dropdown(choices=["major", "minor"], value="major", label="Seating Mode")
        set_mode_btn = gr.Button("Set Mode")
        mode_status = gr.Textbox(label="Mode Status", interactive=False)
        set_mode_btn.click(set_mode, inputs=[mode_dropdown], outputs=[mode_status])

        allocate_btn = gr.Button("Allocate & Generate Output")
        alloc_status = gr.Textbox(label="Status", interactive=False)
        zip_output = gr.File(label="Download ZIP")
        summary_file = gr.File(label="Download Summary Excel")
        unallocated_file = gr.File(label="Unallocated Students Excel")
        allocate_btn.click(allocate_and_generate, outputs=[alloc_status, zip_output, summary_file, unallocated_file])

    with gr.Tab("🔁 Reset"):
        reset_btn = gr.Button("Clear All Data")
        reset_btn.click(clear_all, outputs=[hall_status, hall_preview, alloc_status, upload_status, parsed_subjects_preview])


if __name__ == "__main__":
    demo.launch()