#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Module: PDF to Cells Extractor
Extracts individual voter cells from PDF pages
"""

import cv2
import numpy as np
from pdf2image import convert_from_path
import os


class PDFToCellsExtractor:
    def __init__(self, pdf_path):
        self.pdf_path = pdf_path

    def detect_grid_cells(self, page_image):
        """Detect voter cells from page using grid detection"""
        gray = cv2.cvtColor(page_image, cv2.COLOR_BGR2GRAY)

        # Detect horizontal lines
        horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (80, 1))
        horizontal_lines = cv2.morphologyEx(
            cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1],
            cv2.MORPH_OPEN, horizontal_kernel, iterations=2
        )

        # Detect vertical lines
        vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 80))
        vertical_lines = cv2.morphologyEx(
            cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1],
            cv2.MORPH_OPEN, vertical_kernel, iterations=2
        )

        # Combine lines
        grid = cv2.add(horizontal_lines, vertical_lines)

        # Find contours
        contours, _ = cv2.findContours(grid, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # Filter and sort by area
        cells = []
        page_h, page_w = page_image.shape[:2]
        min_area = (page_w * page_h) * 0.02  # At least 2% of page

        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > min_area:
                x, y, w, h = cv2.boundingRect(cnt)
                # Reasonable aspect ratio for voter cells
                aspect = float(w) / h if h > 0 else 0
                if 0.3 < aspect < 3.0:
                    cells.append((x, y, w, h))

        # Sort by position (top to bottom, left to right)
        cells = sorted(cells, key=lambda c: (c[1], c[0]))

        return cells

    def extract_page(self, page_num, output_dir):
        """Extract cells from a single page"""
        os.makedirs(output_dir, exist_ok=True)

        # Convert PDF page to image
        pages = convert_from_path(
            self.pdf_path,
            dpi=300,
            first_page=page_num,
            last_page=page_num
        )

        if not pages:
            return 0

        # Convert to OpenCV format
        page_pil = pages[0]
        page_np = np.array(page_pil)
        page_bgr = cv2.cvtColor(page_np, cv2.COLOR_RGB2BGR)

        # Detect cells
        cells = self.detect_grid_cells(page_bgr)

        # Extract and save each cell
        for idx, (x, y, w, h) in enumerate(cells, 1):
            cell_img = page_bgr[y:y+h, x:x+w]
            output_path = f"{output_dir}/page{page_num:03d}_cell{idx:02d}.png"
            cv2.imwrite(output_path, cell_img)

        return len(cells)

    def extract_pages(self, start_page, end_page, output_dir):
        """Extract cells from multiple pages"""
        os.makedirs(output_dir, exist_ok=True)

        total_cells = 0

        for page_num in range(start_page, end_page + 1):
            print(f"  Page {page_num}/{end_page}...", end=' ')

            num_cells = self.extract_page(page_num, output_dir)
            total_cells += num_cells

            print(f"{num_cells} cells extracted")

        return total_cells