import random
import pygame
from scipy.optimize import linear_sum_assignment

# ---------- 설정값 ----------
W, H = 900, 640
ROWS, COLS = 3, 4
DESK_W, DESK_H = 110, 70
GAP = 20
NEAR = DESK_W * 1.3          # 이 거리 안이면 '주변 자리'
SCORE = {1: 100, 2: 80, 3: 60}
INK, SUB = (30, 30, 30), (110, 110, 110)
BLUE, RED, GRAY = (90, 140, 220), (215, 95, 90), (225, 225, 225)

pygame.init()
screen = pygame.display.set_mode((W, H))
pygame.display.set_caption("교실 자리 배치 (단순 버전)")
clock = pygame.time.Clock()
font = pygame.font.SysFont("malgungothic", 20)
small = pygame.font.SysFont("malgungothic", 15)
pygame.key.start_text_input()

# ---------- 책상: 행 x 열 격자로 고정 생성 ----------
def make_desks():
    desks = []
    next_id = 1
    for r in range(ROWS):
        for c in range(COLS):
            x = 40 + c * (DESK_W + GAP)
            y = 80 + r * (DESK_H + GAP)
            desks.append({
                "id": next_id,
                "rect": pygame.Rect(x, y, DESK_W, DESK_H),
                "name": "",
                "prefs": [None, None, None],
            })
            next_id += 1
    return desks

desks = make_desks()
by_id = {d["id"]: d for d in desks}

# ---------- 텍스트 입력창 ----------
def make_box(rect, numeric=False):
    return {"rect": pygame.Rect(rect), "text": "", "active": False, "numeric": numeric}

def handle_box(box, event):
    if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
        box["active"] = box["rect"].collidepoint(event.pos)
    elif box["active"] and event.type == pygame.KEYDOWN and event.key == pygame.K_BACKSPACE:
        box["text"] = box["text"][:-1]
    elif box["active"] and event.type == pygame.TEXTINPUT:   # 한글 조합 완료 시점에 문자가 들어옴
        if not box["numeric"] or event.text.isdigit():
            box["text"] += event.text

def draw_box(box):
    pygame.draw.rect(screen, (255, 255, 255), box["rect"])
    pygame.draw.rect(screen, (60, 140, 220) if box["active"] else (150, 150, 150), box["rect"], 2)
    t = small.render(box["text"], True, INK)
    screen.blit(t, (box["rect"].x + 6, box["rect"].y + (box["rect"].h - t.get_height()) // 2))

# ---------- 책상 클릭 시 뜨는 이름/순위 입력 패널 ----------
panel_desk = None   # 지금 열려있는 책상 (None이면 패널 닫힘)
name_box = make_box((W - 210, 40, 160, 28))
pref_boxes = [make_box((W - 210, 78 + 34 * i, 160, 28), True) for i in range(3)]
save_btn = pygame.Rect(W - 210, 192, 75, 32)
cancel_btn = pygame.Rect(W - 125, 192, 85, 32)
calc_btn = pygame.Rect(20, H - 60, 140, 40)

def open_panel(d):
    global panel_desk
    panel_desk = d
    name_box["text"] = d["name"]
    for box, val in zip(pref_boxes, d["prefs"]):
        box["text"] = str(val) if val else ""

def save_panel():
    global panel_desk
    panel_desk["name"] = name_box["text"].strip()
    panel_desk["prefs"] = [int(b["text"]) if b["text"].strip().isdigit() else None for b in pref_boxes]
    panel_desk = None

# ---------- 만족도 계산 & 최적 배정 ----------
def dist(a, b):
    """두 점 사이의 직선 거리. '주변 자리'인지 판정할 때 씀."""
    return ((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2) ** 0.5

def satisfaction(prefs, desk):
    """이 학생(prefs)이 desk에 앉았을 때 만족도(%). 순위 그대로면 만점, 주변 자리면 -10점."""
    best = 0
    for rank, pid in enumerate(prefs, 1):
        if pid is None or pid not in by_id:
            continue
        base = SCORE[rank]
        if desk["id"] == pid:
            best = max(best, base)
        elif dist(desk["rect"].center, by_id[pid]["rect"].center) <= NEAR:
            best = max(best, base - 10)
    return best

def compute_and_apply():
    """만족도 합이 최대가 되는 배정을 계산해서, 결과를 바로 desks에 적용한다."""
    students = [d for d in desks if d["name"].strip()]
    if not students:   # 아무도 입력 안 했으면 데모용으로 무작위 채우기
        for d in desks:
            d["name"] = f"학생{d['id']}"
            d["prefs"] = random.sample(list(by_id), 3)
        students = desks

    # cost[i][j] = i번째 학생을 j번째 책상에 앉힐 때의 '비용'.
    # linear_sum_assignment는 최소화만 하므로, 만족도에 -를 붙여 최대화 문제로 바꾼다.
    cost = [[-satisfaction(s["prefs"], d) for d in desks] for s in students]
    rows, cols = linear_sum_assignment(cost)   # rows[i], cols[i] = 한 쌍(학생, 배정된 책상)

    # 배정 결과를 값으로 미리 복사해둔다 (딕셔너리를 바로 덮어쓰면 순서에 따라 꼬일 수 있음)
    names, prefs = {}, {}
    for r, c in zip(rows, cols):
        names[desks[c]["id"]] = students[r]["name"]
        prefs[desks[c]["id"]] = list(students[r]["prefs"])
    for d in desks:
        d["name"] = names.get(d["id"], "")
        d["prefs"] = prefs.get(d["id"], [None, None, None])

# ---------- 그리기 ----------
def color_for(d):
    """만족도 0보다 크면 파랑(만족), 아니면 빨강(불만족), 빈 자리는 회색."""
    if not d["name"]:
        return GRAY
    return BLUE if satisfaction(d["prefs"], d) > 0 else RED

def draw_center_text(surf, text, color_, rect):
    t = small.render(text, True, color_)
    surf.blit(t, t.get_rect(center=rect.center))

def draw():
    screen.fill((245, 244, 238))
    screen.blit(small.render("좌클릭: 이름/순위 입력 · '배치 계산': 최적 자리 배정", True, SUB), (20, 20))

    for d in desks:
        pygame.draw.rect(screen, color_for(d), d["rect"], border_radius=6)
        pygame.draw.rect(screen, (60, 60, 60), d["rect"], 2, border_radius=6)
        screen.blit(small.render(str(d["id"]), True, (255, 255, 255)), (d["rect"].x + 4, d["rect"].y + 2))
        if d["name"]:
            t = font.render(d["name"], True, (255, 255, 255))
            screen.blit(t, t.get_rect(center=d["rect"].center))

    pygame.draw.rect(screen, (70, 150, 90), calc_btn, border_radius=8)
    draw_center_text(screen, "배치 계산", (255, 255, 255), calc_btn)

    if panel_desk:
        panel = pygame.Rect(W - 230, 20, 220, 210)
        pygame.draw.rect(screen, (250, 250, 245), panel, border_radius=10)
        pygame.draw.rect(screen, (120, 120, 120), panel, 2, border_radius=10)
        for label, box in zip(["이름", "1순위", "2순위", "3순위"], [name_box] + pref_boxes):
            screen.blit(small.render(label, True, SUB), (box["rect"].x - 55, box["rect"].y + 6))
            draw_box(box)
        pygame.draw.rect(screen, (70, 150, 90), save_btn, border_radius=6)
        pygame.draw.rect(screen, (170, 80, 70), cancel_btn, border_radius=6)
        draw_center_text(screen, "저장", (255, 255, 255), save_btn)
        draw_center_text(screen, "취소", (255, 255, 255), cancel_btn)

    pygame.display.flip()

# ---------- 메인 루프 ----------
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        if panel_desk:   # 패널이 열려있으면 책상 클릭은 무시하고 입력창/저장/취소만 처리
            handle_box(name_box, event)
            for b in pref_boxes:
                handle_box(b, event)
            if event.type == pygame.MOUSEBUTTONDOWN:
                if save_btn.collidepoint(event.pos):
                    save_panel()
                elif cancel_btn.collidepoint(event.pos):
                    panel_desk = None
            continue

        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            if calc_btn.collidepoint(event.pos):
                compute_and_apply()
            else:
                for d in desks:   # 클릭 좌표에 있는 책상을 찾아서 패널 열기
                    if d["rect"].collidepoint(event.pos):
                        open_panel(d)
                        break

    draw()
    clock.tick(60)

pygame.quit()